diff options
Diffstat (limited to 'extras/hs-test/cpu.go')
-rw-r--r-- | extras/hs-test/cpu.go | 59 |
1 files changed, 45 insertions, 14 deletions
diff --git a/extras/hs-test/cpu.go b/extras/hs-test/cpu.go index e17bc11fbe0..49a7dfb02f8 100644 --- a/extras/hs-test/cpu.go +++ b/extras/hs-test/cpu.go @@ -2,43 +2,74 @@ package main import ( "bufio" + "errors" "fmt" + . "github.com/onsi/ginkgo/v2" "os" + "os/exec" + "strings" ) -var CPU_PATH = "/sys/fs/cgroup/cpuset.cpus.effective" +var CgroupPath = "/sys/fs/cgroup/" type CpuContext struct { cpuAllocator *CpuAllocatorT cpus []int } -func (c *CpuContext) Release() { - c.cpuAllocator.cpus = append(c.cpuAllocator.cpus, c.cpus...) - c.cpus = c.cpus[:0] // empty the list -} - type CpuAllocatorT struct { cpus []int } var cpuAllocator *CpuAllocatorT = nil -func (c *CpuAllocatorT) Allocate(nCpus int) (*CpuContext, error) { +func (c *CpuAllocatorT) Allocate(containerCount int, nCpus int) (*CpuContext, error) { var cpuCtx CpuContext - if len(c.cpus) < nCpus { - return nil, fmt.Errorf("could not allocate %d CPUs; available: %d", nCpus, len(c.cpus)) + // 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 } - cpuCtx.cpus = c.cpus[0:nCpus] + 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 - c.cpus = c.cpus[nCpus:] return &cpuCtx, nil } -func (c *CpuAllocatorT) readCpus(fname string) error { +func (c *CpuAllocatorT) readCpus() error { var first, last int - file, err := os.Open(CPU_PATH) + + // 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 } @@ -60,7 +91,7 @@ func (c *CpuAllocatorT) readCpus(fname string) error { func CpuAllocator() (*CpuAllocatorT, error) { if cpuAllocator == nil { cpuAllocator = new(CpuAllocatorT) - err := cpuAllocator.readCpus(CPU_PATH) + err := cpuAllocator.readCpus() if err != nil { return nil, err } |