Skip to content

Commit 712b328

Browse files
sreyasreya
committed
more tests
1 parent 97e68f4 commit 712b328

File tree

5 files changed

+82
-17
lines changed

5 files changed

+82
-17
lines changed

‎agent/agentexec/cli.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ const (
2121

2222
// CLI runs the agent-exec command. It should only be called by the cli package.
2323
func CLI(args []string, environ []string) error {
24+
runtime.LockOSThread()
25+
2426
if runtime.GOOS != "linux" {
2527
return xerrors.Errorf("agent-exec is only supported on Linux")
2628
}
@@ -42,7 +44,6 @@ func CLI(args []string, environ []string) error {
4244
if err != nil {
4345
return xerrors.Errorf("get default nice score: %w", err)
4446
}
45-
fmt.Println("nice score", nice, "pid", pid)
4647
}
4748

4849
oomscore, ok := envValInt(environ, EnvProcOOMScore)
@@ -54,7 +55,7 @@ func CLI(args []string, environ []string) error {
5455
}
5556
}
5657

57-
err = unix.Setpriority(unix.PRIO_PROCESS, pid, nice)
58+
err = syscall.Setpriority(syscall.PRIO_PROCESS, 0, nice)
5859
if err != nil {
5960
return xerrors.Errorf("set nice score: %w", err)
6061
}
@@ -82,8 +83,8 @@ func defaultNiceScore() (int, error) {
8283
if err != nil {
8384
return 0, xerrors.Errorf("get nice score: %w", err)
8485
}
85-
// Priority is niceness + 20.
86-
score -= 20
86+
// See https://linux.die.net/man/2/setpriority#Notes
87+
score = 20 - score
8788

8889
score += 5
8990
if score > 19 {

‎agent/agentexec/cli_test.go

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,15 @@ func TestCLI(t *testing.T) {
2626

2727
ctx := testutil.Context(t, testutil.WaitMedium)
2828
cmd, path := cmd(ctx, t)
29-
cmd.Env = append(cmd.Env, "CODER_PROC_NICE_SCORE=10")
29+
cmd.Env = append(cmd.Env, "CODER_PROC_NICE_SCORE=12")
3030
cmd.Env = append(cmd.Env, "CODER_PROC_OOM_SCORE=123")
3131
err := cmd.Start()
3232
require.NoError(t, err)
3333
go cmd.Wait()
3434

3535
waitForSentinel(ctx, t, cmd, path)
3636
requireOOMScore(t, cmd.Process.Pid, 123)
37-
requireNiceScore(t, cmd.Process.Pid, 10)
37+
requireNiceScore(t, cmd.Process.Pid, 12)
3838
})
3939

4040
t.Run("Defaults", func(t *testing.T) {
@@ -50,7 +50,6 @@ func TestCLI(t *testing.T) {
5050

5151
expectedNice := expectedNiceScore(t)
5252
expectedOOM := expectedOOMScore(t)
53-
fmt.Println("expected nice", expectedNice, "expected oom", expectedOOM)
5453
requireOOMScore(t, cmd.Process.Pid, expectedOOM)
5554
requireNiceScore(t, cmd.Process.Pid, expectedNice)
5655
})
@@ -61,7 +60,8 @@ func requireNiceScore(t *testing.T, pid int, score int) {
6160

6261
nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
6362
require.NoError(t, err)
64-
require.Equal(t, score, nice)
63+
// See https://linux.die.net/man/2/setpriority#Notes
64+
require.Equal(t, score, 20-nice)
6565
}
6666

6767
func requireOOMScore(t *testing.T, pid int, expected int) {
@@ -108,6 +108,10 @@ func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string)
108108
file = filepath.Join(dir, "sentinel")
109109
//nolint:gosec
110110
cmd = exec.CommandContext(ctx, TestBin, "agent-exec", "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
111+
// We set this so we can also easily kill the sleep process the shell spawns.
112+
cmd.SysProcAttr = &syscall.SysProcAttr{
113+
Setpgid: true,
114+
}
111115
}
112116
cmd.Env = os.Environ()
113117
var buf bytes.Buffer
@@ -118,10 +122,10 @@ func cmd(ctx context.Context, t *testing.T, args ...string) (*exec.Cmd, string)
118122
if t.Failed() {
119123
t.Logf("cmd %q output: %s", cmd.Args, buf.String())
120124
}
121-
122-
// if cmd.Process != nil {
123-
// _ = cmd.Process.Kill()
124-
// }
125+
if cmd.Process != nil {
126+
// We use -cmd.Process.Pid to kill the whole process group.
127+
_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGINT)
128+
}
125129
})
126130
return cmd, file
127131
}
@@ -151,7 +155,7 @@ func expectedNiceScore(t *testing.T) int {
151155
require.NoError(t, err)
152156

153157
// Priority is niceness + 20.
154-
score -= 20
158+
score = 20 - score
155159
score += 5
156160
if score > 19 {
157161
return 19

‎agent/agentexec/exec.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,9 @@ const (
2020
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
2121
// is returned. All instances of exec.Cmd should flow through this function to ensure
2222
// proper resource constraints are applied to the child process.
23-
func CommandContext(ctx context.Context, cmd string, args []string) (*exec.Cmd, error) {
24-
_, enabled := envVal(os.Environ(), EnvProcPrioMgmt)
23+
func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
24+
environ := os.Environ()
25+
_, enabled := envVal(environ, EnvProcPrioMgmt)
2526
if runtime.GOOS != "linux" || !enabled {
2627
return exec.CommandContext(ctx, cmd, args...), nil
2728
}

‎agent/agentexec/exec_test.go

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package agentexec_test
2+
3+
import (
4+
"context"
5+
"os"
6+
"os/exec"
7+
"runtime"
8+
"testing"
9+
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/coder/coder/v2/agent/agentexec"
13+
)
14+
15+
func TestExec(t *testing.T) {
16+
t.Run("NonLinux", func(t *testing.T) {
17+
18+
t.Setenv(agentexec.EnvProcPrioMgmt, "true")
19+
20+
if runtime.GOOS == "linux" {
21+
t.Skip("skipping on linux")
22+
}
23+
24+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
25+
require.NoError(t, err)
26+
require.Equal(t, "sh", cmd.Path)
27+
require.Equal(t, []string{"-c", "sleep"}, cmd.Args[1:])
28+
})
29+
30+
t.Run("Linux", func(t *testing.T) {
31+
32+
t.Run("Disabled", func(t *testing.T) {
33+
if runtime.GOOS != "linux" {
34+
t.Skip("skipping on linux")
35+
}
36+
37+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
38+
require.NoError(t, err)
39+
path, err := exec.LookPath("sh")
40+
require.NoError(t, err)
41+
require.Equal(t, path, cmd.Path)
42+
require.Equal(t, []string{"sh", "-c", "sleep"}, cmd.Args)
43+
})
44+
45+
t.Run("Enabled", func(t *testing.T) {
46+
t.Setenv(agentexec.EnvProcPrioMgmt, "hello")
47+
48+
if runtime.GOOS != "linux" {
49+
t.Skip("skipping on linux")
50+
}
51+
52+
executable, err := os.Executable()
53+
require.NoError(t, err)
54+
55+
cmd, err := agentexec.CommandContext(context.Background(), "sh", "-c", "sleep")
56+
require.NoError(t, err)
57+
require.Equal(t, executable, cmd.Path)
58+
require.Equal(t, []string{executable, "agent-exec", "sh", "-c", "sleep"}, cmd.Args)
59+
})
60+
})
61+
}

‎cli/root.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,6 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command {
121121
r.update(),
122122
r.whoami(),
123123

124-
// Hidden
125-
r.agentExec(),
126124
r.expCmd(),
127125
r.gitssh(),
128126
r.support(),

0 commit comments

Comments
 (0)