github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/procid/procid_test.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package procid
    16  
    17  import (
    18  	"os"
    19  	"runtime"
    20  	"testing"
    21  
    22  	"golang.org/x/sys/unix"
    23  	"github.com/SagerNet/gvisor/pkg/sync"
    24  )
    25  
    26  // runOnMain is used to send functions to run on the main (initial) thread.
    27  var runOnMain = make(chan func(), 10)
    28  
    29  func checkProcid(t *testing.T, start *sync.WaitGroup, done *sync.WaitGroup) {
    30  	defer done.Done()
    31  
    32  	runtime.LockOSThread()
    33  	defer runtime.UnlockOSThread()
    34  
    35  	start.Done()
    36  	start.Wait()
    37  
    38  	procID := Current()
    39  	tid := unix.Gettid()
    40  
    41  	if procID != uint64(tid) {
    42  		t.Logf("Bad procid: expected %v, got %v", tid, procID)
    43  		t.Fail()
    44  	}
    45  }
    46  
    47  func TestProcidInitialized(t *testing.T) {
    48  	var start sync.WaitGroup
    49  	var done sync.WaitGroup
    50  
    51  	count := 100
    52  	start.Add(count + 1)
    53  	done.Add(count + 1)
    54  
    55  	// Run the check on the main thread.
    56  	//
    57  	// When cgo is not included, the only case when procid isn't initialized
    58  	// is in the main (initial) thread, so we have to test this case
    59  	// specifically.
    60  	runOnMain <- func() {
    61  		checkProcid(t, &start, &done)
    62  	}
    63  
    64  	// Run the check on a number of different threads.
    65  	for i := 0; i < count; i++ {
    66  		go checkProcid(t, &start, &done)
    67  	}
    68  
    69  	done.Wait()
    70  }
    71  
    72  func TestMain(m *testing.M) {
    73  	// Make sure we remain at the main (initial) thread.
    74  	runtime.LockOSThread()
    75  
    76  	// Start running tests in a different goroutine.
    77  	go func() {
    78  		os.Exit(m.Run())
    79  	}()
    80  
    81  	// Execute any functions that have been sent for execution in the main
    82  	// thread.
    83  	for f := range runOnMain {
    84  		f()
    85  	}
    86  }