package hst import ( "bufio" "errors" "fmt" . "github.com/onsi/ginkgo/v2" "os" "os/exec" "strings" ) var CgroupPath = "/sys/fs/cgroup/" type CpuContext struct { cpuAllocator *CpuAllocatorT cpus []int } type CpuAllocatorT struct { cpus []int } var cpuAllocator *CpuAllocatorT = nil func (c *CpuAllocatorT) Allocate(containerCount int, nCpus int) (*CpuContext, error) { var cpuCtx CpuContext // splitting cpus into equal parts; this will over-allocate cores but it's good enough for now maxContainerCount := 4 // skip CPU 0 minCpu := ((GinkgoParallelProcess() - 1) * maxContainerCount * nCpus) + 1 maxCpu := (GinkgoParallelProcess() * maxContainerCount * nCpus) if len(c.cpus)-1 < maxCpu { err := fmt.Errorf("could not allocate %d CPUs; available: %d; attempted to allocate cores %d-%d", nCpus*containerCount, len(c.cpus)-1, minCpu, maxCpu) return nil, err } if containerCount == 1 { cpuCtx.cpus = c.cpus[minCpu : minCpu+nCpus] } else if containerCount > 1 && containerCount <= maxContainerCount { cpuCtx.cpus = c.cpus[minCpu+(nCpus*(containerCount-1)) : minCpu+(nCpus*containerCount)] } else { return nil, fmt.Errorf("too many containers; CPU allocation for >%d containers is not implemented", maxContainerCount) } cpuCtx.cpuAllocator = c return &cpuCtx, nil } func (c *CpuAllocatorT) readCpus() error { var first, last int // Path depends on cgroup version. We need to check which version is in use. // For that following command can be used: 'stat -fc %T /sys/fs/cgroup/' // In case the output states 'cgroup2fs' then cgroups v2 is used, 'tmpfs' in case cgroups v1. cmd := exec.Command("stat", "-fc", "%T", "/sys/fs/cgroup/") byteOutput, err := cmd.CombinedOutput() if err != nil { return err } CpuPath := CgroupPath if strings.Contains(string(byteOutput), "tmpfs") { CpuPath += "cpuset/cpuset.effective_cpus" } else if strings.Contains(string(byteOutput), "cgroup2fs") { CpuPath += "cpuset.cpus.effective" } else { return errors.New("cgroup unknown fs: " + string(byteOutput)) } file, err := os.Open(CpuPath) if err != nil { return err } defer file.Close() sc := bufio.NewScanner(file) sc.Scan() line := sc.Text() _, err = fmt.Sscanf(line, "%d-%d", &first, &last) if err != nil { return err } for i := first; i <= last; i++ { c.cpus = append(c.cpus, i) } return nil } func CpuAllocator() (*CpuAllocatorT, error) { if cpuAllocator == nil { cpuAllocator = new(CpuAllocatorT) err := cpuAllocator.readCpus() if err != nil { return nil, err } } return cpuAllocator, nil }