Skip to content

feat: implement agent process management #9461

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 21 commits into from
Sep 15, 2023
Merged

feat: implement agent process management #9461

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
4ed4069
feat: implement agent process management
sreya Aug 30, 2023
7e59db6
improve process detection
sreya Aug 31, 2023
8c65216
add agentproc tests
sreya Aug 31, 2023
760cbb3
some minor agent tests
sreya Sep 1, 2023
f4b864e
add a proper test for proc management
sreya Sep 1, 2023
cbcb854
custom nice
sreya Sep 1, 2023
8230247
refactor into build files
sreya Sep 8, 2023
3e1defd
tick tock
sreya Sep 8, 2023
2fe9c70
make fmt
sreya Sep 8, 2023
ef41e9a
pr comments
sreya Sep 12, 2023
05baba0
lint
sreya Sep 12, 2023
478d57c
whoops
sreya Sep 12, 2023
0ced5ce
skip non-linux
sreya Sep 13, 2023
8aaa6d5
skip non-linux
sreya Sep 13, 2023
cea4851
prevent race
sreya Sep 13, 2023
5020eb4
only prioritize ourselves
sreya Sep 14, 2023
11ab047
Revert "only prioritize ourselves"
sreya Sep 14, 2023
46ef05a
only prioritize coder agent
sreya Sep 14, 2023
d132480
remove oom_score_adj
sreya Sep 14, 2023
04ee5cb
defer first
sreya Sep 15, 2023
ffbeab9
avoid resetting niceness for already niced procs
sreya Sep 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
tick tock
  • Loading branch information
@sreya
sreya committed Sep 8, 2023
commit 3e1defd13e92a5878bba6bf2277ce5b51fb7441d
23 changes: 17 additions & 6 deletions agent/agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,10 @@ type Options struct {
ReportMetadataInterval time.Duration
ServiceBannerRefreshInterval time.Duration
Syscaller agentproc.Syscaller
ModifiedProcesses chan []*agentproc.Process
// ModifiedProcesses is used for testing process priority management.
ModifiedProcesses chan []*agentproc.Process
// ProcessManagementTick is used for testing process priority management.
ProcessManagementTick <-chan time.Time
}

type Client interface {
Expand Down Expand Up @@ -157,6 +160,7 @@ func New(options Options) Agent {
addresses: options.Addresses,
syscaller: options.Syscaller,
modifiedProcs: options.ModifiedProcesses,
processManagementTick: options.ProcessManagementTick,

prometheusRegistry: prometheusRegistry,
metrics: newAgentMetrics(prometheusRegistry),
Expand Down Expand Up @@ -211,8 +215,12 @@ type agent struct {

prometheusRegistry *prometheus.Registry
metrics *agentMetrics
modifiedProcs chan []*agentproc.Process
syscaller agentproc.Syscaller

// podifiedProcs is used for testing process priority management.
modifiedProcs chan []*agentproc.Process
// processManagementTick is used for testing process priority management.
processManagementTick <-chan time.Time
}

func (a *agent) TailnetConn() *tailnet.Conn {
Expand Down Expand Up @@ -1283,7 +1291,6 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
}

manage := func() {
// Do once before falling into loop.
procs, err := a.manageProcessPriority(ctx)
if err != nil {
a.logger.Error(ctx, "manage process priority",
Expand All @@ -1296,14 +1303,18 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
}
}

// Do once before falling into loop.
manage()

ticker := time.NewTicker(time.Second)
defer ticker.Stop()
if a.processManagementTick == nil {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
a.processManagementTick = ticker.C
}

for {
select {
case <-ticker.C:
case <-a.processManagementTick:
manage()
case <-ctx.Done():
return
Expand Down
14 changes: 12 additions & 2 deletions agent/agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2410,13 +2410,16 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
expectedProcs = map[int32]agentproc.Process{}
fs = afero.NewMemMapFs()
syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
ticker = make(chan time.Time)
modProcs = make(chan []*agentproc.Process)
logger = slog.Make(sloghuman.Sink(io.Discard))
)

// Create some processes.
for i := 0; i < 4; i++ {
// Create a prioritized process.
// Create a prioritized process. This process should
// have it's oom_score_adj set to -500 and its nice
// score should be untouched.
var proc agentproc.Process
if i == 0 {
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir,
Expand All @@ -2426,6 +2429,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
},
)
} else {
// The rest are peasants.
proc = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
syscaller.EXPECT().SetPriority(proc.PID, 10).Return(nil)
syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
Expand All @@ -2443,6 +2447,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
o.Filesystem = fs
o.Logger = logger
o.ProcessManagementTick = ticker
})
actualProcs := <-modProcs
require.Len(t, actualProcs, 4)
Expand All @@ -2467,20 +2472,22 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
var (
expectedProcs = map[int32]agentproc.Process{}
fs = afero.NewMemMapFs()
ticker = make(chan time.Time)
syscaller = agentproctest.NewMockSyscaller(gomock.NewController(t))
modProcs = make(chan []*agentproc.Process)
logger = slog.Make(sloghuman.Sink(io.Discard))
)

// Create some processes.
for i := 0; i < 2; i++ {
// Create a prioritized process.
proc := agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
syscaller.EXPECT().
Kill(proc.PID, syscall.Signal(0)).
Return(nil)

if i == 0 {
// Set a random nice score. This one should not be adjusted by
// our management loop.
syscaller.EXPECT().GetPriority(proc.PID).Return(25, nil)
} else {
syscaller.EXPECT().GetPriority(proc.PID).Return(20, nil)
Expand All @@ -2496,6 +2503,7 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
o.Filesystem = fs
o.Logger = logger
o.ProcessManagementTick = ticker
})
actualProcs := <-modProcs
// We should ignore the process with a custom nice score.
Expand Down Expand Up @@ -2533,6 +2541,8 @@ func TestAgent_ManageProcessPriority(t *testing.T) {

_, _, _, _, _ = setupAgent(t, agentsdk.Manifest{}, 0, func(c *agenttest.Client, o *agent.Options) {
o.Logger = log
// Try to enable it so that we can assert that non-linux
// environments are truly disabled.
o.EnvironmentVariables = map[string]string{agent.EnvProcMemNice: "1"}
})
require.Eventually(t, func() bool {
Expand Down