Skip to content

Commit e1c1fe9

Browse files
sreyasreya
committed
add a proper test for proc management
1 parent ba137c1 commit e1c1fe9

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)
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