From 198fcc3c86ad72176c963f5c24136b7106dea192 Mon Sep 17 00:00:00 2001
From: Boming Zhang <bomingzh@sjtu.edu.cn>
Date: Sat, 11 Jan 2025 20:38:08 -0500
Subject: [PATCH] feat(executor/local): rlimit on CPU, Memory, Stack

---
 internal/executor/local/executor.go | 38 ++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/internal/executor/local/executor.go b/internal/executor/local/executor.go
index 78fe467..2051816 100644
--- a/internal/executor/local/executor.go
+++ b/internal/executor/local/executor.go
@@ -16,6 +16,32 @@ import (
 
 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) {
 	var results []stage.ExecutorResult
 
@@ -28,6 +54,16 @@ func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
 		}
 		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.Content != nil {
 				execCmd.Stdin = strings.NewReader(*cmd.Stdin.Content)
@@ -45,7 +81,7 @@ func (e *Local) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
 		execCmd.Stderr = &stderrBuffer
 
 		startTime := time.Now()
-		err := execCmd.Start()
+		err = execCmd.Start()
 		if err != nil {
 			return nil, fmt.Errorf("failed to start command: %v", err)
 		}