fix: add agent exec abstraction (#15717)

agent/agentexec/exec.go

@@ -20,60 +20,101 @@ const (
 	EnvProcPrioMgmt  = "CODER_PROC_PRIO_MGMT"
 	EnvProcOOMScore  = "CODER_PROC_OOM_SCORE"
 	EnvProcNiceScore = "CODER_PROC_NICE_SCORE"
+
+	// unset is set to an invalid value for nice and oom scores.
+	unset = -2000
 )
 
-// CommandContext returns an exec.Cmd that calls "coder agent-exec" prior to exec'ing
-// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
-// is returned. All instances of exec.Cmd should flow through this function to ensure
-// proper resource constraints are applied to the child process.
-func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
-	cmd, args, err := agentExecCmd(cmd, args...)
-	if err != nil {
-		return nil, xerrors.Errorf("agent exec cmd: %w", err)
-	}
-	return exec.CommandContext(ctx, cmd, args...), nil
+var DefaultExecer Execer = execer{}
+
+// Execer defines an abstraction for creating exec.Cmd variants. It's unfortunately
+// necessary because we need to be able to wrap child processes with "coder agent-exec"
+// for templates that expect the agent to manage process priority.
+type Execer interface {
+	// CommandContext returns an exec.Cmd that calls "coder agent-exec" prior to exec'ing
+	// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
+	// is returned. All instances of exec.Cmd should flow through this function to ensure
+	// proper resource constraints are applied to the child process.
+	CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd
+	// PTYCommandContext returns an pty.Cmd that calls "coder agent-exec" prior to exec'ing
+	// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal pty.Cmd
+	// is returned. All instances of pty.Cmd should flow through this function to ensure
+	// proper resource constraints are applied to the child process.
+	PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd
 }
 
-// PTYCommandContext returns an pty.Cmd that calls "coder agent-exec" prior to exec'ing
-// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal pty.Cmd
-// is returned. All instances of pty.Cmd should flow through this function to ensure
-// proper resource constraints are applied to the child process.
-func PTYCommandContext(ctx context.Context, cmd string, args ...string) (*pty.Cmd, error) {
-	cmd, args, err := agentExecCmd(cmd, args...)
-	if err != nil {
-		return nil, xerrors.Errorf("agent exec cmd: %w", err)
-	}
-	return pty.CommandContext(ctx, cmd, args...), nil
-}
-
-func agentExecCmd(cmd string, args ...string) (string, []string, error) {
+func NewExecer() (Execer, error) {
 	_, enabled := os.LookupEnv(EnvProcPrioMgmt)
 	if runtime.GOOS != "linux" || !enabled {
-		return cmd, args, nil
+		return DefaultExecer, nil
 	}
 
 	executable, err := os.Executable()
 	if err != nil {
-		return "", nil, xerrors.Errorf("get executable: %w", err)
+		return nil, xerrors.Errorf("get executable: %w", err)
 	}
 
 	bin, err := filepath.EvalSymlinks(executable)
 	if err != nil {
-		return "", nil, xerrors.Errorf("eval symlinks: %w", err)
+		return nil, xerrors.Errorf("eval symlinks: %w", err)
 	}
 
+	oomScore, ok := envValInt(EnvProcOOMScore)
+	if !ok {
+		oomScore = unset
+	}
+
+	niceScore, ok := envValInt(EnvProcNiceScore)
+	if !ok {
+		niceScore = unset
+	}
+
+	return priorityExecer{
+		binPath:   bin,
+		oomScore:  oomScore,
+		niceScore: niceScore,
+	}, nil
+}
+
+type execer struct{}
+
+func (execer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
+	return exec.CommandContext(ctx, cmd, args...)
+}
+
+func (execer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
+	return pty.CommandContext(ctx, cmd, args...)
+}
+
+type priorityExecer struct {
+	binPath   string
+	oomScore  int
+	niceScore int
+}
+
+func (e priorityExecer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
+	cmd, args = e.agentExecCmd(cmd, args...)
+	return exec.CommandContext(ctx, cmd, args...)
+}
+
+func (e priorityExecer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
+	cmd, args = e.agentExecCmd(cmd, args...)
+	return pty.CommandContext(ctx, cmd, args...)
+}
+
+func (e priorityExecer) agentExecCmd(cmd string, args ...string) (string, []string) {
 	execArgs := []string{"agent-exec"}
-	if score, ok := envValInt(EnvProcOOMScore); ok {
-		execArgs = append(execArgs, oomScoreArg(score))
+	if e.oomScore != unset {
+		execArgs = append(execArgs, oomScoreArg(e.oomScore))
 	}
 
-	if score, ok := envValInt(EnvProcNiceScore); ok {
-		execArgs = append(execArgs, niceScoreArg(score))
+	if e.niceScore != unset {
+		execArgs = append(execArgs, niceScoreArg(e.niceScore))
 	}
 	execArgs = append(execArgs, "--", cmd)
 	execArgs = append(execArgs, args...)
 
-	return bin, execArgs, nil
+	return e.binPath, execArgs
 }
 
 // envValInt searches for a key in a list of environment variables and parses it to an int.