Skip to content

Commit 8230247

Browse files
sreyasreya
committed
refactor into build files
1 parent cbcb854 commit 8230247

File tree

6 files changed

+88
-49
lines changed

6 files changed

+88
-49
lines changed

‎agent/agent.go

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ const (
5353
ProtocolDial = "dial"
5454
)
5555

56+
// EnvProcMemNice determines whether we attempt to manage
57+
// process CPU and OOM Killer priority.
5658
const EnvProcMemNice = "CODER_PROC_MEMNICE_ENABLE"
5759

5860
type Options struct {
@@ -73,7 +75,6 @@ type Options struct {
7375
ReportMetadataInterval time.Duration
7476
ServiceBannerRefreshInterval time.Duration
7577
Syscaller agentproc.Syscaller
76-
ProcessManagementTick <-chan time.Time
7778
ModifiedProcesses chan []*agentproc.Process
7879
}
7980

@@ -128,7 +129,7 @@ func New(options Options) Agent {
128129
}
129130

130131
if options.Syscaller == nil {
131-
options.Syscaller = agentproc.UnixSyscaller{}
132+
options.Syscaller = agentproc.NewSyscaller()
132133
}
133134

134135
ctx, cancelFunc := context.WithCancel(context.Background())
@@ -155,7 +156,6 @@ func New(options Options) Agent {
155156
subsystems: options.Subsystems,
156157
addresses: options.Addresses,
157158
syscaller: options.Syscaller,
158-
processManagementTick: options.ProcessManagementTick,
159159
modifiedProcs: options.ModifiedProcesses,
160160

161161
prometheusRegistry: prometheusRegistry,
@@ -209,11 +209,10 @@ type agent struct {
209209

210210
connCountReconnectingPTY atomic.Int64
211211

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

219218
func (a *agent) TailnetConn() *tailnet.Conn {
@@ -1299,9 +1298,12 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
12991298

13001299
manage()
13011300

1301+
ticker := time.NewTicker(time.Second)
1302+
defer ticker.Stop()
1303+
13021304
for {
13031305
select {
1304-
case <-a.processManagementTick:
1306+
case <-ticker.C:
13051307
manage()
13061308
case <-ctx.Done():
13071309
return
@@ -1313,7 +1315,7 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
13131315
const (
13141316
procDir = agentproc.DefaultProcDir
13151317
niceness = 10
1316-
oomScoreAdj = -1000
1318+
oomScoreAdj = -500
13171319
)
13181320

13191321
procs, err := agentproc.List(a.filesystem, a.syscaller, agentproc.DefaultProcDir)
@@ -1357,6 +1359,11 @@ func (a *agent) manageProcessPriority(ctx context.Context) ([]*agentproc.Process
13571359
)
13581360
continue
13591361
}
1362+
1363+
// We only want processes that don't have a nice value set
1364+
// so we don't override user nice values.
1365+
// Getpriority actually returns priority for the nice value
1366+
// which is niceness + 20, so here 20 = a niceness of 0 (aka unset).
13601367
if score != 20 {
13611368
a.logger.Error(ctx, "skipping process due to custom niceness",
13621369
slog.F("name", proc.Name()),

‎agent/agent_test.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2409,7 +2409,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24092409
var (
24102410
expectedProcs = map[int32]agentproc.Process{}
24112411
fs = afero.NewMemMapFs()
2412-
ticker = make(chan time.Time)
24132412
syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
24142413
modProcs = make(chan []*agentproc.Process)
24152414
logger = slog.Make(sloghuman.Sink(io.Discard))
@@ -2439,7 +2438,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24392438
}
24402439

24412440
_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
2442-
o.ProcessManagementTick = ticker
24432441
o.Syscaller = syscaller
24442442
o.ModifiedProcesses = modProcs
24452443
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
@@ -2454,7 +2452,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24542452
expected, ok := expectedProcs[actual.PID]
24552453
require.True(t, ok)
24562454
if expected.PID == 1 {
2457-
expectedScore = "-1000"
2455+
expectedScore = "-500"
24582456
}
24592457

24602458
score, err := afero.ReadFile(fs, filepath.Join(actual.Dir, "oom_score_adj"))
@@ -2469,7 +2467,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24692467
var (
24702468
expectedProcs = map[int32]agentproc.Process{}
24712469
fs = afero.NewMemMapFs()
2472-
ticker = make(chan time.Time)
24732470
syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
24742471
modProcs = make(chan []*agentproc.Process)
24752472
logger = slog.Make(sloghuman.Sink(io.Discard))
@@ -2494,7 +2491,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
24942491
}
24952492

24962493
_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
2497-
o.ProcessManagementTick = ticker
24982494
o.Syscaller = syscaller
24992495
o.ModifiedProcesses = modProcs
25002496
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}

‎agent/agentproc/syscaller.go

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -2,40 +2,10 @@ package agentproc
22

33
import (
44
"syscall"
5-
6-
"golang.org/x/sys/unix"
7-
"golang.org/x/xerrors"
85
)
96

107
type Syscaller interface {
118
SetPriority(pid int32, priority int) error
129
GetPriority(pid int32) (int, error)
1310
Kill(pid int32, sig syscall.Signal) error
1411
}
15-
16-
type UnixSyscaller struct{}
17-
18-
func (UnixSyscaller) SetPriority(pid int32, nice int) error {
19-
err := unix.Setpriority(unix.PRIO_PROCESS, int(pid), nice)
20-
if err != nil {
21-
return xerrors.Errorf("set priority: %w", err)
22-
}
23-
return nil
24-
}
25-
26-
func (UnixSyscaller) GetPriority(pid int32) (int, error) {
27-
nice, err := unix.Getpriority(0, int(pid))
28-
if err != nil {
29-
return 0, xerrors.Errorf("get priority: %w", err)
30-
}
31-
return nice, nil
32-
}
33-
34-
func (UnixSyscaller) Kill(pid int32, sig syscall.Signal) error {
35-
err := syscall.Kill(int(pid), sig)
36-
if err != nil {
37-
return xerrors.Errorf("kill: %w", err)
38-
}
39-
40-
return nil
41-
}

‎agent/agentproc/syscaller_other.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
//go:build !linux
2+
// +build !linux
3+
4+
package agentproc
5+
6+
import "syscall"
7+
8+
func NewSyscaller() Syscaller {
9+
return nopSyscaller{}
10+
}
11+
12+
type nopSyscaller struct{}
13+
14+
func (nopSyscaller) SetPriority(pid int32, priority int) error {
15+
return nil
16+
}
17+
18+
func (nopSyscaller) GetPriority(pid int32) (int, error) {
19+
return 0, nil
20+
}
21+
22+
func (nopSyscaller) Kill(pid int32, sig syscall.Signal) error {
23+
return nil
24+
}

‎agent/agentproc/syscaller_unix.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
//go:build linux
2+
// +build linux
3+
4+
package agentproc
5+
6+
import (
7+
"syscall"
8+
9+
"golang.org/x/sys/unix"
10+
"golang.org/x/xerrors"
11+
)
12+
13+
func NewSyscaller() Syscaller {
14+
return UnixSyscaller{}
15+
}
16+
17+
type UnixSyscaller struct{}
18+
19+
func (UnixSyscaller) SetPriority(pid int32, nice int) error {
20+
err := unix.Setpriority(unix.PRIO_PROCESS, int(pid), nice)
21+
if err != nil {
22+
return xerrors.Errorf("set priority: %w", err)
23+
}
24+
return nil
25+
}
26+
27+
func (UnixSyscaller) GetPriority(pid int32) (int, error) {
28+
nice, err := unix.Getpriority(0, int(pid))
29+
if err != nil {
30+
return 0, xerrors.Errorf("get priority: %w", err)
31+
}
32+
return nice, nil
33+
}
34+
35+
func (UnixSyscaller) Kill(pid int32, sig syscall.Signal) error {
36+
err := syscall.Kill(int(pid), sig)
37+
if err != nil {
38+
return xerrors.Errorf("kill: %w", err)
39+
}
40+
41+
return nil
42+
}

‎cli/agent.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -287,15 +287,15 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
287287
return resp.SessionToken, nil
288288
},
289289
EnvironmentVariables: map[string]string{
290-
"GIT_ASKPASS": executablePath,
290+
"GIT_ASKPASS": executablePath,
291+
agent.EnvProcMemNice: os.Getenv(agent.EnvProcMemNice),
291292
},
292293
IgnorePorts: ignorePorts,
293294
SSHMaxTimeout: sshMaxTimeout,
294295
Subsystems: subsystems,
295296

296-
PrometheusRegistry: prometheusRegistry,
297-
ProcessManagementTick: procTicker.C,
298-
Syscaller: agentproc.UnixSyscaller{},
297+
PrometheusRegistry: prometheusRegistry,
298+
Syscaller: agentproc.NewSyscaller(),
299299
// Intentionally set this to nil. It's mainly used
300300
// for testing.
301301
ModifiedProcesses: nil,

0 commit comments

Comments
 (0)