From 7c57b248bbd7641b5e87d023d1be5a651878f9ad Mon Sep 17 00:00:00 2001 From: criyle Date: Sat, 3 Aug 2019 23:25:59 -0700 Subject: [PATCH] add cgroup v1 support for resource stat --- README.md | 3 +- cgroup/cgroup.go | 92 ++++++++++++++++++++++++++++++++++++ cgroup/consts.go | 7 +++ cgroup/subcgroup.go | 38 +++++++++++++++ cgroup/utils.go | 22 +++++++++ cmd/run_program/main.go | 45 +++++++++++++++++- runprogram/runprogram.go | 3 ++ runprogram/runprogram_run.go | 15 +++--- rununshared/run.go | 1 + rununshared/rununshared.go | 3 ++ tracer/tracer_track.go | 7 +-- 11 files changed, 220 insertions(+), 16 deletions(-) create mode 100644 cgroup/cgroup.go create mode 100644 cgroup/consts.go create mode 100644 cgroup/subcgroup.go create mode 100644 cgroup/utils.go diff --git a/README.md b/README.md index 049b88f..76ebef3 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ New Features: 3. Allow multiple traced programs in different threads 4. Allow pipes as input / output files 5. Use Linux Namespace to isolate file access (elimilate ptrace) +6. Use Linux Control Groups to limit & acct CPU & memory (elimilate wait4.rusage) Default file access action: @@ -66,4 +67,4 @@ It seems unshare net or ipc takes time, maybe limits action by seccomp instead. TODO: -1. Use Linux Control Groups to limit & acct CPU & memory (elimilate wait4 rusage) +1. Add ability to pre-fork container deamons diff --git a/cgroup/cgroup.go b/cgroup/cgroup.go new file mode 100644 index 0000000..426254d --- /dev/null +++ b/cgroup/cgroup.go @@ -0,0 +1,92 @@ +// Package cgroup provices basic resource control over cgroups +// it measure +// cpu: cpuacct.usage (ns) +// memory: memory.max_usage_in_bytes +// it limits: +// memory: memory.limit_in_bytes +// # of tasks: pids.max +package cgroup + +import "os" + +// additional ideas: +// cpu share(not used): cpu.share +// reclaim pages from old process: memory.force_empty +// (tasks kill are managed out of cgroup as freeze takes some time) +// freeze: freezer.state + +// CGroup is the combination of sub-cgroups +type CGroup struct { + prefix string + cpuacct, memory, pids *SubCGroup +} + +// NewCGroup creates new cgrouup directories +func NewCGroup(prefix string) (*CGroup, error) { + cpuacctPath, err := CreateSubCGroupPath("cpuacct", prefix) + if err != nil { + return nil, err + } + memoryPath, err := CreateSubCGroupPath("memory", prefix) + if err != nil { + return nil, err + } + pidsPath, err := CreateSubCGroupPath("pids", prefix) + if err != nil { + return nil, err + } + return &CGroup{ + prefix: prefix, + cpuacct: NewSubCGroup(cpuacctPath), + memory: NewSubCGroup(memoryPath), + pids: NewSubCGroup(pidsPath), + }, nil +} + +// 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 + } + return nil +} + +// Destroy removes dir for sub-cggroup +func (c *CGroup) Destroy() error { + if err := os.Remove(c.cpuacct.path); err != nil { + return err + } + if err := os.Remove(c.memory.path); err != nil { + return err + } + if err := os.Remove(c.pids.path); err != nil { + return err + } + return nil +} + +// CpuacctUsage read cpuacct.usage in ns +func (c *CGroup) CpuacctUsage() (uint64, error) { + return c.cpuacct.ReadUint("cpuacct.usage") +} + +// MemoryMaxUsageInBytes read memory.max_usage_in_bytes +func (c *CGroup) MemoryMaxUsageInBytes() (uint64, error) { + return c.memory.ReadUint("memory.max_usage_in_bytes") +} + +// SetMemoryLimitInBytes write memory.limit_in_bytes +func (c *CGroup) SetMemoryLimitInBytes(i uint64) error { + return c.memory.WriteUint("memory.limit_in_bytes", i) +} + +// SetPidsMax write pids.max +func (c *CGroup) SetPidsMax(i uint64) error { + return c.pids.WriteUint("pids.max", i) +} diff --git a/cgroup/consts.go b/cgroup/consts.go new file mode 100644 index 0000000..fffdb40 --- /dev/null +++ b/cgroup/consts.go @@ -0,0 +1,7 @@ +package cgroup + +const ( + // systemd mounted cgroups + basePath = "/sys/fs/cgroup" + cgroupProcs = "cgroup.procs" +) diff --git a/cgroup/subcgroup.go b/cgroup/subcgroup.go new file mode 100644 index 0000000..16e08c8 --- /dev/null +++ b/cgroup/subcgroup.go @@ -0,0 +1,38 @@ +package cgroup + +import ( + "io/ioutil" + "path" + "strconv" + "strings" +) + +// SubCGroup is the sub-cgroup +type SubCGroup struct { + path string +} + +// NewSubCGroup creates a sug CGroup +func NewSubCGroup(p string) *SubCGroup { + return &SubCGroup{ + path: p, + } +} + +// WriteUint writes uint64 into given file +func (c *SubCGroup) WriteUint(filename string, i uint64) error { + return ioutil.WriteFile(path.Join(c.path, filename), []byte(strconv.FormatUint(i, 10)), 644) +} + +// ReadUint read uint64 into given file +func (c *SubCGroup) ReadUint(filename string) (uint64, error) { + b, err := ioutil.ReadFile(path.Join(c.path, filename)) + if err != nil { + return 0, err + } + s, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64) + if err != nil { + return 0, err + } + return s, nil +} diff --git a/cgroup/utils.go b/cgroup/utils.go new file mode 100644 index 0000000..7fda2a9 --- /dev/null +++ b/cgroup/utils.go @@ -0,0 +1,22 @@ +package cgroup + +import ( + "io/ioutil" + "os" + "path" +) + +// EnsureDirExists creates dir if not exists +func EnsureDirExists(path string) error { + if _, err := os.Stat(path); os.IsNotExist(err) { + return os.Mkdir(path, os.ModePerm) + } + return nil +} + +// CreateSubCGroupPath creates path for sub-cgroup +func CreateSubCGroupPath(group, prefix string) (string, error) { + base := path.Join(basePath, group, prefix) + EnsureDirExists(base) + return ioutil.TempDir(base, "") +} diff --git a/cmd/run_program/main.go b/cmd/run_program/main.go index 0680388..a78c3d0 100644 --- a/cmd/run_program/main.go +++ b/cmd/run_program/main.go @@ -5,7 +5,9 @@ import ( "fmt" "io/ioutil" "os" + "time" + "github.com/criyle/go-judger/cgroup" "github.com/criyle/go-judger/runconfig" "github.com/criyle/go-judger/runprogram" "github.com/criyle/go-judger/rununshared" @@ -31,11 +33,13 @@ func printUsage() { func main() { var ( addReadable, addWritable, addRawReadable, addRawWritable arrayFlags - allowProc, unsafe, showDetails, namespace bool + allowProc, unsafe, showDetails, namespace, useCGroup bool pType, result string timeLimit, realTimeLimit, memoryLimit, outputLimit, stackLimit uint64 inputFileName, outputFileName, errorFileName, workPath string runner Runner + cg *cgroup.CGroup + err error ) flag.Usage = printUsage @@ -58,6 +62,7 @@ func main() { flag.Var(&addRawReadable, "add-readable-raw", "Add a readable file (don't transform to its real path)") flag.Var(&addRawWritable, "add-writable-raw", "Add a writable file (don't transform to its real path)") flag.BoolVar(&namespace, "ns", false, "Use namespace to restrict file accesses") + flag.BoolVar(&useCGroup, "cgroup", false, "Use cgroup to colloct resource usage") flag.Parse() args := flag.Args() @@ -85,6 +90,26 @@ func main() { addWrite := runconfig.GetExtraSet(addWritable, addRawWritable) h := runconfig.GetConf(pType, workPath, args, addRead, addWrite, allowProc, showDetails) + if useCGroup { + cg, err = cgroup.NewCGroup("run_program") + if err != nil { + panic(err) + } + defer cg.Destroy() + if err = cg.SetMemoryLimitInBytes(memoryLimit << 20); err != nil { + panic(err) + } + } + + syncFunc := func(pid int) error { + if cg != nil { + if err := cg.AddProc(pid); err != nil { + return err + } + } + return nil + } + // open input / output / err files files, err := prepareFiles(inputFileName, outputFileName, errorFileName) if err != nil { @@ -135,7 +160,8 @@ func main() { Target: "w", }, }), - ShowDetails: true, + ShowDetails: showDetails, + SyncFunc: syncFunc, } } else { runner = &runprogram.RunProgram{ @@ -154,6 +180,7 @@ func main() { ShowDetails: showDetails, Unsafe: unsafe, Handler: h, + SyncFunc: syncFunc, } } @@ -175,6 +202,20 @@ func main() { rt, err := runner.Start() println("results:", rt, err) + if useCGroup { + cpu, err := cg.CpuacctUsage() + if err != nil { + panic(err) + } + memory, err := cg.MemoryMaxUsageInBytes() + if err != nil { + panic(err) + } + println("cgroup: cpu: ", cpu, " memory: ", memory) + rt.UserTime = cpu / uint64(time.Millisecond) + rt.UserMem = memory >> 10 + } + if err != nil { c, ok := err.(specs.TraceCode) if !ok { diff --git a/runprogram/runprogram.go b/runprogram/runprogram.go index ef651c2..83ba96a 100644 --- a/runprogram/runprogram.go +++ b/runprogram/runprogram.go @@ -36,6 +36,9 @@ type RunProgram struct { // ShowDetails / Unsafe debug flag ShowDetails, Unsafe bool + + // Use by cgroup to add proc + SyncFunc func(pid int) error } // TraceAction defines action against a syscall check diff --git a/runprogram/runprogram_run.go b/runprogram/runprogram_run.go index 7aa8d65..960c86b 100644 --- a/runprogram/runprogram_run.go +++ b/runprogram/runprogram_run.go @@ -26,13 +26,14 @@ func (r *RunProgram) Start() (rt specs.TraceResult, err error) { } ch := &forkexec.Runner{ - Args: r.Args, - Env: r.Env, - RLimits: r.RLimits.PrepareRLimit(), - Files: r.Files, - WorkDir: r.WorkDir, - Seccomp: bpf, - Ptrace: true, + Args: r.Args, + Env: r.Env, + RLimits: r.RLimits.PrepareRLimit(), + Files: r.Files, + WorkDir: r.WorkDir, + Seccomp: bpf, + Ptrace: true, + SyncFunc: r.SyncFunc, } th := &tracerHandler{ diff --git a/rununshared/run.go b/rununshared/run.go index 7d2dfc8..3685dac 100644 --- a/rununshared/run.go +++ b/rununshared/run.go @@ -46,6 +46,7 @@ func (r *RunUnshared) Start() (rt specs.TraceResult, err error) { Mounts: r.Mounts, PivotRoot: r.Root, DropCaps: true, + SyncFunc: r.SyncFunc, } return r.Trace(ch) } diff --git a/rununshared/rununshared.go b/rununshared/rununshared.go index b4945f3..1a48a1c 100644 --- a/rununshared/rununshared.go +++ b/rununshared/rununshared.go @@ -35,4 +35,7 @@ type RunUnshared struct { // Show Details ShowDetails bool + + // Use by cgroup to add proc + SyncFunc func(pid int) error } diff --git a/tracer/tracer_track.go b/tracer/tracer_track.go index 7f97d6a..849f3ff 100644 --- a/tracer/tracer_track.go +++ b/tracer/tracer_track.go @@ -26,7 +26,6 @@ func Trace(handler Handler, runner Runner, limits specs.ResLimit) (result specs. traced = make(map[int]bool) // store all process that have set ptrace options execved = false // store whether the runner process have successfully execvd pid int // store pid of wait4 result - initMem uint64 // initial memory usage (likely due to original process) sTime = time.Now().UnixNano() // records start time for trace process fTime int64 // records finish time for execve ) @@ -85,13 +84,9 @@ func Trace(handler Handler, runner Runner, limits specs.ResLimit) (result specs. status := specs.TraceCodeNormal if pid == pgid { - if initMem == 0 { - initMem = uint64(rusage.Maxrss) - handler.Debug("initial memory:", initMem) - } // update resource usage and check against limits userTime := uint64(rusage.Utime.Sec*1e3 + rusage.Utime.Usec/1e3) // ms - userMem := uint64(rusage.Maxrss) - initMem // kb + userMem := uint64(rusage.Maxrss) // kb // check tle / mle if userTime > limits.TimeLimit {