feat(executor/local): rlimit on CPU, Memory, Stack
All checks were successful
submodules sync / sync (push) Successful in 51s
build / build (push) Successful in 1m47s
build / trigger-build-image (push) Successful in 9s

This commit is contained in:
张泊明518370910136 2025-01-11 20:38:08 -05:00
parent c1d9111ea9
commit 198fcc3c86
GPG Key ID: D47306D7062CDA9D

View File

@ -16,6 +16,32 @@ import (
type Local struct{} type Local struct{}
func ToRlimit(c stage.Cmd) (map[int]syscall.Rlimit, error) {
limits := make(map[int]syscall.Rlimit)
if c.CPULimit > 0 {
// ns to s
timeLimit := (uint64(c.CPULimit) + 1e9 - 1) / 1e9
if timeLimit < 1 {
timeLimit = 1
}
limits[syscall.RLIMIT_CPU] = syscall.Rlimit{
Cur: timeLimit, Max: timeLimit,
}
}
if c.MemoryLimit > 0 {
limits[syscall.RLIMIT_AS] = syscall.Rlimit{
Cur: c.MemoryLimit, Max: c.MemoryLimit, // bytes
}
}
if c.StackLimit > 0 {
limits[syscall.RLIMIT_STACK] = syscall.Rlimit{
Cur: c.StackLimit, Max: c.StackLimit, // bytes
}
}
return limits, nil
}
func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
var results []stage.ExecutorResult var results []stage.ExecutorResult
@ -28,6 +54,16 @@ func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
} }
execCmd.Env = env execCmd.Env = env
limits, err := ToRlimit(cmd)
if err != nil {
return nil, fmt.Errorf("failed to convert rlimits: %v", err)
}
for resource, limit := range limits {
if err := syscall.Setrlimit(resource, &limit); err != nil {
return nil, fmt.Errorf("failed to set rlimit: %v", err)
}
}
if cmd.Stdin != nil { if cmd.Stdin != nil {
if cmd.Stdin.Content != nil { if cmd.Stdin.Content != nil {
execCmd.Stdin = strings.NewReader(*cmd.Stdin.Content) execCmd.Stdin = strings.NewReader(*cmd.Stdin.Content)
@ -45,7 +81,7 @@ func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
execCmd.Stderr = &stderrBuffer execCmd.Stderr = &stderrBuffer
startTime := time.Now() startTime := time.Now()
err := execCmd.Start() err = execCmd.Start()
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to start command: %v", err) return nil, fmt.Errorf("failed to start command: %v", err)
} }