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) {
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)
}

View File

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

View File

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

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"
filePerm = 0644
dirPerm = 0755
)

View File

@ -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:
//

View File

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