Add cpu cgroup support

This commit is contained in:
criyle 2020-12-12 21:28:41 -08:00
parent 4ae1d6d760
commit 5343e365c3
8 changed files with 180 additions and 52 deletions

View File

@ -6,7 +6,7 @@ import (
) )
func BenchmarkCgroup(b *testing.B) { 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 { if err != nil {
b.Fatal(err) b.Fatal(err)
} }
@ -36,11 +36,11 @@ func BenchmarkCgroup(b *testing.B) {
} }
func TestCgroupAll(t *testing.T) { func TestCgroupAll(t *testing.T) {
// ensure root privillege when testing // ensure root privilege when testing
if os.Getuid() != 0 { if os.Getuid() != 0 {
t.Skip("no root privillege") 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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }

View File

@ -10,6 +10,7 @@ import (
type Builder struct { type Builder struct {
Prefix string Prefix string
CPU bool
CPUSet bool CPUSet bool
CPUAcct bool CPUAcct bool
Memory 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 // WithCPUSet includes cpuset cgroup
func (b *Builder) WithCPUSet() *Builder { func (b *Builder) WithCPUSet() *Builder {
b.CPUSet = true b.CPUSet = true

View File

@ -13,62 +13,77 @@ import (
type Cgroup struct { type Cgroup struct {
prefix string prefix string
cpu *SubCgroup
cpuset *SubCgroup cpuset *SubCgroup
cpuacct *SubCgroup cpuacct *SubCgroup
memory *SubCgroup memory *SubCgroup
pids *SubCgroup pids *SubCgroup
all []*SubCgroup
} }
// Build creates new cgrouup directories // Build creates new cgrouup directories
func (b *Builder) Build() (cg *Cgroup, err error) { func (b *Builder) Build() (cg *Cgroup, err error) {
var ( var (
cpuSetPath, cpuacctPath, memoryPath, pidsPath string paths []string
subCgroup = make(map[int]*SubCgroup)
info = getCachedCgroupHierarchy()
) )
// if failed, remove potential created directory // if failed, remove potential created directory
defer func() { defer func() {
if err != nil { if err != nil {
remove(cpuSetPath) removeAll(paths...)
remove(cpuacctPath)
remove(memoryPath)
remove(pidsPath)
} }
}() }()
for _, c := range []struct { for _, c := range []struct {
enable bool enable bool
name string name string
path *string
}{ }{
{b.CPUSet, "cpuset", &cpuSetPath}, {b.CPU, "cpu"},
{b.CPUAcct, "cpuacct", &cpuacctPath}, {b.CPUSet, "cpuset"},
{b.Memory, "memory", &memoryPath}, {b.CPUAcct, "cpuacct"},
{b.Pids, "pids", &pidsPath}, {b.Memory, "memory"},
{b.Pids, "pids"},
} { } {
if !c.enable { if !c.enable {
continue continue
} }
if *c.path, err = CreateSubCgroupPath(c.name, b.Prefix); err != nil { h := info[c.name]
return 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 b.CPUSet {
if err = initCpuset(cpuSetPath); err != nil { if err = initCpuset(subCgroup[info["cpuset"]].path); err != nil {
return return
} }
} }
var all []*SubCgroup
for _, v := range subCgroup {
all = append(all, v)
}
return &Cgroup{ return &Cgroup{
prefix: b.Prefix, prefix: b.Prefix,
cpuset: NewSubCgroup(cpuSetPath), cpu: subCgroup[info["cpu"]],
cpuacct: NewSubCgroup(cpuacctPath), cpuset: subCgroup[info["cpuset"]],
memory: NewSubCgroup(memoryPath), cpuacct: subCgroup[info["cpuacct"]],
pids: NewSubCgroup(pidsPath), memory: subCgroup[info["memory"]],
pids: subCgroup[info["pids"]],
all: all,
}, nil }, nil
} }
// AddProc writes cgroup.procs to all sub-cgroup // AddProc writes cgroup.procs to all sub-cgroup
func (c *Cgroup) AddProc(pid int) error { 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 { if err := s.WriteUint(cgroupProcs, uint64(pid)); err != nil {
return err 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 // Destroy removes dir for sub-cgroup, errors are ignored if remove one failed
func (c *Cgroup) Destroy() error { func (c *Cgroup) Destroy() error {
var err1 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 { if err := remove(s.path); err != nil {
err1 = err err1 = err
} }
@ -87,6 +102,16 @@ func (c *Cgroup) Destroy() error {
return err1 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 // SetCpusetCpus set cpuset.cpus
func (c *Cgroup) SetCpusetCpus(b []byte) error { func (c *Cgroup) SetCpusetCpus(b []byte) error {
return c.cpuset.WriteFile("cpuset.cpus", b) 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) 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 // SetPidsMax write pids.max
func (c *Cgroup) SetPidsMax(i uint64) error { func (c *Cgroup) SetPidsMax(i uint64) error {
return c.pids.WriteUint("pids.max", i) return c.pids.WriteUint("pids.max", i)
@ -183,3 +218,14 @@ func remove(name string) error {
} }
return nil 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
}

75
pkg/cgroup/cgroup_info.go Normal file
View File

@ -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
}

View File

@ -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
}

View File

@ -7,4 +7,5 @@ const (
procCgroupsPath = "/proc/cgroups" procCgroupsPath = "/proc/cgroups"
filePerm = 0644 filePerm = 0644
dirPerm = 0755
) )

View File

@ -2,12 +2,13 @@
// under systemd defined mount path (i.e.,sys/fs/cgroup). // under systemd defined mount path (i.e.,sys/fs/cgroup).
// //
// Current available: // Current available:
// cpu
// cpuset
// cpuacct // cpuacct
// memory // memory
// pids // 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: // Additional ideas:
// //

View File

@ -1,17 +1,15 @@
package cgroup package cgroup
import ( import (
"bufio"
"io/ioutil" "io/ioutil"
"os" "os"
"path" "path"
"strings"
) )
// EnsureDirExists creates directories if the path not exists // EnsureDirExists creates directories if the path not exists
func EnsureDirExists(path string) error { func EnsureDirExists(path string) error {
if _, err := os.Stat(path); os.IsNotExist(err) { if _, err := os.Stat(path); os.IsNotExist(err) {
return os.MkdirAll(path, os.ModePerm) return os.MkdirAll(path, dirPerm)
} }
return nil return nil
} }
@ -22,28 +20,3 @@ func CreateSubCgroupPath(group, prefix string) (string, error) {
EnsureDirExists(base) EnsureDirExists(base)
return ioutil.TempDir(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
}