Skip to content

Commit d795bf2

Browse files
sreyasreya
committed
move setupagent into agenttest
1 parent 6d601ee commit d795bf2

File tree

4 files changed

+267
-174
lines changed

4 files changed

+267
-174
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.

0 commit comments

Comments
 (0)