coder · sreya · Apr 3, 2024 · Mar 19, 2024 · Mar 19, 2024 · Mar 20, 2024
diff --git a/agent/agent.go b/agent/agent.go
@@ -62,7 +62,10 @@ const (
 
 // EnvProcPrioMgmt determines whether we attempt to manage
 // process CPU and OOM Killer priority.
-const EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"
+const (
+	EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"
+	EnvProcOOMScore = "CODER_PROC_OOM_SCORE"
+)
 
 type Options struct {
 	Filesystem                   afero.Fs
@@ -1575,8 +1578,16 @@ func (a *agent) manageProcessPriorityUntilGracefulShutdown() {
 		a.processManagementTick = ticker.C
 	}
 
+	oomScore := unsetOOMScore
+	if scoreStr, ok := a.environmentVariables[EnvProcOOMScore]; ok {
+		score, err := strconv.Atoi(strings.TrimSpace(scoreStr))
+		if err == nil {
+			oomScore = score
+		}
+	}
+
 	for {
-		procs, err := a.manageProcessPriority(ctx)
+		procs, err := a.manageProcessPriority(ctx, oomScore)
 		if err != nil {
 			a.logger.Error(ctx, "manage process priority",
 				slog.Error(err),
@@ -1594,11 +1605,23 @@ func (a *agent) manageProcessPriorityUntilGracefulShutdown() {
 	}
 }
 
-func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process, error) {
+const unsetOOMScore = 1001
+
+func (a *agent) manageProcessPriority(ctx context.Context, oomScore int) ([]*agentproc.Process, error) {
 	const (
 		niceness = 10
 	)
 
+	agentScore, err := a.getAgentOOMScore()
+	if err != nil {
+		agentScore = unsetOOMScore
+	}
+	if oomScore == unsetOOMScore && agentScore != unsetOOMScore {
+		// If the child score has not been explicitly specified we should
+		// set it to a score relative to the agent score.
+		oomScore = childOOMScore(agentScore)
+	}
+
 	procs, err := agentproc.List(a.filesystem, a.syscaller)
 	if err != nil {
 		return nil, xerrors.Errorf("list: %w", err)
@@ -1622,14 +1645,14 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
 
 		// If the process is prioritized we should adjust
 		// it's oom_score_adj and avoid lowering its niceness.
-		if slices.ContainsFunc[[]string, string](prioritizedProcs, containsFn) {
+		if slices.ContainsFunc(prioritizedProcs, containsFn) {
 			continue
 		}
 
-		score, err := proc.Niceness(a.syscaller)
-		if err != nil {
+		score, niceErr := proc.Niceness(a.syscaller)
+		if niceErr != nil && !xerrors.Is(niceErr, os.ErrPermission) {
 			logger.Warn(ctx, "unable to get proc niceness",
-				slog.Error(err),
+				slog.Error(niceErr),
 			)
 			continue
 		}
@@ -1639,19 +1662,31 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
 		// Getpriority actually returns priority for the nice value
 		// which is niceness + 20, so here 20 = a niceness of 0 (aka unset).
 		if score != 20 {
-			// We don't log here since it can get spammy
 			continue
 		}
 
-		err = proc.SetNiceness(a.syscaller, niceness)
-		if err != nil {
-			logger.Warn(ctx, "unable to set proc niceness",
-				slog.F("niceness", niceness),
-				slog.Error(err),
-			)
-			continue
+		if niceErr == nil {
+			err := proc.SetNiceness(a.syscaller, niceness)
+			if err != nil && !xerrors.Is(err, os.ErrPermission) {
+				logger.Warn(ctx, "unable to set proc niceness",
+					slog.F("niceness", niceness),
+					slog.Error(err),
+				)
+			}
 		}
 
+		// If the oom score is valid and it's not already set and isn't a custom value set by another process
+		// then it's ok to update it.
+		if oomScore != unsetOOMScore && oomScore != proc.OOMScoreAdj && !isCustomOOMScore(agentScore, proc) {
+			oomScoreStr := strconv.Itoa(oomScore)
+			err := afero.WriteFile(a.filesystem, fmt.Sprintf("/proc/%d/oom_score_adj", proc.PID), []byte(oomScoreStr), 0o644)
+			if err != nil && !xerrors.Is(err, os.ErrPermission) {
+				logger.Warn(ctx, "unable to set oom_score_adj",
+					slog.F("score", "0"),
+					slog.Error(err),
+				)
+			}
+		}
 		modProcs = append(modProcs, proc)
 	}
 	return modProcs, nil
@@ -2005,3 +2040,44 @@ func PrometheusMetricsHandler(prometheusRegistry *prometheus.Registry, logger sl
 		}
 	})
 }
+
+// childOOMScore returns the oom_score_adj for a child process. It is based
+// on the oom_score_adj of the agent process.
+func childOOMScore(agentScore int) int {
+	// If the agent has a negative oom_score_adj, we set the child to 0
+	// so it's treated like every other process.
+	if agentScore < 0 {
+		return 0
+	}
+
+	// If the agent is already almost at the maximum then set it to the max.
+	if agentScore >= 998 {
+		return 1000
+	}
+
+	// If the agent oom_score_adj is >=0, we set the child to slightly
+	// less than the maximum. If users want a different score they set it
+	// directly.
+	return 998
+}
+
+func (a *agent) getAgentOOMScore() (int, error) {
+	scoreStr, err := afero.ReadFile(a.filesystem, "/proc/self/oom_score_adj")
+	if err != nil {
+		return 0, xerrors.Errorf("read file: %w", err)
+	}
+
+	score, err := strconv.Atoi(strings.TrimSpace(string(scoreStr)))
+	if err != nil {
+		return 0, xerrors.Errorf("parse int: %w", err)
+	}
+
+	return score, nil
+}
+
+// isCustomOOMScore checks to see if the oom_score_adj is not a value that would
+// originate from an agent-spawned process.
+func isCustomOOMScore(agentScore int, process *agentproc.Process) bool {
+	score := process.OOMScoreAdj
+	return agentScore != score && score != 1000 && score != 0 && score != 998
+}
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2529,11 +2529,11 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 			logger        = slog.Make(sloghuman.Sink(io.Discard))
 		)
 
+		requireFileWrite(t, fs, "/proc/self/oom_score_adj", "-500")
+
 		// Create some processes.
 		for i := 0; i < 4; i++ {
-			// Create a prioritized process. This process should
-			// have it's oom_score_adj set to -500 and its nice
-			// score should be untouched.
+			// Create a prioritized process.
 			var proc agentproc.Process
 			if i == 0 {
 				proc = agentproctest.GenerateProcess(t, fs,
@@ -2551,8 +2551,8 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 					},
 				)
 
-				syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
 				syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
+				syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
 			}
 			syscaller.EXPECT().
 				Kill(proc.PID, syscall.Signal(0)).
@@ -2571,6 +2571,9 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		})
 		actualProcs := <-modProcs
 		require.Len(t, actualProcs, len(expectedProcs)-1)
+		for _, proc := range actualProcs {
+			requireFileEquals(t, fs, fmt.Sprintf("/proc/%d/oom_score_adj", proc.PID), "0")
+		}
 	})
 
 	t.Run("IgnoreCustomNice", func(t *testing.T) {
@@ -2589,8 +2592,11 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 			logger        = slog.Make(sloghuman.Sink(io.Discard))
 		)
 
+		err := afero.WriteFile(fs, "/proc/self/oom_score_adj", []byte("0"), 0o644)
+		require.NoError(t, err)
+
 		// Create some processes.
-		for i := 0; i < 2; i++ {
+		for i := 0; i < 3; i++ {
 			proc := agentproctest.GenerateProcess(t, fs)
 			syscaller.EXPECT().
 				Kill(proc.PID, syscall.Signal(0)).
@@ -2618,7 +2624,59 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		})
 		actualProcs := <-modProcs
 		// We should ignore the process with a custom nice score.
-		require.Len(t, actualProcs, 1)
+		require.Len(t, actualProcs, 2)
+		for _, proc := range actualProcs {
+			_, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			requireFileEquals(t, fs, fmt.Sprintf("/proc/%d/oom_score_adj", proc.PID), "998")
+		}
+	})
+
+	t.Run("CustomOOMScore", func(t *testing.T) {
+		t.Parallel()
+
+		if runtime.GOOS != "linux" {
+			t.Skip("Skipping non-linux environment")
+		}
+
+		var (
+			fs        = afero.NewMemMapFs()
+			ticker    = make(chan time.Time)
+			syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			modProcs  = make(chan []*agentproc.Process)
+			logger    = slog.Make(sloghuman.Sink(io.Discard))
+		)
+
+		err := afero.WriteFile(fs, "/proc/self/oom_score_adj", []byte("0"), 0o644)
+		require.NoError(t, err)
+
+		// Create some processes.
+		for i := 0; i < 3; i++ {
+			proc := agentproctest.GenerateProcess(t, fs)
+			syscaller.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+			syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
+			syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
+		}
+
+		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
+			o.Syscaller = syscaller
+			o.ModifiedProcesses = modProcs
+			o.EnvironmentVariables = map[string]string{
+				agent.EnvProcPrioMgmt: "1",
+				agent.EnvProcOOMScore: "-567",
+			}
+			o.Filesystem = fs
+			o.Logger = logger
+			o.ProcessManagementTick = ticker
+		})
+		actualProcs := <-modProcs
+		// We should ignore the process with a custom nice score.
+		require.Len(t, actualProcs, 3)
+		for _, proc := range actualProcs {
+			requireFileEquals(t, fs, fmt.Sprintf("/proc/%d/oom_score_adj", proc.PID), "-567")
+		}
 	})
 
 	t.Run("DisabledByDefault", func(t *testing.T) {
@@ -2739,3 +2797,17 @@ func requireEcho(t *testing.T, conn net.Conn) {
 	require.NoError(t, err)
 	require.Equal(t, "test", string(b))
 }
+
+func requireFileWrite(t testing.TB, fs afero.Fs, fp, data string) {
+	t.Helper()
+	err := afero.WriteFile(fs, fp, []byte(data), 0o600)
+	require.NoError(t, err)
+}
+
+func requireFileEquals(t testing.TB, fs afero.Fs, fp, expect string) {
+	t.Helper()
+	actual, err := afero.ReadFile(fs, fp)
+	require.NoError(t, err)
+
+	require.Equal(t, expect, string(actual))
+}
diff --git a/agent/agentproc/agentproctest/proc.go b/agent/agentproc/agentproctest/proc.go
@@ -2,6 +2,7 @@ package agentproctest
 
 import (
 	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/spf13/afero"
@@ -29,8 +30,9 @@ func GenerateProcess(t *testing.T, fs afero.Fs, muts ...func(*agentproc.Process)
 	cmdline := fmt.Sprintf("%s\x00%s\x00%s", arg1, arg2, arg3)
 
 	process := agentproc.Process{
-		CmdLine: cmdline,
-		PID:     int32(pid),
+		CmdLine:     cmdline,
+		PID:         int32(pid),
+		OOMScoreAdj: 0,
 	}
 
 	for _, mut := range muts {
@@ -45,5 +47,9 @@ func GenerateProcess(t *testing.T, fs afero.Fs, muts ...func(*agentproc.Process)
 	err = afero.WriteFile(fs, fmt.Sprintf("%s/cmdline", process.Dir), []byte(process.CmdLine), 0o444)
 	require.NoError(t, err)
 
+	score := strconv.Itoa(process.OOMScoreAdj)
+	err = afero.WriteFile(fs, fmt.Sprintf("%s/oom_score_adj", process.Dir), []byte(score), 0o444)
+	require.NoError(t, err)
+
 	return process
 }
diff --git a/agent/agentproc/proc_unix.go b/agent/agentproc/proc_unix.go
@@ -5,6 +5,7 @@ package agentproc
 
 import (
 	"errors"
+	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -50,10 +51,26 @@ func List(fs afero.Fs, syscaller Syscaller) ([]*Process, error) {
 			}
 			return nil, xerrors.Errorf("read cmdline: %w", err)
 		}
+
+		oomScore, err := afero.ReadFile(fs, filepath.Join(defaultProcDir, entry, "oom_score_adj"))
+		if err != nil {
+			if xerrors.Is(err, os.ErrPermission) {
+				continue
+			}
+
+			return nil, xerrors.Errorf("read oom_score_adj: %w", err)
+		}
+
+		oom, err := strconv.Atoi(strings.TrimSpace(string(oomScore)))
+		if err != nil {
+			return nil, xerrors.Errorf("convert oom score: %w", err)
+		}
+
 		processes = append(processes, &Process{
-			PID:     int32(pid),
-			CmdLine: string(cmdline),
-			Dir:     filepath.Join(defaultProcDir, entry),
+			PID:         int32(pid),
+			CmdLine:     string(cmdline),
+			Dir:         filepath.Join(defaultProcDir, entry),
+			OOMScoreAdj: oom,
 		})
 	}
 

diff --git a/agent/agentproc/syscaller.go b/agent/agentproc/syscaller.go
@@ -14,7 +14,8 @@ type Syscaller interface {
 const defaultProcDir = "/proc"
 
 type Process struct {
-	Dir     string
-	CmdLine string
-	PID     int32
+	Dir         string
+	CmdLine     string
+	PID         int32
+	OOMScoreAdj int
 }