Skip to content

Commit 46ef05a

Browse files
sreyasreya
committed
only prioritize coder agent
1 parent 11ab047 commit 46ef05a

File tree

6 files changed

+205
-18
lines changed

6 files changed

+205
-18
lines changed

‎agent/agent.go

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"os/user"
1717
"path/filepath"
1818
"runtime"
19+
"runtime/debug"
1920
"sort"
2021
"strconv"
2122
"strings"
@@ -1278,7 +1279,7 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
12781279
}
12791280
}
12801281

1281-
var prioritizedProcs = []string{"coder"}
1282+
var prioritizedProcs = []string{"coder agent"}
12821283

12831284
func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
12841285
if val := a.envVars[EnvProcPrioMgmt]; val == "" || runtime.GOOS != "linux" {
@@ -1290,6 +1291,15 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
12901291
return
12911292
}
12921293

1294+
defer func() {
1295+
if r := recover(); r != nil {
1296+
a.logger.Critical(ctx, "recovered from panic",
1297+
slog.F("panic", r),
1298+
slog.F("stack", string(debug.Stack())),
1299+
)
1300+
}
1301+
}()
1302+
12931303
if a.processManagementTick == nil {
12941304
ticker := time.NewTicker(time.Second)
12951305
defer ticker.Stop()
@@ -1333,15 +1343,17 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
13331343

13341344
for _, proc := range procs {
13351345
logger = a.logger.With(
1336-
slog.F("name", proc.Name()),
1346+
slog.F("cmd", proc.Cmd()),
13371347
slog.F("pid", proc.PID),
13381348
)
13391349

1340-
// Trim off the path e.g. "./coder" -> "coder"
1341-
name := filepath.Base(proc.Name())
1350+
containsFn := func(e string) bool {
1351+
contains := strings.Contains(proc.Cmd(), e)
1352+
return contains
1353+
}
13421354
// If the process is prioritized we should adjust
13431355
// it's oom_score_adj and avoid lowering its niceness.
1344-
if slices.Contains(prioritizedProcs, name) {
1356+
if slices.ContainsFunc[[]string, string](prioritizedProcs, containsFn) {
13451357
err = proc.SetOOMAdj(oomScoreAdj)
13461358
if err != nil {
13471359
logger.Warn(ctx, "unable to set proc oom_score_adj",
@@ -1366,7 +1378,7 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
13661378
// so we don't override user nice values.
13671379
// Getpriority actually returns priority for the nice value
13681380
// which is niceness + 20, so here 20 = a niceness of 0 (aka unset).
1369-
if score != 20 {
1381+
if score != 20 && score != niceness {
13701382
logger.Debug(ctx, "skipping process due to custom niceness",
13711383
slog.F("niceness", score),
13721384
)

‎agent/agent_test.go

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2429,12 +2429,18 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24292429
proc = agentproctest.GenerateProcess(t, fs,
24302430
func(p *agentproc.Process) {
24312431
p.CmdLine = "./coder\x00agent\x00--no-reap"
2432-
p.PID = 1
2432+
p.PID = int32(i)
24332433
},
24342434
)
24352435
} else {
2436-
// The rest are peasants.
2437-
proc = agentproctest.GenerateProcess(t, fs)
2436+
proc = agentproctest.GenerateProcess(t, fs,
2437+
func(p *agentproc.Process) {
2438+
// Make the cmd something similar to a prioritized
2439+
// process but differentiate the arguments.
2440+
p.CmdLine = "./coder\x00stat"
2441+
},
2442+
)
2443+
24382444
syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
24392445
syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
24402446
}
@@ -2460,7 +2466,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24602466
expectedScore := "0"
24612467
expected, ok := expectedProcs[actual.PID]
24622468
require.True(t, ok)
2463-
if expected.PID == 1 {
2469+
if expected.PID == 0 {
24642470
expectedScore = "-500"
24652471
}
24662472

‎agent/agentproc/proc_other.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ func (p *Process) SetNiceness(sc Syscaller, score int) error {
1919
return errUnimplemented
2020
}
2121

22-
func (p *Process) Name() string {
22+
func (p *Process) Cmd() string {
2323
return ""
2424
}
2525

‎agent/agentproc/proc_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,16 +170,16 @@ func TestProcess(t *testing.T) {
170170
require.NoError(t, err)
171171
})
172172

173-
t.Run("Name", func(t *testing.T) {
173+
t.Run("Cmd", func(t *testing.T) {
174174
t.Parallel()
175175

176176
var (
177177
proc = &agentproc.Process{
178178
CmdLine: "helloworld\x00--arg1\x00--arg2",
179179
}
180-
expectedName = "helloworld"
180+
expectedName = "helloworld --arg1 --arg2"
181181
)
182182

183-
require.Equal(t, expectedName, proc.Name())
183+
require.Equal(t, expectedName, proc.Cmd())
184184
})
185185
}

‎agent/agentproc/proc_unix.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,10 @@ func (p *Process) SetNiceness(sc Syscaller, score int) error {
115115
return nil
116116
}
117117

118-
func (p *Process) Name() string {
119-
args := strings.Split(p.CmdLine, "\x00")
120-
// Split will always return at least one element.
121-
return args[0]
118+
func (p *Process) Cmd() string {
119+
return strings.Join(p.cmdLine(), " ")
120+
}
121+
122+
func (p *Process) cmdLine() []string {
123+
return strings.Split(p.CmdLine, "\x00")
122124
}

0 commit comments

Comments
 (0)