open-lambda · m0mosenpai
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -55,7 +55,7 @@ jobs:
       run: cargo install cross --locked
     - name: Build OpenLambda
       run: |
-        make ol imgs/ol-wasm wasm-worker wasm-functions native-functions container-proxy
+        make ol-bin imgs/ol-wasm wasm-worker wasm-functions native-functions container-proxy
         make install-python-bindings
         make sudo-install
     - name: Test Go Common

diff --git a/.github/workflows/pkg.yml b/.github/workflows/pkg.yml
@@ -47,7 +47,7 @@ jobs:
       run: cargo install cross --locked
     - name: Build OpenLambda
       run: |
-        make ol imgs/lambda wasm-worker wasm-functions native-functions container-proxy
+        make ol-bin imgs/lambda wasm-worker wasm-functions native-functions container-proxy
         sudo make install-python-bindings
     - name: Test PyPI Packages
       run: sudo ./scripts/package_tests.py

diff --git a/.gitignore b/.gitignore
@@ -60,7 +60,7 @@
 /test.json
 
 # binaries
-ol
+ol-bin
 ol-wasm
 ol-container-proxy
 

diff --git a/Makefile b/Makefile
@@ -35,7 +35,7 @@ endif
 .PHONY: container-proxy
 .PHONY: fmt check-fmt
 
-all: ol imgs/ol-wasm wasm-worker wasm-functions native-functions container-proxy
+all: ol-bin imgs/ol-wasm wasm-worker wasm-functions native-functions container-proxy
 
 wasm-worker:
 	cd wasm-worker && cargo build ${BUILD_FLAGS}
@@ -76,19 +76,21 @@ container-proxy:
 	cd container-proxy && cargo build ${BUILD_FLAGS}
 	cp ./container-proxy/target/${BUILDTYPE}/open-lambda-container-proxy ./ol-container-proxy
 
-ol: ${OL_GO_FILES}
-	cd ${OL_DIR} && ${GO} build -o ../ol
+ol-bin: ${OL_GO_FILES}
+	cd ${OL_DIR} && ${GO} build -o ../ol-bin
 
-build: ol wasm-worker container-proxy
+build: ol-bin wasm-worker container-proxy
 
 install: build
 	cp ol ${INSTALL_PREFIX}/bin/
+	cp ol-bin ${INSTALL_PREFIX}/bin/
 	cp ol-wasm ${INSTALL_PREFIX}/bin/
 	cp ol-container-proxy ${INSTALL_PREFIX}/bin/
 	cp autocomplete/bash_autocomplete /etc/bash_completion.d/ol 
 
 sudo-install: build
 	sudo cp ol ${INSTALL_PREFIX}/bin/
+	sudo cp ol-bin ${INSTALL_PREFIX}/bin/
 	sudo cp ol-wasm ${INSTALL_PREFIX}/bin/
 	sudo cp ol-container-proxy ${INSTALL_PREFIX}/bin/
 	sudo cp autocomplete/bash_autocomplete /etc/bash_completion.d/ol 
@@ -127,6 +129,6 @@ lint-wasm-worker:
 lint: lint-wasm-worker lint-functions lint-python lint-go
 
 clean:
-	rm -f ol imgs/ol-min imgs/ol-wasm
+	rm -f ol-bin imgs/ol-min imgs/ol-wasm
 	${MAKE} -C lambda clean
 	${MAKE} -C sock clean
diff --git a/go/common/cgroup.go b/go/common/cgroup.go
@@ -0,0 +1,35 @@
+package common
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+// EnsureCgroupConfigured fails if neither cgroup_root nor OL_SYSTEMD is set.
+func EnsureCgroupConfigured() error {
+	if Conf != nil && Conf.Cgroup_root != "" {
+		return nil
+	}
+	if os.Getenv("OL_SYSTEMD") == "1" {
+		return nil
+	}
+	return fmt.Errorf("no cgroup root: set cgroup_root in config or run the worker via 'ol'")
+}
+
+func CgroupRoot() (string, error) {
+	if Conf != nil && Conf.Cgroup_root != "" {
+		return Conf.Cgroup_root, nil
+	}
+	data, err := os.ReadFile("/proc/self/cgroup")
+	if err != nil {
+		return "", fmt.Errorf("read /proc/self/cgroup: %w", err)
+	}
+	line := strings.TrimSpace(string(data))
+	parts := strings.SplitN(line, "::", 2)
+	if len(parts) != 2 {
+		return "", fmt.Errorf("unexpected /proc/self/cgroup format: %q", line)
+	}
+	return filepath.Join("/sys/fs/cgroup", parts[1]), nil
+}
diff --git a/go/common/config.go b/go/common/config.go
@@ -21,6 +21,9 @@ type Config struct {
 	// worker directory, which contains handler code, pid file, logs, etc.
 	Worker_dir string `json:"worker_dir"`
 
+	// configure the cgroup pool root - defaults to '/proc/self/cgroup'
+	Cgroup_root string `json:"cgroup_root"`
+
 	// Url/ip the worker server listens to
 	Worker_url string `json:"worker_url"`
 
@@ -492,8 +495,3 @@ func GetOlPath(ctx *cli.Context) (string, error) {
 	}
 	return filepath.Abs(olPath)
 }
-
-// CgroupPoolPath returns the cgroup pool root path for the given OL directory.
-func CgroupPoolPath(olPath string) string {
-	return filepath.Join("/sys/fs/cgroup", filepath.Base(olPath)+"-sandboxes")
-}
diff --git a/go/worker/commands.go b/go/worker/commands.go
@@ -16,7 +16,6 @@ import (
 
 	"github.com/open-lambda/open-lambda/go/common"
 	"github.com/open-lambda/open-lambda/go/worker/event"
-	"github.com/open-lambda/open-lambda/go/worker/sandbox/cgroups"
 
 	"github.com/urfave/cli/v2"
 )
@@ -40,27 +39,18 @@ func udsGet(requestPath string) (*http.Response, error) {
 
 // initCmd corresponds to the "init" command of the admin tool.
 func initCmd(ctx *cli.Context) error {
-	if os.Getuid() != 0 {
-		return fmt.Errorf("'ol worker init' must be run with sudo")
-	}
-
 	olPath, err := common.GetOlPath(ctx)
 	if err != nil {
-		return fmt.Errorf("init failed to get OL path: %w", err)
+		return err
 	}
 
 	if err := common.LoadDefaults(olPath); err != nil {
-		return fmt.Errorf("init failed to load config defaults: %w", err)
+		return err
 	}
 
 	if err := initOLDir(olPath, ctx.String("image"), ctx.Bool("newbase")); err != nil {
-		return fmt.Errorf("init failed to create OL directory: %w", err)
-	}
-
-	if err := cgroups.InitPoolRoot(common.CgroupPoolPath(olPath)); err != nil {
-		return fmt.Errorf("init failed to create cgroup pool: %w", err)
+		return err
 	}
-
 	fmt.Printf("\nYou may optionally modify the defaults here: %s\n\n",
 		filepath.Join(olPath, "config.json"))
 	fmt.Printf("Next start a worker using the \"ol worker up\" command.\n")

diff --git a/go/worker/helpers.go b/go/worker/helpers.go
@@ -21,7 +21,6 @@ import (
 	"github.com/open-lambda/open-lambda/go/worker/embedded"
 )
 
-
 func initOLBaseDir(baseDir string, dockerBaseImage string) error {
 	if dockerBaseImage == "" {
 		dockerBaseImage = "ol-wasm"
@@ -280,8 +279,14 @@ func runningToStoppedClean() error {
 // It cleans up cgroups and mounts associated with the OpenLambda instance at `olPath`.
 // Returns errors encountered during cleanup operations.
 func stoppedDirtyToStoppedClean(olPath string) error {
-	// Clean up child cgroups, preserving the pool root
-	cgRoot := common.CgroupPoolPath(olPath)
+	// Clean up cgroups associated with sandboxes
+	if err := common.EnsureCgroupConfigured(); err != nil {
+		return err
+	}
+	cgRoot, err := common.CgroupRoot()
+	if err != nil {
+		return fmt.Errorf("failed to resolve cgroup root: %s", err)
+	}
 	fmt.Printf("Attempting to clean up cgroups at %s\n", cgRoot)
 
 	cgroupErrorCount := 0
@@ -301,21 +306,25 @@ func stoppedDirtyToStoppedClean(olPath string) error {
 		if err != nil {
 			return fmt.Errorf("error reading cgroup root: %s", err.Error())
 		}
-		kill := filepath.Join(cgRoot, "cgroup.kill")
-		if err := os.WriteFile(kill, []byte(fmt.Sprintf("%d", 1)), os.ModeAppend); err != nil {
-			fmt.Printf("Could not kill processes in cgroup: %s\n", err.Error())
-			cgroupErrorCount += 1
-		}
 		for _, file := range files {
 			if strings.HasPrefix(file.Name(), "cg-") {
 				cg := filepath.Join(cgRoot, file.Name())
 				fmt.Printf("Attempting to remove %s\n", cg)
+				kill := filepath.Join(cg, "cgroup.kill")
+				if err := os.WriteFile(kill, []byte(fmt.Sprintf("%d", 1)), os.ModeAppend); err != nil {
+					// Print an error if killing processes in the cgroup fails.
+					fmt.Printf("Could not kill processes in cgroup: %s\n", err.Error())
+					cgroupErrorCount += 1
+				}
 				if err := syscall.Rmdir(cg); err != nil {
+					// Print an error if removing a cgroup fails.
 					fmt.Printf("could not remove cgroup: %s\n", err.Error())
 					cgroupErrorCount += 1
 				}
 			}
 		}
+		_ = syscall.Rmdir(filepath.Join(cgRoot, "worker"))
+		_ = syscall.Rmdir(cgRoot)
 	}
 
 	sandboxErrorCount := 0

diff --git a/go/worker/sandbox/cgroups/cgroup.go b/go/worker/sandbox/cgroups/cgroup.go
@@ -78,9 +78,9 @@ func (cg *CgroupImpl) Destroy() {
 	}
 }
 
-// GroupPath returns the path to this cgroup directory.
+// GroupPath returns the path to the Cgroup pool for OpenLambda
 func (cg *CgroupImpl) GroupPath() string {
-	return fmt.Sprintf("%s/%s", cg.pool.poolPath, cg.name)
+	return fmt.Sprintf("%s/%s", cg.pool.GroupPath(), cg.name)
 }
 
 func (cg *CgroupImpl) MemoryEvents() map[string]int64 {
@@ -107,7 +107,7 @@ func (cg *CgroupImpl) MemoryEvents() map[string]int64 {
 
 // ResourcePath returns the path to a specific resource in this cgroup
 func (cg *CgroupImpl) ResourcePath(resource string) string {
-	return fmt.Sprintf("%s/%s/%s", cg.pool.poolPath, cg.name, resource)
+	return fmt.Sprintf("%s/%s/%s", cg.pool.GroupPath(), cg.name, resource)
 }
 
 func (cg *CgroupImpl) TryWriteInt(resource string, val int64) error {

diff --git a/go/worker/sandbox/cgroups/pool.go b/go/worker/sandbox/cgroups/pool.go
@@ -1,11 +1,13 @@
 package cgroups
 
 import (
+	"errors"
 	"fmt"
 	"log/slog"
 	"os"
+	"path"
 	"path/filepath"
-	"strconv"
+	"slices"
 	"strings"
 	"syscall"
 
@@ -25,49 +27,73 @@ type CgroupPool struct {
 	nextID   int
 }
 
-// InitPoolRoot creates the cgroup pool root directory and enables controllers.
-func InitPoolRoot(poolPath string) error {
+// NewCgroupPool creates a new CgroupPool with the specified name.
+func NewCgroupPool(name, poolPath string) (*CgroupPool, error) {
+	pool := &CgroupPool{
+		Name:     path.Base(path.Dir(common.Conf.Worker_dir)) + "-" + name,
+		poolPath: poolPath,
+		ready:    make(chan *CgroupImpl, CGROUP_RESERVE),
+		recycled: make(chan *CgroupImpl, CGROUP_RESERVE),
+		quit:     make(chan chan bool),
+		nextID:   0,
+	}
 
+	pool.printf("using parent cgroup %s", poolPath)
+
+	// create cgroup pool parent - no-op if already exists
 	if err := os.MkdirAll(poolPath, 0700); err != nil {
-		return fmt.Errorf("failed to create cgroup pool root %s: %w", poolPath, err)
+		return nil, fmt.Errorf("mkdir %s: %w", poolPath, err)
 	}
 
-	ctrlPath := filepath.Join(poolPath, "cgroup.subtree_control")
-	if err := os.WriteFile(ctrlPath, []byte("+pids +io +memory +cpu"), os.ModeAppend); err != nil {
-		return fmt.Errorf("failed to enable controllers at %s: %w", ctrlPath, err)
+	// enumerate controllers delegated to the parent
+	ctrlData, err := os.ReadFile(filepath.Join(poolPath, "cgroup.controllers"))
+	if err != nil {
+		return nil, fmt.Errorf("read %s/cgroup.controllers: %w", poolPath, err)
+	}
+	available := strings.Fields(string(ctrlData))
+
+	// require cpu, memory, pids
+	for _, req := range []string{"cpu", "memory", "pids"} {
+		if !slices.Contains(available, req) {
+			return nil, fmt.Errorf(
+				"cannot delegate required cgroup controller %q to %s (available: %v); ",
+				req, poolPath, available,
+			)
+		}
 	}
 
-	uidStr := os.Getenv("SUDO_UID")
-	if uidStr == "" {
-		return fmt.Errorf("SUDO_UID not set; worker must be run with sudo")
+	// move self pid to worker cgroup before enabling subtree_control
+	workerPath := filepath.Join(poolPath, "worker")
+	if err := os.Mkdir(workerPath, 0700); err != nil && !os.IsExist(err) {
+		return nil, fmt.Errorf("mkdir %s: %w", workerPath, err)
 	}
-	uid, err := strconv.Atoi(uidStr)
+	// migrate every pid in parent to worker/ - --detach leaves >1 in the scope
+	parentProcs, err := os.ReadFile(filepath.Join(poolPath, "cgroup.procs"))
 	if err != nil {
-		return fmt.Errorf("invalid SUDO_UID value %q: %w", uidStr, err)
+		return nil, fmt.Errorf("read %s/cgroup.procs: %w", poolPath, err)
 	}
-	if err := os.Chown(poolPath, uid, uid); err != nil {
-		return fmt.Errorf("failed to chown cgroup pool root: %w", err)
+	workerProcsPath := filepath.Join(workerPath, "cgroup.procs")
+	for _, pidStr := range strings.Fields(string(parentProcs)) {
+		// ESRCH: pid exited between read and write - ignore
+		if err := os.WriteFile(workerProcsPath, []byte(pidStr), 0); err != nil && !errors.Is(err, syscall.ESRCH) {
+			return nil, fmt.Errorf("move pid %s into %s: %w", pidStr, workerPath, err)
+		}
 	}
 
-	fmt.Printf("\tCreated cgroup pool root at %s\n", poolPath)
-	return nil
-}
-
-func NewCgroupPool(name string, poolPath string) (*CgroupPool, error) {
-	pool := &CgroupPool{
-		Name:     name,
-		poolPath: poolPath,
-		ready:    make(chan *CgroupImpl, CGROUP_RESERVE),
-		recycled: make(chan *CgroupImpl, CGROUP_RESERVE),
-		quit:     make(chan chan bool),
-		nextID:   0,
+	// enable every available controller on parent's subtree_control
+	var enable strings.Builder
+	for i, c := range available {
+		if i > 0 {
+			enable.WriteByte(' ')
+		}
+		enable.WriteByte('+')
+		enable.WriteString(c)
 	}
-
-	if st, err := os.Stat(poolPath); err != nil || !st.IsDir() {
-		return nil, fmt.Errorf("cgroup pool root %s does not exist.", poolPath)
+	subPath := filepath.Join(poolPath, "cgroup.subtree_control")
+	if err := os.WriteFile(subPath, []byte(enable.String()), 0); err != nil {
+		return nil, fmt.Errorf("enable controllers at %s: %w", subPath, err)
 	}
 
-	pool.printf("reusing pool root %s", poolPath)
 	go pool.cgTask()
 	return pool, nil
 }
@@ -156,14 +182,13 @@ Empty:
 	done <- true
 }
 
-// Destroy drains all child cgroups but preserves the pool root.
+// Destroy this entire cgroup pool
 func (pool *CgroupPool) Destroy() {
 	// signal cgTask, then wait for it to finish
 	ch := make(chan bool)
 	pool.quit <- ch
 	<-ch
-
-	pool.printf("destroyed all child cgroups, pool root preserved")
+	pool.printf("cgroup pool drained")
 }
 
 // GetCg retrieves a cgroup from the pool, setting its memory limit and CPU percentage.
@@ -184,3 +209,8 @@ func (pool *CgroupPool) GetCg(memLimitMB int, moveMemCharge bool, cpuPercent int
 
 	return cg
 }
+
+// GroupPath returns the path to the Cgroup pool for OpenLambda
+func (pool *CgroupPool) GroupPath() string {
+	return pool.poolPath
+}