forked from coder/coder
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexec.go
More file actions
149 lines (122 loc) · 4.11 KB
/
exec.go
File metadata and controls
149 lines (122 loc) · 4.11 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
package agentexec
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strconv"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/pty"
)
const (
// EnvProcPrioMgmt is the environment variable that determines whether
// we attempt to manage process CPU and OOM Killer priority.
EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"
EnvProcOOMScore = "CODER_PROC_OOM_SCORE"
EnvProcNiceScore = "CODER_PROC_NICE_SCORE"
// unset is set to an invalid value for nice and oom scores.
unset = -2000
)
var DefaultExecer Execer = execer{}
// Execer defines an abstraction for creating exec.Cmd variants. It's unfortunately
// necessary because we need to be able to wrap child processes with "coder agent-exec"
// for templates that expect the agent to manage process priority.
type Execer interface {
// CommandContext returns an exec.Cmd that calls "coder agent-exec" prior to exec'ing
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal exec.Cmd
// is returned. All instances of exec.Cmd should flow through this function to ensure
// proper resource constraints are applied to the child process.
CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd
// PTYCommandContext returns an pty.Cmd that calls "coder agent-exec" prior to exec'ing
// the provided command if CODER_PROC_PRIO_MGMT is set, otherwise a normal pty.Cmd
// is returned. All instances of pty.Cmd should flow through this function to ensure
// proper resource constraints are applied to the child process.
PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd
}
func NewExecer() (Execer, error) {
_, enabled := os.LookupEnv(EnvProcPrioMgmt)
if runtime.GOOS != "linux" || !enabled {
return DefaultExecer, nil
}
executable, err := os.Executable()
if err != nil {
return nil, xerrors.Errorf("get executable: %w", err)
}
bin, err := filepath.EvalSymlinks(executable)
if err != nil {
return nil, xerrors.Errorf("eval symlinks: %w", err)
}
oomScore, ok := envValInt(EnvProcOOMScore)
if !ok {
oomScore = unset
}
niceScore, ok := envValInt(EnvProcNiceScore)
if !ok {
niceScore = unset
}
return priorityExecer{
binPath: bin,
oomScore: oomScore,
niceScore: niceScore,
}, nil
}
type execer struct{}
func (execer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
return exec.CommandContext(ctx, cmd, args...)
}
func (execer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
return pty.CommandContext(ctx, cmd, args...)
}
type priorityExecer struct {
binPath string
oomScore int
niceScore int
}
func (e priorityExecer) CommandContext(ctx context.Context, cmd string, args ...string) *exec.Cmd {
cmd, args = e.agentExecCmd(cmd, args...)
return exec.CommandContext(ctx, cmd, args...)
}
func (e priorityExecer) PTYCommandContext(ctx context.Context, cmd string, args ...string) *pty.Cmd {
cmd, args = e.agentExecCmd(cmd, args...)
return pty.CommandContext(ctx, cmd, args...)
}
func (e priorityExecer) agentExecCmd(cmd string, args ...string) (string, []string) {
execArgs := []string{"agent-exec"}
if e.oomScore != unset {
execArgs = append(execArgs, oomScoreArg(e.oomScore))
}
if e.niceScore != unset {
execArgs = append(execArgs, niceScoreArg(e.niceScore))
}
execArgs = append(execArgs, "--", cmd)
execArgs = append(execArgs, args...)
return e.binPath, execArgs
}
// envValInt searches for a key in a list of environment variables and parses it to an int.
// If the key is not found or cannot be parsed, returns 0 and false.
func envValInt(key string) (int, bool) {
val, ok := os.LookupEnv(key)
if !ok {
return 0, false
}
i, err := strconv.Atoi(val)
if err != nil {
return 0, false
}
return i, true
}
// The following are flags used by the agent-exec command. We use flags instead of
// environment variables to avoid having to deal with a caller overriding the
// environment variables.
const (
niceFlag = "coder-nice"
oomFlag = "coder-oom"
)
func niceScoreArg(score int) string {
return fmt.Sprintf("--%s=%d", niceFlag, score)
}
func oomScoreArg(score int) string {
return fmt.Sprintf("--%s=%d", oomFlag, score)
}