coder · sreya · Sep 15, 2023 · Aug 30, 2023 · Aug 31, 2023 · Aug 31, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -15,6 +15,7 @@ import (
 	"os/exec"
 	"os/user"
 	"path/filepath"
+	"runtime"
 	"sort"
 	"strconv"
 	"strings"
@@ -34,6 +35,7 @@ import (
 	"tailscale.com/types/netlogtype"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/agent/agentproc"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/agent/reconnectingpty"
 	"github.com/coder/coder/v2/buildinfo"
@@ -51,6 +53,10 @@ const (
 	ProtocolDial            = "dial"
 )
 
+// EnvProcMemNice determines whether we attempt to manage
+// process CPU and OOM Killer priority.
+const EnvProcMemNice = "CODER_PROC_MEMNICE_ENABLE"
+
 type Options struct {
 	Filesystem                   afero.Fs
 	LogDir                       string
@@ -68,6 +74,11 @@ type Options struct {
 	PrometheusRegistry           *prometheus.Registry
 	ReportMetadataInterval       time.Duration
 	ServiceBannerRefreshInterval time.Duration
+	Syscaller                    agentproc.Syscaller
+	// ModifiedProcesses is used for testing process priority management.
+	ModifiedProcesses chan []*agentproc.Process
+	// ProcessManagementTick is used for testing process priority management.
+	ProcessManagementTick <-chan time.Time
 }
 
 type Client interface {
@@ -120,6 +131,10 @@ func New(options Options) Agent {
 		prometheusRegistry = prometheus.NewRegistry()
 	}
 
+	if options.Syscaller == nil {
+		options.Syscaller = agentproc.NewSyscaller()
+	}
+
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	a := &agent{
 		tailnetListenPort:            options.TailnetListenPort,
@@ -143,6 +158,9 @@ func New(options Options) Agent {
 		sshMaxTimeout:                options.SSHMaxTimeout,
 		subsystems:                   options.Subsystems,
 		addresses:                    options.Addresses,
+		syscaller:                    options.Syscaller,
+		modifiedProcs:                options.ModifiedProcesses,
+		processManagementTick:        options.ProcessManagementTick,
 
 		prometheusRegistry: prometheusRegistry,
 		metrics:            newAgentMetrics(prometheusRegistry),
@@ -197,6 +215,12 @@ type agent struct {
 
 	prometheusRegistry *prometheus.Registry
 	metrics            *agentMetrics
+	syscaller          agentproc.Syscaller
+
+	// podifiedProcs is used for testing process priority management.
+	modifiedProcs chan []*agentproc.Process
+	// processManagementTick is used for testing process priority management.
+	processManagementTick <-chan time.Time
 }
 
 func (a *agent) TailnetConn() *tailnet.Conn {
@@ -225,6 +249,7 @@ func (a *agent) runLoop(ctx context.Context) {
 	go a.reportLifecycleLoop(ctx)
 	go a.reportMetadataLoop(ctx)
 	go a.fetchServiceBannerLoop(ctx)
+	go a.manageProcessPriorityLoop(ctx)
 
 	for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
 		a.logger.Info(ctx, "connecting to coderd")
@@ -1253,6 +1278,133 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 	}
 }
 
+var prioritizedProcs = []string{"coder"}
+
+func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
+	if val := a.envVars[EnvProcMemNice]; val == "" || runtime.GOOS != "linux" {
+		a.logger.Info(ctx, "process priority not enabled, agent will not manage process niceness/oom_score_adj ",
+			slog.F("env_var", EnvProcMemNice),
+			slog.F("value", val),
+			slog.F("goos", runtime.GOOS),
+		)
+		return
+	}
+
+	manage := func() {
+		procs, err := a.manageProcessPriority(ctx)
+		if err != nil {
+			a.logger.Error(ctx, "manage process priority",
+				slog.F("dir", agentproc.DefaultProcDir),
+				slog.Error(err),
+			)
+		}
+		if a.modifiedProcs != nil {
+			a.modifiedProcs <- procs
+		}
+	}
+
+	// Do once before falling into loop.
+	manage()
+
+	if a.processManagementTick == nil {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+		a.processManagementTick = ticker.C
+	}
+
+	for {
+		select {
+		case <-a.processManagementTick:
+			manage()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process, error) {
+	const (
+		procDir     = agentproc.DefaultProcDir
+		niceness    = 10
+		oomScoreAdj = -500
+	)
+
+	procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
+	if err != nil {
+		return nil, xerrors.Errorf("list: %w", err)
+	}
+
+	modProcs := []*agentproc.Process{}
+	for _, proc := range procs {
+		// 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) {
+			err = proc.SetOOMAdj(oomScoreAdj)
+			if err != nil {
+				a.logger.Error(ctx, "unable to set proc oom_score_adj",
+					slog.F("name", proc.Name()),
+					slog.F("pid", proc.PID),
+					slog.F("oom_score_adj", oomScoreAdj),
+					slog.Error(err),
+				)
+				continue
+			}
+			modProcs = append(modProcs, proc)
+
+			a.logger.Debug(ctx, "decreased process oom_score",
+				slog.F("name", proc.Name()),
+				slog.F("pid", proc.PID),
+				slog.F("oom_score_adj", oomScoreAdj),
+			)
+			continue
+		}
+
+		score, err := proc.Niceness(a.syscaller)
+		if err != nil {
+			a.logger.Error(ctx, "unable to get proc niceness",
+				slog.F("name", proc.Name()),
+				slog.F("pid", proc.PID),
+				slog.Error(err),
+			)
+			continue
+		}
+
+		// We only want processes that don't have a nice value set
+		// so we don't override user nice values.
+		// Getpriority actually returns priority for the nice value
+		// which is niceness + 20, so here 20 = a niceness of 0 (aka unset).
+		if score != 20 {
+			a.logger.Error(ctx, "skipping process due to custom niceness",
+				slog.F("name", proc.Name()),
+				slog.F("pid", proc.PID),
+				slog.F("niceness", score),
+			)
+			continue
+		}
+
+		err = proc.SetNiceness(a.syscaller, niceness)
+		if err != nil {
+			a.logger.Error(ctx, "unable to set proc niceness",
+				slog.F("name", proc.Name()),
+				slog.F("pid", proc.PID),
+				slog.F("niceness", niceness),
+				slog.Error(err),
+			)
+			continue
+		}
+
+		a.logger.Debug(ctx, "deprioritized process",
+			slog.F("name", proc.Name()),
+			slog.F("pid", proc.PID),
+			slog.F("niceness", niceness),
+		)
+		modProcs = append(modProcs, proc)
+	}
+	return modProcs, nil
+}
+
 // isClosed returns whether the API is closed or not.
 func (a *agent) isClosed() bool {
 	select {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -21,10 +21,12 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"testing"
 	"time"
 
 	scp "github.com/bramvdbogaerde/go-scp"
+	"github.com/golang/mock/gomock"
 	"github.com/google/uuid"
 	"github.com/pion/udp"
 	"github.com/pkg/sftp"
@@ -41,8 +43,11 @@ import (
 	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/agent"
+	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/agent/agentproc/agentproctest"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/coderd/httpapi"
@@ -2395,6 +2400,157 @@ func TestAgent_Metrics_SSH(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestAgent_ManageProcessPriority(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		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))
+		)
+
+		// 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.
+			var proc agentproc.Process
+			if i == 0 {
+				proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir,
+					func(p *agentproc.Process) {
+						p.CmdLine = "./coder\x00agent\x00--no-reap"
+						p.PID = 1
+					},
+				)
+			} else {
+				// The rest are peasants.
+				proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
+				syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
+				syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
+			}
+			syscaller.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+
+			expectedProcs[proc.PID] = proc
+		}
+
+		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
+			o.Syscaller = syscaller
+			o.ModifiedProcesses = modProcs
+			o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
+			o.Filesystem = fs
+			o.Logger = logger
+			o.ProcessManagementTick = ticker
+		})
+		actualProcs := <-modProcs
+		require.Len(t, actualProcs, 4)
+
+		for _, actual := range actualProcs {
+			expectedScore := "0"
+			expected, ok := expectedProcs[actual.PID]
+			require.True(t, ok)
+			if expected.PID == 1 {
+				expectedScore = "-500"
+			}
+
+			score, err := afero.ReadFile(fs, filepath.Join(actual.Dir, "oom_score_adj"))
+			require.NoError(t, err)
+			require.Equal(t, expectedScore, strings.TrimSpace(string(score)))
+		}
+	})
+
+	t.Run("IgnoreCustomNice", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			expectedProcs = map[int32]agentproc.Process{}
+			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))
+		)
+
+		// Create some processes.
+		for i := 0; i < 2; i++ {
+			proc := agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
+			syscaller.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+
+			if i == 0 {
+				// Set a random nice score. This one should not be adjusted by
+				// our management loop.
+				syscaller.EXPECT().GetPriority(proc.PID).Return(25, nil)
+			} else {
+				syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
+				syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
+			}
+
+			expectedProcs[proc.PID] = proc
+		}
+
+		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
+			o.Syscaller = syscaller
+			o.ModifiedProcesses = modProcs
+			o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
+			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, 1)
+	})
+
+	t.Run("DisabledByDefault", func(t *testing.T) {
+		t.Parallel()
+
+		if runtime.GOOS != "linux" {
+			t.Skip("Skipping non-linux environment")
+		}
+
+		var buf bytes.Buffer
+		log := slog.Make(sloghuman.Sink(&buf))
+
+		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
+			o.Logger = log
+		})
+
+		require.Eventually(t, func() bool {
+			return strings.Contains(buf.String(), "process priority not enabled")
+		}, testutil.WaitLong, testutil.IntervalFast)
+	})
+
+	t.Run("DisabledForNonLinux", func(t *testing.T) {
+		t.Parallel()
+
+		if runtime.GOOS == "linux" {
+			t.Skip("Skipping linux environment")
+		}
+
+		var buf bytes.Buffer
+		log := slog.Make(sloghuman.Sink(&buf))
+
+		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
+			o.Logger = log
+			// Try to enable it so that we can assert that non-linux
+			// environments are truly disabled.
+			o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
+		})
+		require.Eventually(t, func() bool {
+			return strings.Contains(buf.String(), "process priority not enabled")
+		}, testutil.WaitLong, testutil.IntervalFast)
+	})
+}
+
 func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actual []*promgo.MetricFamily) bool {
 	t.Helper()