Skip to content

feat: implement agent process management #9461

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 21 commits into from
Sep 15, 2023
Merged

feat: implement agent process management #9461

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
4ed4069
feat: implement agent process management
sreya Aug 30, 2023
7e59db6
improve process detection
sreya Aug 31, 2023
8c65216
add agentproc tests
sreya Aug 31, 2023
760cbb3
some minor agent tests
sreya Sep 1, 2023
f4b864e
add a proper test for proc management
sreya Sep 1, 2023
cbcb854
custom nice
sreya Sep 1, 2023
8230247
refactor into build files
sreya Sep 8, 2023
3e1defd
tick tock
sreya Sep 8, 2023
2fe9c70
make fmt
sreya Sep 8, 2023
ef41e9a
pr comments
sreya Sep 12, 2023
05baba0
lint
sreya Sep 12, 2023
478d57c
whoops
sreya Sep 12, 2023
0ced5ce
skip non-linux
sreya Sep 13, 2023
8aaa6d5
skip non-linux
sreya Sep 13, 2023
cea4851
prevent race
sreya Sep 13, 2023
5020eb4
only prioritize ourselves
sreya Sep 14, 2023
11ab047
Revert "only prioritize ourselves"
sreya Sep 14, 2023
46ef05a
only prioritize coder agent
sreya Sep 14, 2023
d132480
remove oom_score_adj
sreya Sep 14, 2023
04ee5cb
defer first
sreya Sep 15, 2023
ffbeab9
avoid resetting niceness for already niced procs
sreya Sep 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
pr comments
  • Loading branch information
@sreya
sreya committed Sep 12, 2023
commit ef41e9a4bcec19817a760bf6cdf48189e4fe73dc
74 changes: 28 additions & 46 deletions agent/agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@ const (
ProtocolDial = "dial"
)

// EnvProcMemNice determines whether we attempt to manage
// EnvProcPrioMgmt determines whether we attempt to manage
// process CPU and OOM Killer priority.
const EnvProcMemNice = "CODER_PROC_MEMNICE_ENABLE"
const EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"

type Options struct {
Filesystem afero.Fs
Expand Down Expand Up @@ -217,7 +217,7 @@ type agent struct {
metrics *agentMetrics
syscaller agentproc.Syscaller

// podifiedProcs is used for testing process priority management.
// modifiedProcs is used for testing process priority management.
modifiedProcs chan []*agentproc.Process
// processManagementTick is used for testing process priority management.
processManagementTick <-chan time.Time
Expand Down Expand Up @@ -1281,41 +1281,34 @@ 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),
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 ",
slog.F("env_var", EnvProcPrioMgmt),
slog.F("value", val),
slog.F("goos", runtime.GOOS),
)
return
}

manage := func() {
if a.processManagementTick == nil {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
a.processManagementTick = ticker.C
}

for {
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
}
Expand All @@ -1324,48 +1317,46 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {

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)
procs, err := agentproc.List(a.filesystem, a.syscaller)
if err != nil {
return nil, xerrors.Errorf("list: %w", err)
}

modProcs := []*agentproc.Process{}
var (
modProcs = []*agentproc.Process{}
logger slog.Logger
)

for _, proc := range procs {
logger = a.logger.With(
slog.F("name", proc.Name()),
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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We want to specifically prioritize the agent and not other coder processes right? If I'm reading this code correctly it would treat coder server and coder stat the same as the agent.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good catch, I don't see that as being a big deal but we can be more discriminate about which processes we want to prioritize by also parsing command arguments. WDYT?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not just check if its the current process?

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),
logger.Warn(ctx, "unable to set proc oom_score_adj",
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),
logger.Warn(ctx, "unable to get proc niceness",
slog.Error(err),
)
continue
Expand All @@ -1376,30 +1367,21 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
// 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),
logger.Debug(ctx, "skipping process due to custom niceness",
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),
logger.Warn(ctx, "unable to set proc niceness",
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
Expand Down
16 changes: 8 additions & 8 deletions agent/agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2422,15 +2422,15 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
// score should be untouched.
var proc agentproc.Process
if i == 0 {
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir,
proc = agentproctest.GenerateProcess(t, fs,
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)
proc = agentproctest.GenerateProcess(t, fs)
syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
}
Expand All @@ -2444,7 +2444,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
_, _, _, _, _ = 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.EnvironmentVariables = map[string]string{agent.EnvProcPrioMgmt: "1"}
o.Filesystem = fs
o.Logger = logger
o.ProcessManagementTick = ticker
Expand Down Expand Up @@ -2480,7 +2480,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {

// Create some processes.
for i := 0; i < 2; i++ {
proc := agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
proc := agentproctest.GenerateProcess(t, fs)
syscaller.EXPECT().
Kill(proc.PID, syscall.Signal(0)).
Return(nil)
Expand All @@ -2500,7 +2500,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
_, _, _, _, _ = 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.EnvironmentVariables = map[string]string{agent.EnvProcPrioMgmt: "1"}
o.Filesystem = fs
o.Logger = logger
o.ProcessManagementTick = ticker
Expand All @@ -2518,7 +2518,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
}

var buf bytes.Buffer
log := slog.Make(sloghuman.Sink(&buf))
log := slog.Make(sloghuman.Sink(&buf)).Leveled(slog.LevelDebug)

_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
o.Logger = log
Expand All @@ -2537,13 +2537,13 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
}

var buf bytes.Buffer
log := slog.Make(sloghuman.Sink(&buf))
log := slog.Make(sloghuman.Sink(&buf)).Leveled(slog.LevelDebug)

_, _, _, _, _ = 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"}
o.EnvironmentVariables = map[string]string{agent.EnvProcPrioMgmt: "1"}
})
require.Eventually(t, func() bool {
return strings.Contains(buf.String(), "process priority not enabled")
Expand Down
4 changes: 2 additions & 2 deletions agent/agentproc/agentproctest/proc.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (
"github.com/coder/coder/v2/cryptorand"
)

func GenerateProcess(t *testing.T, fs afero.Fs, dir string, muts ...func(*agentproc.Process)) agentproc.Process {
func GenerateProcess(t *testing.T, fs afero.Fs, muts ...func(*agentproc.Process)) agentproc.Process {
t.Helper()

pid, err := cryptorand.Intn(1<<31 - 1)
Expand All @@ -38,7 +38,7 @@ func GenerateProcess(t *testing.T, fs afero.Fs, dir string, muts ...func(*agentp
mut(&process)
}

process.Dir = fmt.Sprintf("%s/%d", dir, process.PID)
process.Dir = fmt.Sprintf("%s/%d", "/proc", process.PID)

err = fs.MkdirAll(process.Dir, 0o555)
require.NoError(t, err)
Expand Down
28 changes: 28 additions & 0 deletions agent/agentproc/proc_other.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
//go:build !linux
// +build !linux

package agentproc

import (
"github.com/spf13/afero"
)

func (p *Process) SetOOMAdj(score int) error {
return errUnimplimented
}

func (p *Process) Niceness(sc Syscaller) (int, error) {
return 0, errUnimplimented
}

func (p *Process) SetNiceness(sc Syscaller, score int) error {
return errUnimplimented
}

func (p *Process) Name() string {
return ""
}

func List(fs afero.Fs, syscaller Syscaller) ([]*Process, error) {
return nil, errUnimplimented
}
30 changes: 13 additions & 17 deletions agent/agentproc/proc_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,21 +25,20 @@ func TestList(t *testing.T) {
fs = afero.NewMemMapFs()
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
expectedProcs = make(map[int32]agentproc.Process)
rootDir = agentproc.DefaultProcDir
)

for i := 0; i < 4; i++ {
proc := agentproctest.GenerateProcess(t, fs, rootDir)
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, rootDir)
actualProcs, err := agentproc.List(fs, sc)
require.NoError(t, err)
require.Len(t, actualProcs, 4)
require.Len(t, actualProcs, len(expectedProcs))
for _, proc := range actualProcs {
expected, ok := expectedProcs[proc.PID]
require.True(t, ok)
Expand All @@ -56,11 +55,10 @@ func TestList(t *testing.T) {
fs = afero.NewMemMapFs()
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
expectedProcs = make(map[int32]agentproc.Process)
rootDir = agentproc.DefaultProcDir
)

for i := 0; i < 3; i++ {
proc := agentproctest.GenerateProcess(t, fs, rootDir)
proc := agentproctest.GenerateProcess(t, fs)
expectedProcs[proc.PID] = proc

sc.EXPECT().
Expand All @@ -70,14 +68,14 @@ func TestList(t *testing.T) {

// 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, rootDir)
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, rootDir)
actualProcs, err := agentproc.List(fs, sc)
require.NoError(t, err)
require.Len(t, actualProcs, 3)
require.Len(t, actualProcs, len(expectedProcs))
for _, proc := range actualProcs {
expected, ok := expectedProcs[proc.PID]
require.True(t, ok)
Expand All @@ -94,11 +92,10 @@ func TestList(t *testing.T) {
fs = afero.NewMemMapFs()
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
expectedProcs = make(map[int32]agentproc.Process)
rootDir = agentproc.DefaultProcDir
)

for i := 0; i < 3; i++ {
proc := agentproctest.GenerateProcess(t, fs, rootDir)
proc := agentproctest.GenerateProcess(t, fs)
expectedProcs[proc.PID] = proc

sc.EXPECT().
Expand All @@ -108,14 +105,14 @@ func TestList(t *testing.T) {

// 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, rootDir)
proc := agentproctest.GenerateProcess(t, fs)
sc.EXPECT().
Kill(proc.PID, syscall.Signal(0)).
Return(syscall.ESRCH)

actualProcs, err := agentproc.List(fs, sc, rootDir)
actualProcs, err := agentproc.List(fs, sc)
require.NoError(t, err)
require.Len(t, actualProcs, 3)
require.Len(t, actualProcs, len(expectedProcs))
for _, proc := range actualProcs {
expected, ok := expectedProcs[proc.PID]
require.True(t, ok)
Expand All @@ -136,15 +133,14 @@ func TestProcess(t *testing.T) {

var (
fs = afero.NewMemMapFs()
dir = agentproc.DefaultProcDir
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
proc = agentproctest.GenerateProcess(t, fs)
expectedScore = -1000
)

err := proc.SetOOMAdj(expectedScore)
require.NoError(t, err)

actualScore, err := afero.ReadFile(fs, fmt.Sprintf("%s/%d/oom_score_adj", dir, proc.PID))
actualScore, err := afero.ReadFile(fs, fmt.Sprintf("/proc/%d/oom_score_adj", proc.PID))
require.NoError(t, err)
require.Equal(t, fmt.Sprintf("%d", expectedScore), strings.TrimSpace(string(actualScore)))
})
Expand Down
Loading