add cgroup v1 support for resource stat

This commit is contained in:
criyle 2019-08-03 23:25:59 -07:00
parent 144c38f0e8
commit 7c57b248bb
11 changed files with 220 additions and 16 deletions

View File

@ -23,6 +23,7 @@ New Features:
3. Allow multiple traced programs in different threads 3. Allow multiple traced programs in different threads
4. Allow pipes as input / output files 4. Allow pipes as input / output files
5. Use Linux Namespace to isolate file access (elimilate ptrace) 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: Default file access action:
@ -66,4 +67,4 @@ It seems unshare net or ipc takes time, maybe limits action by seccomp instead.
TODO: TODO:
1. Use Linux Control Groups to limit & acct CPU & memory (elimilate wait4 rusage) 1. Add ability to pre-fork container deamons

92
cgroup/cgroup.go Normal file
View File

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

7
cgroup/consts.go Normal file
View File

@ -0,0 +1,7 @@
package cgroup
const (
// systemd mounted cgroups
basePath = "/sys/fs/cgroup"
cgroupProcs = "cgroup.procs"
)

38
cgroup/subcgroup.go Normal file
View File

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

22
cgroup/utils.go Normal file
View File

@ -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, "")
}

View File

@ -5,7 +5,9 @@ import (
"fmt" "fmt"
"io/ioutil" "io/ioutil"
"os" "os"
"time"
"github.com/criyle/go-judger/cgroup"
"github.com/criyle/go-judger/runconfig" "github.com/criyle/go-judger/runconfig"
"github.com/criyle/go-judger/runprogram" "github.com/criyle/go-judger/runprogram"
"github.com/criyle/go-judger/rununshared" "github.com/criyle/go-judger/rununshared"
@ -31,11 +33,13 @@ func printUsage() {
func main() { func main() {
var ( var (
addReadable, addWritable, addRawReadable, addRawWritable arrayFlags addReadable, addWritable, addRawReadable, addRawWritable arrayFlags
allowProc, unsafe, showDetails, namespace bool allowProc, unsafe, showDetails, namespace, useCGroup bool
pType, result string pType, result string
timeLimit, realTimeLimit, memoryLimit, outputLimit, stackLimit uint64 timeLimit, realTimeLimit, memoryLimit, outputLimit, stackLimit uint64
inputFileName, outputFileName, errorFileName, workPath string inputFileName, outputFileName, errorFileName, workPath string
runner Runner runner Runner
cg *cgroup.CGroup
err error
) )
flag.Usage = printUsage 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(&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.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(&namespace, "ns", false, "Use namespace to restrict file accesses")
flag.BoolVar(&useCGroup, "cgroup", false, "Use cgroup to colloct resource usage")
flag.Parse() flag.Parse()
args := flag.Args() args := flag.Args()
@ -85,6 +90,26 @@ func main() {
addWrite := runconfig.GetExtraSet(addWritable, addRawWritable) addWrite := runconfig.GetExtraSet(addWritable, addRawWritable)
h := runconfig.GetConf(pType, workPath, args, addRead, addWrite, allowProc, showDetails) 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 // open input / output / err files
files, err := prepareFiles(inputFileName, outputFileName, errorFileName) files, err := prepareFiles(inputFileName, outputFileName, errorFileName)
if err != nil { if err != nil {
@ -135,7 +160,8 @@ func main() {
Target: "w", Target: "w",
}, },
}), }),
ShowDetails: true, ShowDetails: showDetails,
SyncFunc: syncFunc,
} }
} else { } else {
runner = &runprogram.RunProgram{ runner = &runprogram.RunProgram{
@ -154,6 +180,7 @@ func main() {
ShowDetails: showDetails, ShowDetails: showDetails,
Unsafe: unsafe, Unsafe: unsafe,
Handler: h, Handler: h,
SyncFunc: syncFunc,
} }
} }
@ -175,6 +202,20 @@ func main() {
rt, err := runner.Start() rt, err := runner.Start()
println("results:", rt, err) 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 { if err != nil {
c, ok := err.(specs.TraceCode) c, ok := err.(specs.TraceCode)
if !ok { if !ok {

View File

@ -36,6 +36,9 @@ type RunProgram struct {
// ShowDetails / Unsafe debug flag // ShowDetails / Unsafe debug flag
ShowDetails, Unsafe bool ShowDetails, Unsafe bool
// Use by cgroup to add proc
SyncFunc func(pid int) error
} }
// TraceAction defines action against a syscall check // TraceAction defines action against a syscall check

View File

@ -26,13 +26,14 @@ func (r *RunProgram) Start() (rt specs.TraceResult, err error) {
} }
ch := &forkexec.Runner{ ch := &forkexec.Runner{
Args: r.Args, Args: r.Args,
Env: r.Env, Env: r.Env,
RLimits: r.RLimits.PrepareRLimit(), RLimits: r.RLimits.PrepareRLimit(),
Files: r.Files, Files: r.Files,
WorkDir: r.WorkDir, WorkDir: r.WorkDir,
Seccomp: bpf, Seccomp: bpf,
Ptrace: true, Ptrace: true,
SyncFunc: r.SyncFunc,
} }
th := &tracerHandler{ th := &tracerHandler{

View File

@ -46,6 +46,7 @@ func (r *RunUnshared) Start() (rt specs.TraceResult, err error) {
Mounts: r.Mounts, Mounts: r.Mounts,
PivotRoot: r.Root, PivotRoot: r.Root,
DropCaps: true, DropCaps: true,
SyncFunc: r.SyncFunc,
} }
return r.Trace(ch) return r.Trace(ch)
} }

View File

@ -35,4 +35,7 @@ type RunUnshared struct {
// Show Details // Show Details
ShowDetails bool ShowDetails bool
// Use by cgroup to add proc
SyncFunc func(pid int) error
} }

View File

@ -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 traced = make(map[int]bool) // store all process that have set ptrace options
execved = false // store whether the runner process have successfully execvd execved = false // store whether the runner process have successfully execvd
pid int // store pid of wait4 result 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 sTime = time.Now().UnixNano() // records start time for trace process
fTime int64 // records finish time for execve 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 status := specs.TraceCodeNormal
if pid == pgid { if pid == pgid {
if initMem == 0 {
initMem = uint64(rusage.Maxrss)
handler.Debug("initial memory:", initMem)
}
// update resource usage and check against limits // update resource usage and check against limits
userTime := uint64(rusage.Utime.Sec*1e3 + rusage.Utime.Usec/1e3) // ms 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 // check tle / mle
if userTime > limits.TimeLimit { if userTime > limits.TimeLimit {