Skip to content

Commit ba137c1

Browse files
sreyasreya
committed
some minor agent tests
1 parent 6d601ee commit ba137c1

File tree

4 files changed

+141
-71
lines changed

4 files changed

+141
-71
lines changed

‎agent/agent.go

Lines changed: 96 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ type Options struct {
7373
ReportMetadataInterval time.Duration
7474
ServiceBannerRefreshInterval time.Duration
7575
Syscaller agentproc.Syscaller
76+
ProcessManagementInterval time.Duration
7677
}
7778

7879
type Client interface {
@@ -125,6 +126,14 @@ func New(options Options) Agent {
125126
prometheusRegistry = prometheus.NewRegistry()
126127
}
127128

129+
if options.Syscaller == nil {
130+
options.Syscaller = agentproc.UnixSyscaller{}
131+
}
132+
133+
if options.ProcessManagementInterval == 0 {
134+
options.ProcessManagementInterval = time.Second
135+
}
136+
128137
ctx, cancelFunc := context.WithCancel(context.Background())
129138
a := &agent{
130139
tailnetListenPort: options.TailnetListenPort,
@@ -148,6 +157,8 @@ func New(options Options) Agent {
148157
sshMaxTimeout: options.SSHMaxTimeout,
149158
subsystems: options.Subsystems,
150159
addresses: options.Addresses,
160+
syscaller: options.Syscaller,
161+
processManagementInterval: options.ProcessManagementInterval,
151162

152163
prometheusRegistry: prometheusRegistry,
153164
metrics: newAgentMetrics(prometheusRegistry),
@@ -200,9 +211,10 @@ type agent struct {
200211

201212
connCountReconnectingPTY atomic.Int64
202213

203-
prometheusRegistry *prometheus.Registry
204-
metrics *agentMetrics
205-
syscaller agentproc.Syscaller
214+
prometheusRegistry *prometheus.Registry
215+
metrics *agentMetrics
216+
processManagementInterval time.Duration
217+
syscaller agentproc.Syscaller
206218
}
207219

208220
func (a *agent) TailnetConn() *tailnet.Conn {
@@ -1263,15 +1275,9 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
12631275
var prioritizedProcs = []string{"coder"}
12641276

12651277
func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
1266-
ticker := time.NewTicker(time.Minute)
1278+
ticker := time.NewTicker(a.processManagementInterval)
12671279
defer ticker.Stop()
12681280

1269-
const (
1270-
procDir = agentproc.DefaultProcDir
1271-
niceness = 10
1272-
oomScoreAdj = 100
1273-
)
1274-
12751281
if val := a.envVars[EnvProcMemNice]; val == "" || runtime.GOOS != "linux" {
12761282
a.logger.Info(ctx, "process priority not enabled, agent will not manage process niceness/oom_score_adj ",
12771283
slog.F("env_var", EnvProcMemNice),
@@ -1281,81 +1287,103 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
12811287
return
12821288
}
12831289

1290+
// Do once before falling into loop.
1291+
if err := a.manageProcessPriority(ctx); err != nil {
1292+
a.logger.Error(ctx, "manage process priority",
1293+
slog.F("dir", agentproc.DefaultProcDir),
1294+
slog.Error(err),
1295+
)
1296+
}
1297+
12841298
for {
12851299
select {
12861300
case <-ticker.C:
1287-
procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
1288-
if err != nil {
1289-
a.logger.Error(ctx, "failed to list procs",
1301+
if err := a.manageProcessPriority(ctx); err != nil {
1302+
a.logger.Error(ctx, "manage process priority",
12901303
slog.F("dir", agentproc.DefaultProcDir),
12911304
slog.Error(err),
12921305
)
1293-
continue
12941306
}
1295-
for _, proc := range procs {
1296-
// Trim off the path e.g. "./coder" -> "coder"
1297-
name := filepath.Base(proc.Name())
1298-
// If the process is prioritized we should adjust
1299-
// it's oom_score_adj and avoid lowering its niceness.
1300-
if slices.Contains(prioritizedProcs, name) {
1301-
err = proc.SetOOMAdj(oomScoreAdj)
1302-
if err != nil {
1303-
a.logger.Error(ctx, "unable to set proc oom_score_adj",
1304-
slog.F("name", proc.Name()),
1305-
slog.F("pid", proc.PID),
1306-
slog.F("oom_score_adj", oomScoreAdj),
1307-
slog.Error(err),
1308-
)
1309-
continue
1310-
}
13111307

1312-
a.logger.Debug(ctx, "decreased process oom_score",
1313-
slog.F("name", proc.Name()),
1314-
slog.F("pid", proc.PID),
1315-
slog.F("oom_score_adj", oomScoreAdj),
1316-
)
1317-
continue
1318-
}
1308+
case <-ctx.Done():
1309+
return
1310+
}
1311+
}
1312+
}
13191313

1320-
score, err := proc.Niceness(a.syscaller)
1321-
if err != nil {
1322-
a.logger.Error(ctx, "unable to get proc niceness",
1323-
slog.F("name", proc.Name()),
1324-
slog.F("pid", proc.PID),
1325-
slog.Error(err),
1326-
)
1327-
continue
1328-
}
1329-
if score != 20 {
1330-
a.logger.Error(ctx, "skipping process due to custom niceness",
1331-
slog.F("name", proc.Name()),
1332-
slog.F("pid", proc.PID),
1333-
slog.F("niceness", score),
1334-
)
1335-
continue
1336-
}
1314+
func (a *agent) manageProcessPriority(ctx context.Context) error {
1315+
const (
1316+
procDir = agentproc.DefaultProcDir
1317+
niceness = 10
1318+
oomScoreAdj = 100
1319+
)
13371320

1338-
err = proc.SetNiceness(a.syscaller, niceness)
1339-
if err != nil {
1340-
a.logger.Error(ctx, "unable to set proc niceness",
1341-
slog.F("name", proc.Name()),
1342-
slog.F("pid", proc.PID),
1343-
slog.F("niceness", niceness),
1344-
slog.Error(err),
1345-
)
1346-
continue
1347-
}
1321+
procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
1322+
if err != nil {
1323+
return xerrors.Errorf("list: %w", err)
1324+
}
13481325

1349-
a.logger.Debug(ctx, "deprioritized process",
1326+
for _, proc := range procs {
1327+
// Trim off the path e.g. "./coder" -> "coder"
1328+
name := filepath.Base(proc.Name())
1329+
// If the process is prioritized we should adjust
1330+
// it's oom_score_adj and avoid lowering its niceness.
1331+
if slices.Contains(prioritizedProcs, name) {
1332+
err = proc.SetOOMAdj(oomScoreAdj)
1333+
if err != nil {
1334+
a.logger.Error(ctx, "unable to set proc oom_score_adj",
13501335
slog.F("name", proc.Name()),
13511336
slog.F("pid", proc.PID),
1352-
slog.F("niceness", niceness),
1337+
slog.F("oom_score_adj", oomScoreAdj),
1338+
slog.Error(err),
13531339
)
1340+
continue
13541341
}
1355-
case <-ctx.Done():
1356-
return
1342+
1343+
a.logger.Debug(ctx, "decreased process oom_score",
1344+
slog.F("name", proc.Name()),
1345+
slog.F("pid", proc.PID),
1346+
slog.F("oom_score_adj", oomScoreAdj),
1347+
)
1348+
continue
13571349
}
1350+
1351+
score, err := proc.Niceness(a.syscaller)
1352+
if err != nil {
1353+
a.logger.Error(ctx, "unable to get proc niceness",
1354+
slog.F("name", proc.Name()),
1355+
slog.F("pid", proc.PID),
1356+
slog.Error(err),
1357+
)
1358+
continue
1359+
}
1360+
if score != 20 {
1361+
a.logger.Error(ctx, "skipping process due to custom niceness",
1362+
slog.F("name", proc.Name()),
1363+
slog.F("pid", proc.PID),
1364+
slog.F("niceness", score),
1365+
)
1366+
continue
1367+
}
1368+
1369+
err = proc.SetNiceness(a.syscaller, niceness)
1370+
if err != nil {
1371+
a.logger.Error(ctx, "unable to set proc niceness",
1372+
slog.F("name", proc.Name()),
1373+
slog.F("pid", proc.PID),
1374+
slog.F("niceness", niceness),
1375+
slog.Error(err),
1376+
)
1377+
continue
1378+
}
1379+
1380+
a.logger.Debug(ctx, "deprioritized process",
1381+
slog.F("name", proc.Name()),
1382+
slog.F("pid", proc.PID),
1383+
slog.F("niceness", niceness),
1384+
)
13581385
}
1386+
return nil
13591387
}
13601388

13611389
// isClosed returns whether the API is closed or not.

‎agent/agent_test.go

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import (
4141
"tailscale.com/tailcfg"
4242

4343
"cdr.dev/slog"
44+
"cdr.dev/slog/sloggers/sloghuman"
4445
"cdr.dev/slog/sloggers/slogtest"
4546
"github.com/coder/coder/v2/agent"
4647
"github.com/coder/coder/v2/agent/agentssh"
@@ -2395,6 +2396,48 @@ func TestAgent_Metrics_SSH(t *testing.T) {
23952396
require.NoError(t, err)
23962397
}
23972398

2399+
func TestAgent_ManageProcessPriority(t *testing.T) {
2400+
t.Parallel()
2401+
2402+
t.Run("DisabledByDefault", func(t *testing.T) {
2403+
t.Parallel()
2404+
2405+
if runtime.GOOS != "linux" {
2406+
t.Skip("Skipping non-linux environment")
2407+
}
2408+
2409+
var buf bytes.Buffer
2410+
log := slog.Make(sloghuman.Sink(&buf))
2411+
2412+
_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
2413+
o.Logger = log
2414+
})
2415+
2416+
require.Eventually(t, func() bool {
2417+
return strings.Contains(buf.String(), "process priority not enabled")
2418+
}, testutil.WaitLong, testutil.IntervalFast)
2419+
})
2420+
2421+
t.Run("DisabledForNonLinux", func(t *testing.T) {
2422+
t.Parallel()
2423+
2424+
if runtime.GOOS == "linux" {
2425+
t.Skip("Skipping linux environment")
2426+
}
2427+
2428+
var buf bytes.Buffer
2429+
log := slog.Make(sloghuman.Sink(&buf))
2430+
2431+
_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
2432+
o.Logger = log
2433+
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
2434+
})
2435+
require.Eventually(t, func() bool {
2436+
return strings.Contains(buf.String(), "process priority not enabled")
2437+
}, testutil.WaitLong, testutil.IntervalFast)
2438+
})
2439+
}
2440+
23982441
func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actual []*promgo.MetricFamily) bool {
23992442
t.Helper()
24002443

‎agent/agentproc/agentproctest/proc.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,4 @@ func GenerateProcess(t *testing.T, fs afero.Fs, dir string) agentproc.Process {
4040
Dir: fmt.Sprintf("%s/%d", dir, pid),
4141
FS: fs,
4242
}
43-
4443
}

‎agent/agentproc/syscaller.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ func (UnixSyscaller) GetPriority(pid int32) (int, error) {
3131
return nice, nil
3232
}
3333

34-
func (UnixSyscaller) Kill(pid int, sig syscall.Signal) error {
35-
err := syscall.Kill(pid, sig)
34+
func (UnixSyscaller) Kill(pid int32, sig syscall.Signal) error {
35+
err := syscall.Kill(int(pid), sig)
3636
if err != nil {
3737
return xerrors.Errorf("kill: %w", err)
3838
}

0 commit comments

Comments
 (0)