diff --git a/pkg/cgroup/benchmark_test.go b/pkg/cgroup/benchmark_test.go index ce23c4b..0b388db 100644 --- a/pkg/cgroup/benchmark_test.go +++ b/pkg/cgroup/benchmark_test.go @@ -6,68 +6,64 @@ import ( ) func BenchmarkCgroup(b *testing.B) { - builder, err := NewBuilder("benchmark").WithCPUAcct().WithMemory().WithPids().FilterByEnv() + builder, err := NewBuilder("benchmark").WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() if err != nil { - b.Error(err) - return + b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { cg, err := builder.Build() if err != nil { - b.Error(err) - return + b.Fatal(err) + } + if err := cg.SetCpusetCpus([]byte("0")); err != nil { + b.Fatal(err) } if err := cg.SetMemoryLimitInBytes(4096); err != nil { - b.Error(err) - return + b.Fatal(err) } if err := cg.SetPidsMax(1); err != nil { - b.Error(err) - return + b.Fatal(err) } if _, err := cg.CpuacctUsage(); err != nil { - b.Error(err) - return + b.Fatal(err) } if _, err := cg.MemoryMaxUsageInBytes(); err != nil { - b.Error(err) - return + b.Fatal(err) } cg.Destroy() } } -func TestCgroup(t *testing.T) { +func TestCgroupAll(t *testing.T) { // ensure root privillege when testing - if os.Getuid() != 1 { + if os.Getuid() != 0 { t.Skip("no root privillege") } - builder, err := NewBuilder("test").WithCPUAcct().WithMemory().WithPids().FilterByEnv() + builder, err := NewBuilder("test").WithCPUSet().WithCPUAcct().WithMemory().WithPids().FilterByEnv() if err != nil { - t.Error(err) - return + t.Fatal(err) } cg, err := builder.Build() if err != nil { - t.Error(err) - return + t.Fatal(err) + } + t.Cleanup(func() { + cg.Destroy() + }) + if err := cg.SetCpusetCpus([]byte("0")); err != nil { + t.Fatal(err) } if err := cg.SetMemoryLimitInBytes(4096); err != nil { - t.Error(err) - return + t.Fatal(err) } if err := cg.SetPidsMax(1); err != nil { - t.Error(err) - return + t.Fatal(err) } if _, err := cg.CpuacctUsage(); err != nil { - t.Error(err) - return + t.Fatal(err) } if _, err := cg.MemoryMaxUsageInBytes(); err != nil { - t.Error(err) - return + t.Fatal(err) } - cg.Destroy() } diff --git a/pkg/cgroup/builder.go b/pkg/cgroup/builder.go index cd6dff6..ab7290c 100644 --- a/pkg/cgroup/builder.go +++ b/pkg/cgroup/builder.go @@ -8,8 +8,12 @@ import ( // Builder builds cgroup directories // available: cpuacct, memory, pids type Builder struct { - Prefix string - CPUAcct, Memory, Pids bool + Prefix string + + CPUSet bool + CPUAcct bool + Memory bool + Pids bool } // NewBuilder return a dumb builder without any sub-cgroup @@ -19,6 +23,12 @@ func NewBuilder(prefix string) *Builder { } } +// WithCPUSet includes cpuset cgroup +func (b *Builder) WithCPUSet() *Builder { + b.CPUSet = true + return b +} + // WithCPUAcct includes cpuacct cgroup func (b *Builder) WithCPUAcct() *Builder { b.CPUAcct = true @@ -43,6 +53,7 @@ func (b *Builder) FilterByEnv() (*Builder, error) { if err != nil { return b, err } + b.CPUSet = b.CPUSet && m["cpuset"] b.CPUAcct = b.CPUAcct && m["cpuacct"] b.Memory = b.Memory && m["memory"] b.Pids = b.Pids && m["pids"] @@ -56,6 +67,7 @@ func (b *Builder) String() string { name string enabled bool }{ + {"cpuset", b.CPUSet}, {"cpuacct", b.CPUAcct}, {"memory", b.Memory}, {"pids", b.Pids}, diff --git a/pkg/cgroup/cgroup.go b/pkg/cgroup/cgroup.go index 80992c9..26b3ad8 100644 --- a/pkg/cgroup/cgroup.go +++ b/pkg/cgroup/cgroup.go @@ -3,46 +3,63 @@ package cgroup import ( "bytes" "fmt" + "io/ioutil" "os" + "path/filepath" + "strings" ) // Cgroup is the combination of sub-cgroups type Cgroup struct { - prefix string - cpuacct, memory, pids *SubCgroup + prefix string + + cpuset *SubCgroup + cpuacct *SubCgroup + memory *SubCgroup + pids *SubCgroup } // Build creates new cgrouup directories func (b *Builder) Build() (cg *Cgroup, err error) { var ( - cpuacctPath, memoryPath, pidsPath string + cpuSetPath, cpuacctPath, memoryPath, pidsPath string ) // if failed, remove potential created directory defer func() { if err != nil { + remove(cpuSetPath) remove(cpuacctPath) remove(memoryPath) remove(pidsPath) } }() - if b.CPUAcct { - if cpuacctPath, err = CreateSubCgroupPath("cpuacct", b.Prefix); err != nil { + 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}, + } { + if !c.enable { + continue + } + if *c.path, err = CreateSubCgroupPath(c.name, b.Prefix); err != nil { return } } - if b.Memory { - if memoryPath, err = CreateSubCgroupPath("memory", b.Prefix); err != nil { - return - } - } - if b.Pids { - if pidsPath, err = CreateSubCgroupPath("pids", b.Prefix); err != nil { + + if b.CPUSet { + if err = initCpuset(cpuSetPath); err != nil { return } } return &Cgroup{ prefix: b.Prefix, + cpuset: NewSubCgroup(cpuSetPath), cpuacct: NewSubCgroup(cpuacctPath), memory: NewSubCgroup(memoryPath), pids: NewSubCgroup(pidsPath), @@ -51,14 +68,10 @@ func (b *Builder) Build() (cg *Cgroup, err error) { // AddProc writes cgroup.procs to all sub-cgroup func (c *Cgroup) AddProc(pid int) error { - if err := c.cpuacct.WriteUint(cgroupProcs, uint64(pid)); err != nil { - return err - } - if err := c.memory.WriteUint(cgroupProcs, uint64(pid)); err != nil { - return err - } - if err := c.pids.WriteUint(cgroupProcs, uint64(pid)); err != nil { - return err + for _, s := range []*SubCgroup{c.cpuset, c.cpuset, c.memory, c.pids} { + if err := s.WriteUint(cgroupProcs, uint64(pid)); err != nil { + return err + } } return nil } @@ -66,18 +79,24 @@ 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 - if err := remove(c.cpuacct.path); err != nil { - err1 = err - } - if err := remove(c.memory.path); err != nil { - err1 = err - } - if err := remove(c.pids.path); err != nil { - err1 = err + for _, s := range []*SubCgroup{c.cpuset, c.cpuset, c.memory, c.pids} { + if err := remove(s.path); err != nil { + err1 = err + } } return err1 } +// SetCpusetCpus set cpuset.cpus +func (c *Cgroup) SetCpusetCpus(b []byte) error { + return c.cpuset.WriteFile("cpuset.cpus", b) +} + +// SetCpusetMems set cpuset.mems +func (c *Cgroup) SetCpusetMems(b []byte) error { + return c.cpuset.WriteFile("cpuset.mems", b) +} + // CpuacctUsage read cpuacct.usage in ns func (c *Cgroup) CpuacctUsage() (uint64, error) { return c.cpuacct.ReadUint("cpuacct.usage") @@ -128,6 +147,36 @@ func (c *Cgroup) FindMemoryStatProperty(prop string) (uint64, error) { } } +// initCpuset will copy the config from the parent cpu sets if not exists +func initCpuset(path string) error { + for _, f := range []string{"cpuset.cpus", "cpuset.mems"} { + if err := copyCgroupPropertyFromParent(path, f); err != nil { + return err + } + } + return nil +} + +func copyCgroupPropertyFromParent(path, name string) error { + // ensure current one empty + b, err := ioutil.ReadFile(filepath.Join(path, name)) + if err != nil { + return err + } + if len(strings.TrimSpace(string(b))) > 0 { + return nil + } + // otherwise copy from parent, first to ensure it is empty by recurssion + if err := copyCgroupPropertyFromParent(filepath.Dir(path), name); err != nil { + return err + } + b, err = ioutil.ReadFile(filepath.Join(filepath.Dir(path), name)) + if err != nil { + return err + } + return ioutil.WriteFile(filepath.Join(path, name), b, filePerm) +} + func remove(name string) error { if name != "" { return os.Remove(name) diff --git a/pkg/cgroup/consts.go b/pkg/cgroup/consts.go index 11c04f0..b25c3ba 100644 --- a/pkg/cgroup/consts.go +++ b/pkg/cgroup/consts.go @@ -5,4 +5,6 @@ const ( basePath = "/sys/fs/cgroup" cgroupProcs = "cgroup.procs" procCgroupsPath = "/proc/cgroups" + + filePerm = 0644 ) diff --git a/pkg/cgroup/doc.go b/pkg/cgroup/doc.go index b65ee35..2712f8e 100644 --- a/pkg/cgroup/doc.go +++ b/pkg/cgroup/doc.go @@ -5,6 +5,7 @@ // cpuacct // memory // pids +// cpuset // // Current not available: cpu, cpuset, devices, freezer, net_cls, blkio, perf_event, net_prio, huge_tlb, rdma // diff --git a/pkg/cgroup/subcgroup.go b/pkg/cgroup/subcgroup.go index 28bb7fa..d424069 100644 --- a/pkg/cgroup/subcgroup.go +++ b/pkg/cgroup/subcgroup.go @@ -49,10 +49,13 @@ func (c *SubCgroup) ReadUint(filename string) (uint64, error) { // WriteFile writes cgroup file and handles potential EINTR error while writes to // the slow device (cgroup) func (c *SubCgroup) WriteFile(name string, content []byte) error { + if c.path == "" { + return ErrNotInitialized + } p := path.Join(c.path, name) - err := ioutil.WriteFile(p, content, 0664) + err := ioutil.WriteFile(p, content, filePerm) for err != nil && errors.Is(err, syscall.EINTR) { - err = ioutil.WriteFile(p, content, 0664) + err = ioutil.WriteFile(p, content, filePerm) } return err } @@ -60,6 +63,9 @@ func (c *SubCgroup) WriteFile(name string, content []byte) error { // ReadFile reads cgroup file and handles potential EINTR error while read to // the slow device (cgroup) func (c *SubCgroup) ReadFile(name string) ([]byte, error) { + if c.path == "" { + return nil, nil + } p := path.Join(c.path, name) data, err := ioutil.ReadFile(p) for err != nil && errors.Is(err, syscall.EINTR) { diff --git a/pkg/cgroup/utils.go b/pkg/cgroup/utils.go index 8c97168..b296892 100644 --- a/pkg/cgroup/utils.go +++ b/pkg/cgroup/utils.go @@ -11,7 +11,7 @@ import ( // EnsureDirExists creates directories if the path not exists func EnsureDirExists(path string) error { if _, err := os.Stat(path); os.IsNotExist(err) { - return os.Mkdir(path, os.ModePerm) + return os.MkdirAll(path, os.ModePerm) } return nil }