only prioritize ourselves

sreya · sreya · commit 5020eb4d11e5 · 2023-09-14T20:30:32.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -79,6 +79,9 @@ type Options struct {
 	ModifiedProcesses chan []*agentproc.Process
 	// ProcessManagementTick is used for testing process priority management.
 	ProcessManagementTick <-chan time.Time
+	// PrioritizedPIDs are processes that should have preferred CPU allocation
+	// and be de-prioritized for OOM Killer selection.
+	PrioritizedPIDs []int
 }
 
 type Client interface {
@@ -161,6 +164,7 @@ func New(options Options) Agent {
 		syscaller:                    options.Syscaller,
 		modifiedProcs:                options.ModifiedProcesses,
 		processManagementTick:        options.ProcessManagementTick,
+		prioritizedPIDs:              options.PrioritizedPIDs,
 
 		prometheusRegistry: prometheusRegistry,
 		metrics:            newAgentMetrics(prometheusRegistry),
@@ -221,6 +225,9 @@ type agent struct {
 	modifiedProcs chan []*agentproc.Process
 	// processManagementTick is used for testing process priority management.
 	processManagementTick <-chan time.Time
+	// prioritizedPIDs are processes that should have preferred CPU allocation
+	// and be de-prioritized for OOM Killer selection.
+	prioritizedPIDs []int
 }
 
 func (a *agent) TailnetConn() *tailnet.Conn {
@@ -1278,8 +1285,6 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 	}
 }
 
-var prioritizedProcs = []string{"coder"}
-
 func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
 	if val := a.envVars[EnvProcPrioMgmt]; val == "" || runtime.GOOS != "linux" {
 		a.logger.Debug(ctx, "process priority not enabled, agent will not manage process niceness/oom_score_adj ",
@@ -1337,11 +1342,9 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
 			slog.F("pid", proc.PID),
 		)
 
-		// Trim off the path e.g. "./coder" -> "coder"
-		name := filepath.Base(proc.Name())
 		// If the process is prioritized we should adjust
 		// it's oom_score_adj and avoid lowering its niceness.
-		if slices.Contains(prioritizedProcs, name) {
+		if slices.Contains(a.prioritizedPIDs, int(proc.PID)) {
 			err = proc.SetOOMAdj(oomScoreAdj)
 			if err != nil {
 				logger.Warn(ctx, "unable to set proc oom_score_adj",
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2411,12 +2411,13 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		}
 
 		var (
-			expectedProcs = map[int32]agentproc.Process{}
-			fs            = afero.NewMemMapFs()
-			syscaller     = agentproctest.NewMockSyscaller(gomock.NewController(t))
-			ticker        = make(chan time.Time)
-			modProcs      = make(chan []*agentproc.Process)
-			logger        = slog.Make(sloghuman.Sink(io.Discard))
+			expectedProcs   = map[int32]agentproc.Process{}
+			fs              = afero.NewMemMapFs()
+			syscaller       = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			ticker          = make(chan time.Time)
+			modProcs        = make(chan []*agentproc.Process)
+			logger          = slog.Make(sloghuman.Sink(io.Discard))
+			prioritizedProc = 123
 		)
 
 		// Create some processes.
@@ -2429,12 +2430,15 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 				proc = agentproctest.GenerateProcess(t, fs,
 					func(p *agentproc.Process) {
 						p.CmdLine = "./coder\x00agent\x00--no-reap"
-						p.PID = 1
+						p.PID = int32(prioritizedProc)
 					},
 				)
 			} else {
 				// The rest are peasants.
 				proc = agentproctest.GenerateProcess(t, fs)
+				if proc.PID == int32(prioritizedProc) {
+					proc.PID = 1234
+				}
 				syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
 				syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
 			}
@@ -2452,6 +2456,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 			o.Filesystem = fs
 			o.Logger = logger
 			o.ProcessManagementTick = ticker
+			o.PrioritizedPIDs = []int{prioritizedProc}
 		})
 		actualProcs := <-modProcs
 		require.Len(t, actualProcs, 4)
@@ -2460,7 +2465,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 			expectedScore := "0"
 			expected, ok := expectedProcs[actual.PID]
 			require.True(t, ok)
-			if expected.PID == 1 {
+			if expected.PID == int32(prioritizedProc) {
 				expectedScore = "-500"
 			}
 
diff --git a/cli/agent.go b/cli/agent.go
@@ -299,6 +299,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				// Intentionally set this to nil. It's mainly used
 				// for testing.
 				ModifiedProcesses: nil,
+				PrioritizedPIDs:   []int{os.Getpid()},
 			})
 
 			prometheusSrvClose := ServeHandler(ctx, logger, prometheusMetricsHandler(prometheusRegistry, logger), prometheusAddress, "prometheus")
