diff --git a/pkg/cgroup/benchmark_test.go b/pkg/cgroup/benchmark_test.go index 0b388db..b775d4d 100644 --- a/pkg/cgroup/benchmark_test.go +++ b/pkg/cgroup/benchmark_test.go @@ -6,7 +6,7 @@ import ( ) func BenchmarkCgroup(b *testing.B) { - builder, err := NewBuilder("benchmark").WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() + builder, err := NewBuilder("benchmark").WithCPU().WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() if err != nil { b.Fatal(err) } @@ -36,11 +36,11 @@ func BenchmarkCgroup(b *testing.B) { } func TestCgroupAll(t *testing.T) { - // ensure root privillege when testing + // ensure root privilege when testing if os.Getuid() != 0 { t.Skip("no root privillege") } - builder, err := NewBuilder("test").WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() + builder, err := NewBuilder("test").WithCPU().WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() if err != nil { t.Fatal(err) } diff --git a/pkg/cgroup/builder.go b/pkg/cgroup/builder.go index ab7290c..82a4bfd 100644 --- a/pkg/cgroup/builder.go +++ b/pkg/cgroup/builder.go @@ -10,6 +10,7 @@ import ( type Builder struct { Prefix string + CPU bool CPUSet bool CPUAcct bool Memory bool @@ -23,6 +24,12 @@ func NewBuilder(prefix string) *Builder { } } +// WithCPU includes cpu cgroup +func (b *Builder) WithCPU() *Builder { + b.CPU = true + return b +} + // WithCPUSet includes cpuset cgroup func (b *Builder) WithCPUSet() *Builder { b.CPUSet = true diff --git a/pkg/cgroup/cgroup.go b/pkg/cgroup/cgroup.go index 83569e3..53ef7f7 100644 --- a/pkg/cgroup/cgroup.go +++ b/pkg/cgroup/cgroup.go @@ -13,62 +13,77 @@ import ( type Cgroup struct { prefix string + cpu *SubCgroup cpuset *SubCgroup cpuacct *SubCgroup memory *SubCgroup pids *SubCgroup + + all []*SubCgroup } // Build creates new cgrouup directories func (b *Builder) Build() (cg *Cgroup, err error) { var ( - cpuSetPath, cpuacctPath, memoryPath, pidsPath string + paths []string + subCgroup = make(map[int]*SubCgroup) + info = getCachedCgroupHierarchy() ) // if failed, remove potential created directory defer func() { if err != nil { - remove(cpuSetPath) - remove(cpuacctPath) - remove(memoryPath) - remove(pidsPath) + removeAll(paths...) } }() for _, c := range []struct { enable bool name string - path *string }{ - {b.CPUSet, "cpuset", &cpuSetPath}, - {b.CPUAcct, "cpuacct", &cpuacctPath}, - {b.Memory, "memory", &memoryPath}, - {b.Pids, "pids", &pidsPath}, + {b.CPU, "cpu"}, + {b.CPUSet, "cpuset"}, + {b.CPUAcct, "cpuacct"}, + {b.Memory, "memory"}, + {b.Pids, "pids"}, } { if !c.enable { continue } - if *c.path, err = CreateSubCgroupPath(c.name, b.Prefix); err != nil { - return + h := info[c.name] + if subCgroup[h] == nil { + var path string + if path, err = CreateSubCgroupPath(c.name, b.Prefix); err != nil { + return + } + paths = append(paths, path) + subCgroup[h] = NewSubCgroup(path) } } if b.CPUSet { - if err = initCpuset(cpuSetPath); err != nil { + if err = initCpuset(subCgroup[info["cpuset"]].path); err != nil { return } } + var all []*SubCgroup + for _, v := range subCgroup { + all = append(all, v) + } + return &Cgroup{ prefix: b.Prefix, - cpuset: NewSubCgroup(cpuSetPath), - cpuacct: NewSubCgroup(cpuacctPath), - memory: NewSubCgroup(memoryPath), - pids: NewSubCgroup(pidsPath), + cpu: subCgroup[info["cpu"]], + cpuset: subCgroup[info["cpuset"]], + cpuacct: subCgroup[info["cpuacct"]], + memory: subCgroup[info["memory"]], + pids: subCgroup[info["pids"]], + all: all, }, nil } // AddProc writes cgroup.procs to all sub-cgroup func (c *Cgroup) AddProc(pid int) error { - for _, s := range []*SubCgroup{c.cpuset, c.cpuacct, c.memory, c.pids} { + for _, s := range c.all { if err := s.WriteUint(cgroupProcs, uint64(pid)); err != nil { return err } @@ -79,7 +94,7 @@ func (c *Cgroup) AddProc(pid int) error { // Destroy removes dir for sub-cgroup, errors are ignored if remove one failed func (c *Cgroup) Destroy() error { var err1 error - for _, s := range []*SubCgroup{c.cpuset, c.cpuset, c.memory, c.pids} { + for _, s := range c.all { if err := remove(s.path); err != nil { err1 = err } @@ -87,6 +102,16 @@ func (c *Cgroup) Destroy() error { return err1 } +// SetCPUCfsPeriod set cpu.cfs_period_us in us +func (c *Cgroup) SetCPUCfsPeriod(p uint64) error { + return c.cpu.WriteUint("cpu.cfs_period_us", p) +} + +// SetCPUCfsQuota set cpu.cfs_quota_us in us +func (c *Cgroup) SetCPUCfsQuota(p uint64) error { + return c.cpu.WriteUint("cpu.cfs_quota_us", p) +} + // SetCpusetCpus set cpuset.cpus func (c *Cgroup) SetCpusetCpus(b []byte) error { return c.cpuset.WriteFile("cpuset.cpus", b) @@ -112,6 +137,16 @@ func (c *Cgroup) SetMemoryLimitInBytes(i uint64) error { return c.memory.WriteUint("memory.limit_in_bytes", i) } +// MemoryMemswMaxUsageInBytes read memory.memsw.max_usage_in_bytes +func (c *Cgroup) MemoryMemswMaxUsageInBytes() (uint64, error) { + return c.memory.ReadUint("memory.memsw.max_usage_in_bytes") +} + +// SetMemoryMemswLimitInBytes write memory.memsw.limit_in_bytes +func (c *Cgroup) SetMemoryMemswLimitInBytes(i uint64) error { + return c.memory.WriteUint("memory.memsw.limit_in_bytes", i) +} + // SetPidsMax write pids.max func (c *Cgroup) SetPidsMax(i uint64) error { return c.pids.WriteUint("pids.max", i) @@ -183,3 +218,14 @@ func remove(name string) error { } return nil } + +func removeAll(name ...string) error { + var err1 error + for _, n := range name { + err := remove(n) + if err != nil && err1 == nil { + err1 = err + } + } + return err1 +} diff --git a/pkg/cgroup/cgroup_info.go b/pkg/cgroup/cgroup_info.go new file mode 100644 index 0000000..5b591e3 --- /dev/null +++ b/pkg/cgroup/cgroup_info.go @@ -0,0 +1,75 @@ +package cgroup + +import ( + "bufio" + "os" + "strconv" + "strings" +) + +// Info reads the cgroup mount info from /proc/cgroups +type Info struct { + Hierarchy int + NumCgroups int + Enabled bool +} + +// GetCgroupInfo read /proc/cgroups and return the result +func GetCgroupInfo() (map[string]Info, error) { + f, err := os.Open(procCgroupsPath) + if err != nil { + return nil, err + } + defer f.Close() + + rt := make(map[string]Info) + s := bufio.NewScanner(f) + for s.Scan() { + text := s.Text() + if text[0] == '#' { + continue + } + parts := strings.Fields(text) + if len(parts) < 4 { + continue + } + + // format: subsys_name hierarchy num_cgroups enabled + name := parts[0] + hierarchy, err := strconv.Atoi(parts[1]) + if err != nil { + return nil, err + } + numCgroups, err := strconv.Atoi(parts[2]) + if err != nil { + return nil, err + } + enabled := parts[3] != "0" + rt[name] = Info{ + Hierarchy: hierarchy, + NumCgroups: numCgroups, + Enabled: enabled, + } + } + if err := s.Err(); err != nil { + return nil, err + } + return rt, nil +} + +// GetAllSubCgroup reads /proc/cgroups and get all available sub-cgroup as set +func GetAllSubCgroup() (map[string]bool, error) { + info, err := GetCgroupInfo() + if err != nil { + return nil, err + } + + rt := make(map[string]bool) + for k, v := range info { + if !v.Enabled { + continue + } + rt[k] = true + } + return rt, nil +} diff --git a/pkg/cgroup/cgroup_info_cache.go b/pkg/cgroup/cgroup_info_cache.go new file mode 100644 index 0000000..f9e051e --- /dev/null +++ b/pkg/cgroup/cgroup_info_cache.go @@ -0,0 +1,25 @@ +package cgroup + +import "sync" + +var ( + infoOnce sync.Once + cacheInfo map[string]int +) + +func getCachedCgroupHierarchy() map[string]int { + infoOnce.Do(func() { + info, err := GetCgroupInfo() + if err != nil { + return + } + cacheInfo = make(map[string]int) + for k, v := range info { + if !v.Enabled { + continue + } + cacheInfo[k] = v.Hierarchy + } + }) + return cacheInfo +} diff --git a/pkg/cgroup/consts.go b/pkg/cgroup/consts.go index b25c3ba..1fe3a8a 100644 --- a/pkg/cgroup/consts.go +++ b/pkg/cgroup/consts.go @@ -7,4 +7,5 @@ const ( procCgroupsPath = "/proc/cgroups" filePerm = 0644 + dirPerm = 0755 ) diff --git a/pkg/cgroup/doc.go b/pkg/cgroup/doc.go index 2712f8e..967f92d 100644 --- a/pkg/cgroup/doc.go +++ b/pkg/cgroup/doc.go @@ -2,12 +2,13 @@ // under systemd defined mount path (i.e.,sys/fs/cgroup). // // Current available: +// cpu +// cpuset // cpuacct // memory // pids -// cpuset // -// Current not available: cpu, cpuset, devices, freezer, net_cls, blkio, perf_event, net_prio, huge_tlb, rdma +// Current not available: devices, freezer, net_cls, blkio, perf_event, net_prio, huge_tlb, rdma // // Additional ideas: // diff --git a/pkg/cgroup/utils.go b/pkg/cgroup/utils.go index b296892..96686b5 100644 --- a/pkg/cgroup/utils.go +++ b/pkg/cgroup/utils.go @@ -1,17 +1,15 @@ package cgroup import ( - "bufio" "io/ioutil" "os" "path" - "strings" ) // EnsureDirExists creates directories if the path not exists func EnsureDirExists(path string) error { if _, err := os.Stat(path); os.IsNotExist(err) { - return os.MkdirAll(path, os.ModePerm) + return os.MkdirAll(path, dirPerm) } return nil } @@ -22,28 +20,3 @@ func CreateSubCgroupPath(group, prefix string) (string, error) { EnsureDirExists(base) return ioutil.TempDir(base, "") } - -// GetAllSubCgroup reads /proc/cgroups and get all available sub-cgroup as set -func GetAllSubCgroup() (map[string]bool, error) { - f, err := os.Open(procCgroupsPath) - if err != nil { - return nil, err - } - defer f.Close() - - rt := make(map[string]bool) - s := bufio.NewScanner(f) - for s.Scan() { - text := s.Text() - if text[0] != '#' { - parts := strings.Fields(text) - if len(parts) >= 4 && parts[3] != "0" { - rt[parts[0]] = true - } - } - } - if err := s.Err(); err != nil { - return nil, err - } - return rt, nil -}