Skip to content

Commit f4b864e

Browse files
sreyasreya
committed
add a proper test for proc management
1 parent 760cbb3 commit f4b864e

File tree

8 files changed

+136
-50
lines changed

8 files changed

+136
-50
lines changed

‎agent/agent.go

Lines changed: 32 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ type Options struct {
7373
ReportMetadataInterval time.Duration
7474
ServiceBannerRefreshInterval time.Duration
7575
Syscaller agentproc.Syscaller
76-
ProcessManagementInterval time.Duration
76+
ProcessManagementTick <-chan time.Time
77+
ModifiedProcesses chan []*agentproc.Process
7778
}
7879

7980
type Client interface {
@@ -130,10 +131,6 @@ func New(options Options) Agent {
130131
options.Syscaller = agentproc.UnixSyscaller{}
131132
}
132133

133-
if options.ProcessManagementInterval == 0 {
134-
options.ProcessManagementInterval = time.Second
135-
}
136-
137134
ctx, cancelFunc := context.WithCancel(context.Background())
138135
a := &agent{
139136
tailnetListenPort: options.TailnetListenPort,
@@ -158,7 +155,8 @@ func New(options Options) Agent {
158155
subsystems: options.Subsystems,
159156
addresses: options.Addresses,
160157
syscaller: options.Syscaller,
161-
processManagementInterval: options.ProcessManagementInterval,
158+
processManagementTick: options.ProcessManagementTick,
159+
modifiedProcs: options.ModifiedProcesses,
162160

163161
prometheusRegistry: prometheusRegistry,
164162
metrics: newAgentMetrics(prometheusRegistry),
@@ -211,10 +209,11 @@ type agent struct {
211209

212210
connCountReconnectingPTY atomic.Int64
213211

214-
prometheusRegistry *prometheus.Registry
215-
metrics *agentMetrics
216-
processManagementInterval time.Duration
217-
syscaller agentproc.Syscaller
212+
prometheusRegistry *prometheus.Registry
213+
metrics *agentMetrics
214+
processManagementTick <-chan time.Time
215+
modifiedProcs chan []*agentproc.Process
216+
syscaller agentproc.Syscaller
218217
}
219218

220219
func (a *agent) TailnetConn() *tailnet.Conn {
@@ -1275,9 +1274,6 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
12751274
var prioritizedProcs = []string{"coder"}
12761275

12771276
func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
1278-
ticker := time.NewTicker(a.processManagementInterval)
1279-
defer ticker.Stop()
1280-
12811277
if val := a.envVars[EnvProcMemNice]; val == "" || runtime.GOOS != "linux" {
12821278
a.logger.Info(ctx, "process priority not enabled, agent will not manage process niceness/oom_score_adj ",
12831279
slog.F("env_var", EnvProcMemNice),
@@ -1287,42 +1283,45 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
12871283
return
12881284
}
12891285

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-
)
1286+
manage := func() {
1287+
// Do once before falling into loop.
1288+
procs, err := a.manageProcessPriority(ctx)
1289+
if err != nil {
1290+
a.logger.Error(ctx, "manage process priority",
1291+
slog.F("dir", agentproc.DefaultProcDir),
1292+
slog.Error(err),
1293+
)
1294+
}
1295+
if a.modifiedProcs != nil {
1296+
a.modifiedProcs <- procs
1297+
}
12961298
}
12971299

1300+
manage()
1301+
12981302
for {
12991303
select {
1300-
case <-ticker.C:
1301-
if err := a.manageProcessPriority(ctx); err != nil {
1302-
a.logger.Error(ctx, "manage process priority",
1303-
slog.F("dir", agentproc.DefaultProcDir),
1304-
slog.Error(err),
1305-
)
1306-
}
1307-
1304+
case <-a.processManagementTick:
1305+
manage()
13081306
case <-ctx.Done():
13091307
return
13101308
}
13111309
}
13121310
}
13131311

1314-
func (a *agent) manageProcessPriority(ctx context.Context) error {
1312+
func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process, error) {
13151313
const (
13161314
procDir = agentproc.DefaultProcDir
13171315
niceness = 10
1318-
oomScoreAdj = 100
1316+
oomScoreAdj = -1000
13191317
)
13201318

13211319
procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
13221320
if err != nil {
1323-
return xerrors.Errorf("list: %w", err)
1321+
return nil, xerrors.Errorf("list: %w", err)
13241322
}
13251323

1324+
modProcs := []*agentproc.Process{}
13261325
for _, proc := range procs {
13271326
// Trim off the path e.g. "./coder" -> "coder"
13281327
name := filepath.Base(proc.Name())
@@ -1339,6 +1338,7 @@ func (a *agent) manageProcessPriority(ctx context.Context) error {
13391338
)
13401339
continue
13411340
}
1341+
modProcs = append(modProcs, proc)
13421342

13431343
a.logger.Debug(ctx, "decreased process oom_score",
13441344
slog.F("name", proc.Name()),
@@ -1382,8 +1382,9 @@ func (a *agent) manageProcessPriority(ctx context.Context) error {
13821382
slog.F("pid", proc.PID),
13831383
slog.F("niceness", niceness),
13841384
)
1385+
modProcs = append(modProcs, proc)
13851386
}
1386-
return nil
1387+
return modProcs, nil
13871388
}
13881389

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

‎agent/agent_test.go

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,12 @@ import (
2121
"strings"
2222
"sync"
2323
"sync/atomic"
24+
"syscall"
2425
"testing"
2526
"time"
2627

2728
scp "github.com/bramvdbogaerde/go-scp"
29+
"github.com/golang/mock/gomock"
2830
"github.com/google/uuid"
2931
"github.com/pion/udp"
3032
"github.com/pkg/sftp"
@@ -44,6 +46,8 @@ import (
4446
"cdr.dev/slog/sloggers/sloghuman"
4547
"cdr.dev/slog/sloggers/slogtest"
4648
"github.com/coder/coder/v2/agent"
49+
"github.com/coder/coder/v2/agent/agentproc"
50+
"github.com/coder/coder/v2/agent/agentproc/agentproctest"
4751
"github.com/coder/coder/v2/agent/agentssh"
4852
"github.com/coder/coder/v2/agent/agenttest"
4953
"github.com/coder/coder/v2/coderd/httpapi"
@@ -2399,6 +2403,66 @@ func TestAgent_Metrics_SSH(t *testing.T) {
23992403
func TestAgent_ManageProcessPriority(t *testing.T) {
24002404
t.Parallel()
24012405

2406+
t.Run("OK", func(t *testing.T) {
2407+
t.Parallel()
2408+
2409+
var (
2410+
expectedProcs = map[int32]agentproc.Process{}
2411+
fs = afero.NewMemMapFs()
2412+
ticker = make(chan time.Time)
2413+
syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
2414+
modProcs = make(chan []*agentproc.Process)
2415+
logger = slog.Make(sloghuman.Sink(io.Discard))
2416+
)
2417+
2418+
// Create some processes.
2419+
for i := 0; i < 4; i++ {
2420+
// Create a prioritized process.
2421+
var proc agentproc.Process
2422+
if i == 0 {
2423+
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir,
2424+
func(p *agentproc.Process) {
2425+
p.CmdLine = "./coder\x00agent\x00--no-reap"
2426+
p.PID = 1
2427+
},
2428+
)
2429+
} else {
2430+
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
2431+
syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
2432+
syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
2433+
}
2434+
syscaller.EXPECT().
2435+
Kill(proc.PID, syscall.Signal(0)).
2436+
Return(nil)
2437+
2438+
expectedProcs[proc.PID] = proc
2439+
}
2440+
2441+
_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
2442+
o.ProcessManagementTick = ticker
2443+
o.Syscaller = syscaller
2444+
o.ModifiedProcesses = modProcs
2445+
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
2446+
o.Filesystem = fs
2447+
o.Logger = logger
2448+
})
2449+
actualProcs := <-modProcs
2450+
require.Len(t, actualProcs, 4)
2451+
2452+
for _, actual := range actualProcs {
2453+
expectedScore := "0"
2454+
expected, ok := expectedProcs[actual.PID]
2455+
require.True(t, ok)
2456+
if expected.PID == 1 {
2457+
expectedScore = "-1000"
2458+
}
2459+
2460+
score, err := afero.ReadFile(fs, filepath.Join(actual.Dir, "oom_score_adj"))
2461+
require.NoError(t, err)
2462+
require.Equal(t, expectedScore, strings.TrimSpace(string(score)))
2463+
}
2464+
})
2465+
24022466
t.Run("DisabledByDefault", func(t *testing.T) {
24032467
t.Parallel()
24042468

‎agent/agentproc/agentproctest/doc.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
// Package agentproctest contains utility functions
2+
// for testing process management in the agent.
3+
package agentproctest
4+
5+
//go:generate mockgen -destination ./syscallermock.go -package agentproctest github.com/coder/coder/v2/agent/agentproc Syscaller

‎agent/agentproc/agentproctest/proc.go

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,12 @@ import (
1111
"github.com/coder/coder/v2/cryptorand"
1212
)
1313

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

1717
pid, err := cryptorand.Intn(1<<31 - 1)
1818
require.NoError(t, err)
1919

20-
err = fs.MkdirAll(fmt.Sprintf("/%s/%d", dir, pid), 0555)
21-
require.NoError(t, err)
22-
2320
arg1, err := cryptorand.String(5)
2421
require.NoError(t, err)
2522

@@ -31,13 +28,26 @@ func GenerateProcess(t *testing.T, fs afero.Fs, dir string) agentproc.Process {
3128

3229
cmdline := fmt.Sprintf("%s\x00%s\x00%s", arg1, arg2, arg3)
3330

34-
err = afero.WriteFile(fs, fmt.Sprintf("/%s/%d/cmdline", dir, pid), []byte(cmdline), 0444)
35-
require.NoError(t, err)
36-
37-
return agentproc.Process{
38-
PID: int32(pid),
31+
process := agentproc.Process{
3932
CmdLine: cmdline,
40-
Dir: fmt.Sprintf("%s/%d", dir, pid),
33+
PID: int32(pid),
4134
FS: fs,
4235
}
36+
37+
for _, mut := range muts {
38+
mut(&process)
39+
}
40+
41+
process.Dir = fmt.Sprintf("%s/%d", dir, process.PID)
42+
43+
err = fs.MkdirAll(process.Dir, 0555)
44+
require.NoError(t, err)
45+
46+
err = afero.WriteFile(fs, fmt.Sprintf("%s/cmdline", process.Dir), []byte(process.CmdLine), 0444)
47+
require.NoError(t, err)
48+
49+
err = afero.WriteFile(fs, fmt.Sprintf("%s/oom_score_adj", process.Dir), []byte("0"), 0444)
50+
require.NoError(t, err)
51+
52+
return process
4353
}

‎agent/agentproc/syscallermock_test.go renamed to ‎agent/agentproc/agentproctest/syscallermock.go

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎agent/agentproc/doc.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
11
// Package agentproc contains logic for interfacing with local
22
// processes running in the same context as the agent.
33
package agentproc
4-
5-
//go:generate mockgen -destination ./syscallermock_test.go -package agentproc_test github.com/coder/coder/v2/agent/agentproc Syscaller

‎agent/agentproc/proc_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ func TestList(t *testing.T) {
2323

2424
var (
2525
fs = afero.NewMemMapFs()
26-
sc = NewMockSyscaller(gomock.NewController(t))
26+
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
2727
expectedProcs = make(map[int32]agentproc.Process)
2828
rootDir = agentproc.DefaultProcDir
2929
)
@@ -54,7 +54,7 @@ func TestList(t *testing.T) {
5454

5555
var (
5656
fs = afero.NewMemMapFs()
57-
sc = NewMockSyscaller(gomock.NewController(t))
57+
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
5858
expectedProcs = make(map[int32]agentproc.Process)
5959
rootDir = agentproc.DefaultProcDir
6060
)
@@ -92,7 +92,7 @@ func TestList(t *testing.T) {
9292

9393
var (
9494
fs = afero.NewMemMapFs()
95-
sc = NewMockSyscaller(gomock.NewController(t))
95+
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
9696
expectedProcs = make(map[int32]agentproc.Process)
9797
rootDir = agentproc.DefaultProcDir
9898
)
@@ -153,7 +153,7 @@ func TestProcess(t *testing.T) {
153153
t.Parallel()
154154

155155
var (
156-
sc = NewMockSyscaller(gomock.NewController(t))
156+
sc = agentproctest.NewMockSyscaller(gomock.NewController(t))
157157
proc = &agentproc.Process{
158158
PID: 32,
159159
}

‎cli/agent.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"cdr.dev/slog/sloggers/slogjson"
3030
"cdr.dev/slog/sloggers/slogstackdriver"
3131
"github.com/coder/coder/v2/agent"
32+
"github.com/coder/coder/v2/agent/agentproc"
3233
"github.com/coder/coder/v2/agent/reaper"
3334
"github.com/coder/coder/v2/buildinfo"
3435
"github.com/coder/coder/v2/cli/clibase"
@@ -267,6 +268,8 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
267268
subsystems = append(subsystems, subsystem)
268269
}
269270

271+
procTicker := time.NewTicker(time.Second)
272+
defer procTicker.Stop()
270273
agnt := agent.New(agent.Options{
271274
Client: client,
272275
Logger: logger,
@@ -290,7 +293,12 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
290293
SSHMaxTimeout: sshMaxTimeout,
291294
Subsystems: subsystems,
292295

293-
PrometheusRegistry: prometheusRegistry,
296+
PrometheusRegistry: prometheusRegistry,
297+
ProcessManagementTick: procTicker.C,
298+
Syscaller: agentproc.UnixSyscaller{},
299+
// Intentionally set this to nil. It's mainly used
300+
// for testing.
301+
ModifiedProcesses: nil,
294302
})
295303

296304
prometheusSrvClose := ServeHandler(ctx, logger, prometheusMetricsHandler(prometheusRegistry, logger), prometheusAddress, "prometheus")

0 commit comments

Comments
 (0)