feat: implement agent process management (#9461)

- An opt-in feature has been added to the agent to allow deprioritizing non coder-related processes for CPU by setting their niceness level to 10. - Opting in to the feature requires setting CODER_PROC_PRIO_MGMT to a non-empty value.
2026-09-24 15:04:27 +08:00 · 2023-09-14 19:45:05 -05:00
parent 79d4179123
commit 7311ffbd9d
13 changed files with 856 additions and 1 deletions
@@ -0,0 +1,5 @@
+// Package agentproctest contains utility functions
+// for testing process management in the agent.
+package agentproctest
+
+//go:generate mockgen -destination ./syscallermock.go -package agentproctest github.com/coder/coder/v2/agent/agentproc Syscaller
@@ -0,0 +1,49 @@
+package agentproctest
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/cryptorand"
+)
+
+func GenerateProcess(t *testing.T, fs afero.Fs, muts ...func(*agentproc.Process)) agentproc.Process {
+	t.Helper()
+
+	pid, err := cryptorand.Intn(1<<31 - 1)
+	require.NoError(t, err)
+
+	arg1, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	arg2, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	arg3, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	cmdline := fmt.Sprintf("%s\x00%s\x00%s", arg1, arg2, arg3)
+
+	process := agentproc.Process{
+		CmdLine: cmdline,
+		PID:     int32(pid),
+	}
+
+	for _, mut := range muts {
+		mut(&process)
+	}
+
+	process.Dir = fmt.Sprintf("%s/%d", "/proc", process.PID)
+
+	err = fs.MkdirAll(process.Dir, 0o555)
+	require.NoError(t, err)
+
+	err = afero.WriteFile(fs, fmt.Sprintf("%s/cmdline", process.Dir), []byte(process.CmdLine), 0o444)
+	require.NoError(t, err)
+
+	return process
+}
@@ -0,0 +1,78 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/coder/coder/v2/agent/agentproc (interfaces: Syscaller)
+
+// Package agentproctest is a generated GoMock package.
+package agentproctest
+
+import (
+	reflect "reflect"
+	syscall "syscall"
+
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockSyscaller is a mock of Syscaller interface.
+type MockSyscaller struct {
+	ctrl     *gomock.Controller
+	recorder *MockSyscallerMockRecorder
+}
+
+// MockSyscallerMockRecorder is the mock recorder for MockSyscaller.
+type MockSyscallerMockRecorder struct {
+	mock *MockSyscaller
+}
+
+// NewMockSyscaller creates a new mock instance.
+func NewMockSyscaller(ctrl *gomock.Controller) *MockSyscaller {
+	mock := &MockSyscaller{ctrl: ctrl}
+	mock.recorder = &MockSyscallerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockSyscaller) EXPECT() *MockSyscallerMockRecorder {
+	return m.recorder
+}
+
+// GetPriority mocks base method.
+func (m *MockSyscaller) GetPriority(arg0 int32) (int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPriority", arg0)
+	ret0, _ := ret[0].(int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetPriority indicates an expected call of GetPriority.
+func (mr *MockSyscallerMockRecorder) GetPriority(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPriority", reflect.TypeOf((*MockSyscaller)(nil).GetPriority), arg0)
+}
+
+// Kill mocks base method.
+func (m *MockSyscaller) Kill(arg0 int32, arg1 syscall.Signal) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Kill", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Kill indicates an expected call of Kill.
+func (mr *MockSyscallerMockRecorder) Kill(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Kill", reflect.TypeOf((*MockSyscaller)(nil).Kill), arg0, arg1)
+}
+
+// SetPriority mocks base method.
+func (m *MockSyscaller) SetPriority(arg0 int32, arg1 int) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SetPriority", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SetPriority indicates an expected call of SetPriority.
+func (mr *MockSyscallerMockRecorder) SetPriority(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetPriority", reflect.TypeOf((*MockSyscaller)(nil).SetPriority), arg0, arg1)
+}
@@ -0,0 +1,3 @@
+// Package agentproc contains logic for interfacing with local
+// processes running in the same context as the agent.
+package agentproc
@@ -0,0 +1,24 @@
+//go:build !linux
+// +build !linux
+
+package agentproc
+
+import (
+	"github.com/spf13/afero"
+)
+
+func (p *Process) Niceness(sc Syscaller) (int, error) {
+	return 0, errUnimplemented
+}
+
+func (p *Process) SetNiceness(sc Syscaller, score int) error {
+	return errUnimplemented
+}
+
+func (p *Process) Cmd() string {
+	return ""
+}
+
+func List(fs afero.Fs, syscaller Syscaller) ([]*Process, error) {
+	return nil, errUnimplemented
+}
@@ -0,0 +1,166 @@
+package agentproc_test
+
+import (
+	"runtime"
+	"syscall"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/agent/agentproc/agentproctest"
+)
+
+func TestList(t *testing.T) {
+	t.Parallel()
+
+	if runtime.GOOS != "linux" {
+		t.Skipf("skipping non-linux environment")
+	}
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+		)
+
+		for i := 0; i < 4; i++ {
+			proc := agentproctest.GenerateProcess(t, fs)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		actualProcs, err := agentproc.List(fs, sc)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, len(expectedProcs))
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
+
+	t.Run("FinishedProcess", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+		)
+
+		for i := 0; i < 3; i++ {
+			proc := agentproctest.GenerateProcess(t, fs)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		// Create a process that's already finished. We're not adding
+		// it to the map because it should be skipped over.
+		proc := agentproctest.GenerateProcess(t, fs)
+		sc.EXPECT().
+			Kill(proc.PID, syscall.Signal(0)).
+			Return(xerrors.New("os: process already finished"))
+
+		actualProcs, err := agentproc.List(fs, sc)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, len(expectedProcs))
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
+
+	t.Run("NoSuchProcess", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+		)
+
+		for i := 0; i < 3; i++ {
+			proc := agentproctest.GenerateProcess(t, fs)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		// Create a process that doesn't exist. We're not adding
+		// it to the map because it should be skipped over.
+		proc := agentproctest.GenerateProcess(t, fs)
+		sc.EXPECT().
+			Kill(proc.PID, syscall.Signal(0)).
+			Return(syscall.ESRCH)
+
+		actualProcs, err := agentproc.List(fs, sc)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, len(expectedProcs))
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
+}
+
+// These tests are not very interesting but they provide some modicum of
+// confidence.
+func TestProcess(t *testing.T) {
+	t.Parallel()
+
+	if runtime.GOOS != "linux" {
+		t.Skipf("skipping non-linux environment")
+	}
+
+	t.Run("SetNiceness", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			sc   = agentproctest.NewMockSyscaller(gomock.NewController(t))
+			proc = &agentproc.Process{
+				PID: 32,
+			}
+			score = 20
+		)
+
+		sc.EXPECT().SetPriority(proc.PID, score).Return(nil)
+		err := proc.SetNiceness(sc, score)
+		require.NoError(t, err)
+	})
+
+	t.Run("Cmd", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			proc = &agentproc.Process{
+				CmdLine: "helloworld\x00--arg1\x00--arg2",
+			}
+			expectedName = "helloworld --arg1 --arg2"
+		)
+
+		require.Equal(t, expectedName, proc.Cmd())
+	})
+}
@@ -0,0 +1,109 @@
+//go:build linux
+// +build linux
+
+package agentproc
+
+import (
+	"errors"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"syscall"
+
+	"github.com/spf13/afero"
+	"golang.org/x/xerrors"
+)
+
+func List(fs afero.Fs, syscaller Syscaller) ([]*Process, error) {
+	d, err := fs.Open(defaultProcDir)
+	if err != nil {
+		return nil, xerrors.Errorf("open dir %q: %w", defaultProcDir, err)
+	}
+	defer d.Close()
+
+	entries, err := d.Readdirnames(0)
+	if err != nil {
+		return nil, xerrors.Errorf("readdirnames: %w", err)
+	}
+
+	processes := make([]*Process, 0, len(entries))
+	for _, entry := range entries {
+		pid, err := strconv.ParseInt(entry, 10, 32)
+		if err != nil {
+			continue
+		}
+
+		// Check that the process still exists.
+		exists, err := isProcessExist(syscaller, int32(pid))
+		if err != nil {
+			return nil, xerrors.Errorf("check process exists: %w", err)
+		}
+		if !exists {
+			continue
+		}
+
+		cmdline, err := afero.ReadFile(fs, filepath.Join(defaultProcDir, entry, "cmdline"))
+		if err != nil {
+			var errNo syscall.Errno
+			if xerrors.As(err, &errNo) && errNo == syscall.EPERM {
+				continue
+			}
+			return nil, xerrors.Errorf("read cmdline: %w", err)
+		}
+		processes = append(processes, &Process{
+			PID:     int32(pid),
+			CmdLine: string(cmdline),
+			Dir:     filepath.Join(defaultProcDir, entry),
+		})
+	}
+
+	return processes, nil
+}
+
+func isProcessExist(syscaller Syscaller, pid int32) (bool, error) {
+	err := syscaller.Kill(pid, syscall.Signal(0))
+	if err == nil {
+		return true, nil
+	}
+	if err.Error() == "os: process already finished" {
+		return false, nil
+	}
+
+	var errno syscall.Errno
+	if !errors.As(err, &errno) {
+		return false, err
+	}
+
+	switch errno {
+	case syscall.ESRCH:
+		return false, nil
+	case syscall.EPERM:
+		return true, nil
+	}
+
+	return false, xerrors.Errorf("kill: %w", err)
+}
+
+func (p *Process) Niceness(sc Syscaller) (int, error) {
+	nice, err := sc.GetPriority(p.PID)
+	if err != nil {
+		return 0, xerrors.Errorf("get priority for %q: %w", p.CmdLine, err)
+	}
+	return nice, nil
+}
+
+func (p *Process) SetNiceness(sc Syscaller, score int) error {
+	err := sc.SetPriority(p.PID, score)
+	if err != nil {
+		return xerrors.Errorf("set priority for %q: %w", p.CmdLine, err)
+	}
+	return nil
+}
+
+func (p *Process) Cmd() string {
+	return strings.Join(p.cmdLine(), " ")
+}
+
+func (p *Process) cmdLine() []string {
+	return strings.Split(p.CmdLine, "\x00")
+}
@@ -0,0 +1,19 @@
+package agentproc
+
+import (
+	"syscall"
+)
+
+type Syscaller interface {
+	SetPriority(pid int32, priority int) error
+	GetPriority(pid int32) (int, error)
+	Kill(pid int32, sig syscall.Signal) error
+}
+
+const defaultProcDir = "/proc"
+
+type Process struct {
+	Dir     string
+	CmdLine string
+	PID     int32
+}
@@ -0,0 +1,30 @@
+//go:build !linux
+// +build !linux
+
+package agentproc
+
+import (
+	"syscall"
+
+	"golang.org/x/xerrors"
+)
+
+func NewSyscaller() Syscaller {
+	return nopSyscaller{}
+}
+
+var errUnimplemented = xerrors.New("unimplemented")
+
+type nopSyscaller struct{}
+
+func (nopSyscaller) SetPriority(pid int32, priority int) error {
+	return errUnimplemented
+}
+
+func (nopSyscaller) GetPriority(pid int32) (int, error) {
+	return 0, errUnimplemented
+}
+
+func (nopSyscaller) Kill(pid int32, sig syscall.Signal) error {
+	return errUnimplemented
+}
@@ -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
+}