github.com/cilium/ebpf@v0.15.1-0.20240517100537-8079b37aa138/internal/testutils/cpu.go (about)

     1  package testutils
     2  
     3  import (
     4  	"runtime"
     5  	"testing"
     6  
     7  	"github.com/cilium/ebpf/internal/unix"
     8  
     9  	"github.com/go-quicktest/qt"
    10  )
    11  
    12  // LockOSThreadToSingleCPU force the current goroutine to run on a single CPU.
    13  func LockOSThreadToSingleCPU(tb testing.TB) {
    14  	tb.Helper()
    15  
    16  	runtime.LockOSThread()
    17  	tb.Cleanup(runtime.UnlockOSThread)
    18  
    19  	var old unix.CPUSet
    20  	err := unix.SchedGetaffinity(0, &old)
    21  	qt.Assert(tb, qt.IsNil(err))
    22  
    23  	// Schedule test to run on only CPU 0
    24  	var first unix.CPUSet
    25  	first.Set(0)
    26  	err = unix.SchedSetaffinity(0, &first)
    27  	qt.Assert(tb, qt.IsNil(err))
    28  
    29  	tb.Cleanup(func() {
    30  		_ = unix.SchedSetaffinity(0, &old)
    31  	})
    32  }