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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy