github.com/MerlinKodo/gvisor@v0.0.0-20231110090155-957f62ecf90e/pkg/sentry/kernel/sched/cpuset.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 sched
    16  
    17  import "math/bits"
    18  
    19  const (
    20  	bitsPerByte  = 8
    21  	bytesPerLong = 8 // only for 64-bit architectures
    22  )
    23  
    24  // CPUSet contains a bitmap to record CPU information.
    25  //
    26  // Note that this definition is only correct for little-endian architectures,
    27  // since Linux's cpumask_t uses unsigned long.
    28  type CPUSet []byte
    29  
    30  // CPUSetSize returns the size in bytes of a CPUSet that can contain num cpus.
    31  func CPUSetSize(num uint) uint {
    32  	// NOTE(b/68859821): Applications may expect that the size of a CPUSet in
    33  	// bytes is always a multiple of sizeof(unsigned long), since this is true
    34  	// in Linux. Thus we always round up.
    35  	bytes := (num + bitsPerByte - 1) / bitsPerByte
    36  	longs := (bytes + bytesPerLong - 1) / bytesPerLong
    37  	return longs * bytesPerLong
    38  }
    39  
    40  // NewCPUSet returns a CPUSet for the given number of CPUs which initially
    41  // contains no CPUs.
    42  func NewCPUSet(num uint) CPUSet {
    43  	return CPUSet(make([]byte, CPUSetSize(num)))
    44  }
    45  
    46  // NewFullCPUSet returns a CPUSet for the given number of CPUs, all of which
    47  // are present in the set.
    48  func NewFullCPUSet(num uint) CPUSet {
    49  	c := NewCPUSet(num)
    50  	var i uint
    51  	for ; i < num/bitsPerByte; i++ {
    52  		c[i] = 0xff
    53  	}
    54  	if rem := num % bitsPerByte; rem != 0 {
    55  		c[i] = (1 << rem) - 1
    56  	}
    57  	return c
    58  }
    59  
    60  // Size returns the size of 'c' in bytes.
    61  func (c CPUSet) Size() uint {
    62  	return uint(len(c))
    63  }
    64  
    65  // NumCPUs returns how many cpus are set in the CPUSet.
    66  func (c CPUSet) NumCPUs() uint {
    67  	var n int
    68  	for _, b := range c {
    69  		n += bits.OnesCount8(b)
    70  	}
    71  	return uint(n)
    72  }
    73  
    74  // Copy returns a copy of the CPUSet.
    75  func (c CPUSet) Copy() CPUSet {
    76  	return append(CPUSet(nil), c...)
    77  }
    78  
    79  // Set sets the bit corresponding to cpu.
    80  func (c *CPUSet) Set(cpu uint) {
    81  	(*c)[cpu/bitsPerByte] |= 1 << (cpu % bitsPerByte)
    82  }
    83  
    84  // ClearAbove clears bits corresponding to cpu and all higher cpus.
    85  func (c *CPUSet) ClearAbove(cpu uint) {
    86  	i := cpu / bitsPerByte
    87  	if i >= c.Size() {
    88  		return
    89  	}
    90  	(*c)[i] &^= 0xff << (cpu % bitsPerByte)
    91  	for i++; i < c.Size(); i++ {
    92  		(*c)[i] = 0
    93  	}
    94  }
    95  
    96  // ForEachCPU iterates over the CPUSet and calls fn with the cpu index if
    97  // it's set.
    98  func (c CPUSet) ForEachCPU(fn func(uint)) {
    99  	for i := uint(0); i < c.Size()*bitsPerByte; i++ {
   100  		bit := uint(1) << (i & (bitsPerByte - 1))
   101  		if uint(c[i/bitsPerByte])&bit == bit {
   102  			fn(i)
   103  		}
   104  	}
   105  }