1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
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
}
|
