refactor into build files

coder · sreya · Sep 15, 2023 · Aug 30, 2023 · Aug 31, 2023 · Aug 31, 2023
commit 8230247e32cec554bbcee9de2c09e785ef8adf6d
diff --git a/agent/agent.go b/agent/agent.go
@@ -53,6 +53,8 @@ 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 {
@@ -73,7 +75,6 @@ type Options struct {
 	ReportMetadataInterval       time.Duration
 	ServiceBannerRefreshInterval time.Duration
 	Syscaller                    agentproc.Syscaller
-	ProcessManagementTick        <-chan time.Time
 	ModifiedProcesses            chan []*agentproc.Process
 }
 
@@ -128,7 +129,7 @@ func New(options Options) Agent {
 	}
 
 	if options.Syscaller == nil {
-		options.Syscaller = agentproc.UnixSyscaller{}
+		options.Syscaller = agentproc.NewSyscaller()
 	}
 
 	ctx, cancelFunc := context.WithCancel(context.Background())
@@ -155,7 +156,6 @@ func New(options Options) Agent {
 		subsystems:                   options.Subsystems,
 		addresses:                    options.Addresses,
 		syscaller:                    options.Syscaller,
-		processManagementTick:        options.ProcessManagementTick,
 		modifiedProcs:                options.ModifiedProcesses,
 
 		prometheusRegistry: prometheusRegistry,
@@ -209,11 +209,10 @@ type agent struct {
 
 	connCountReconnectingPTY atomic.Int64
 
-	prometheusRegistry    *prometheus.Registry
-	metrics               *agentMetrics
-	processManagementTick <-chan time.Time
-	modifiedProcs         chan []*agentproc.Process
-	syscaller             agentproc.Syscaller
+	prometheusRegistry *prometheus.Registry
+	metrics            *agentMetrics
+	modifiedProcs      chan []*agentproc.Process
+	syscaller          agentproc.Syscaller
 }
 
 func (a *agent) TailnetConn() *tailnet.Conn {
@@ -1299,9 +1298,12 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
 
 	manage()
 
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
 	for {
 		select {
-		case <-a.processManagementTick:
+		case <-ticker.C:
 			manage()
 		case <-ctx.Done():
 			return
@@ -1313,7 +1315,7 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
 	const (
 		procDir     = agentproc.DefaultProcDir
 		niceness    = 10
-		oomScoreAdj = -1000
+		oomScoreAdj = -500
 	)
 
 	procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
@@ -1357,6 +1359,11 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
 			)
 			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()),

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2409,7 +2409,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		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))
@@ -2439,7 +2438,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		}
 
 		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
-			o.ProcessManagementTick = ticker
 			o.Syscaller = syscaller
 			o.ModifiedProcesses = modProcs
 			o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
@@ -2454,7 +2452,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 			expected, ok := expectedProcs[actual.PID]
 			require.True(t, ok)
 			if expected.PID == 1 {
-				expectedScore = "-1000"
+				expectedScore = "-500"
 			}
 
 			score, err := afero.ReadFile(fs, filepath.Join(actual.Dir, "oom_score_adj"))
@@ -2469,7 +2467,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		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))
@@ -2494,7 +2491,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 		}
 
 		_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
-			o.ProcessManagementTick = ticker
 			o.Syscaller = syscaller
 			o.ModifiedProcesses = modProcs
 			o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}

diff --git a/agent/agentproc/syscaller.go b/agent/agentproc/syscaller.go
@@ -2,40 +2,10 @@ package agentproc
 
 import (
 	"syscall"
-
-	"golang.org/x/sys/unix"
-	"golang.org/x/xerrors"
 )
 
 type Syscaller interface {
 	SetPriority(pid int32, priority int) error
 	GetPriority(pid int32) (int, error)
 	Kill(pid int32, sig syscall.Signal) error
 }
-
-type UnixSyscaller struct{}
-
-func (UnixSyscaller) SetPriority(pid int32, nice int) error {
-	err := unix.Setpriority(unix.PRIO_PROCESS, int(pid), nice)
-	if err != nil {
-		return xerrors.Errorf("set priority: %w", err)
-	}
-	return nil
-}
-
-func (UnixSyscaller) GetPriority(pid int32) (int, error) {
-	nice, err := unix.Getpriority(0, int(pid))
-	if err != nil {
-		return 0, xerrors.Errorf("get priority: %w", err)
-	}
-	return nice, nil
-}
-
-func (UnixSyscaller) Kill(pid int32, sig syscall.Signal) error {
-	err := syscall.Kill(int(pid), sig)
-	if err != nil {
-		return xerrors.Errorf("kill: %w", err)
-	}
-
-	return nil
-}
diff --git a/agent/agentproc/syscaller_other.go b/agent/agentproc/syscaller_other.go
@@ -0,0 +1,24 @@
+//go:build !linux
+// +build !linux
+
+package agentproc
+
+import "syscall"
+
+func NewSyscaller() Syscaller {
+	return nopSyscaller{}
+}
+
+type nopSyscaller struct{}
+
+func (nopSyscaller) SetPriority(pid int32, priority int) error {
+	return nil
+}
+
+func (nopSyscaller) GetPriority(pid int32) (int, error) {
+	return 0, nil
+}
+
+func (nopSyscaller) Kill(pid int32, sig syscall.Signal) error {
+	return nil
+}
diff --git a/agent/agentproc/syscaller_unix.go b/agent/agentproc/syscaller_unix.go
@@ -0,0 +1,42 @@
+//go:build linux
+// +build linux
+
+package agentproc
+
+import (
+	"syscall"
+
+	"golang.org/x/sys/unix"
+	"golang.org/x/xerrors"
+)
+
+func NewSyscaller() Syscaller {
+	return UnixSyscaller{}
+}
+
+type UnixSyscaller struct{}
+
+func (UnixSyscaller) SetPriority(pid int32, nice int) error {
+	err := unix.Setpriority(unix.PRIO_PROCESS, int(pid), nice)
+	if err != nil {
+		return xerrors.Errorf("set priority: %w", err)
+	}
+	return nil
+}
+
+func (UnixSyscaller) GetPriority(pid int32) (int, error) {
+	nice, err := unix.Getpriority(0, int(pid))
+	if err != nil {
+		return 0, xerrors.Errorf("get priority: %w", err)
+	}
+	return nice, nil
+}
+
+func (UnixSyscaller) Kill(pid int32, sig syscall.Signal) error {
+	err := syscall.Kill(int(pid), sig)
+	if err != nil {
+		return xerrors.Errorf("kill: %w", err)
+	}
+
+	return nil
+}
diff --git a/cli/agent.go b/cli/agent.go
@@ -287,15 +287,15 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 					return resp.SessionToken, nil
 				},
 				EnvironmentVariables: map[string]string{
-					"GIT_ASKPASS": executablePath,
+					"GIT_ASKPASS":        executablePath,
+					agent.EnvProcMemNice: os.Getenv(agent.EnvProcMemNice),
 				},
 				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
 				Subsystems:    subsystems,
 
-				PrometheusRegistry:    prometheusRegistry,
-				ProcessManagementTick: procTicker.C,
-				Syscaller:             agentproc.UnixSyscaller{},
+				PrometheusRegistry: prometheusRegistry,
+				Syscaller:          agentproc.NewSyscaller(),
 				// Intentionally set this to nil. It's mainly used
 				// for testing.
 				ModifiedProcesses: nil,