mirrored from https://skia.googlesource.com/buildbot
-
Notifications
You must be signed in to change notification settings - Fork 65
/
exec.go
397 lines (357 loc) · 12.1 KB
/
exec.go
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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
/*
A wrapper around the os/exec package that supports timeouts and testing.
Example usage:
Simple command with argument:
err := Run(&Command{
Name: "touch",
Args: []string{file},
})
More complicated example:
output := bytes.Buffer{}
err := Run(&Command{
Name: "make",
Args: []string{"all"},
// Set environment:
Env: []string{fmt.Sprintf("GOPATH=%s", projectGoPath)},
// Set working directory:
Dir: projectDir,
// Capture output:
CombinedOutput: &output,
// Set a timeout:
Timeout: 10*time.Minute,
})
Inject a Run function for testing:
var actualCommand *Command
SetRunForTesting(func(command *Command) error {
actualCommand = command
return nil
})
defer SetRunForTesting(DefaultRun)
TestCodeCallingRun()
expect.Equal(t, "touch", actualCommand.Name)
expect.Equal(t, 1, len(actualCommand.Args))
expect.Equal(t, file, actualCommand.Args[0])
*/
package exec
import (
"bytes"
"context"
"fmt"
"io"
"os"
osexec "os/exec"
"strings"
"syscall"
"time"
"go.skia.org/infra/go/sklog"
"go.skia.org/infra/go/util"
)
const (
TIMEOUT_ERROR_PREFIX = "Command killed since it took longer than"
)
type Verbosity int
const (
Info Verbosity = iota
Debug
Silent
)
var (
contextKey = &struct{}{}
defaultContext = &execContext{DefaultRun}
WriteInfoLog = WriteLog{LogFunc: sklog.Infof}
WriteWarningLog = WriteLog{LogFunc: sklog.Warningf}
)
// WriteLog implements the io.Writer interface and writes to the given log function.
type WriteLog struct {
LogFunc func(format string, args ...interface{})
}
func (wl WriteLog) Write(p []byte) (n int, err error) {
wl.LogFunc("%s", string(p))
return len(p), nil
}
type Command struct {
// Name of the command, as passed to osexec.Command. Can be the path to a binary or the
// name of a command that osexec.Lookpath can find.
Name string
// Arguments of the command, not including Name.
Args []string
// The environment of the process. If nil, the current process's environment is used.
Env []string
// If Env is non-nil, adds the current process's entire environment to Env, excluding
// variables that are set in Env.
InheritEnv bool
// If Env is non-nil, adds the current process's PATH to Env. Do not include PATH in Env.
InheritPath bool
// The working directory of the command. If nil, runs in the current process's current
// directory.
Dir string
// See docs for osexec.Cmd.Stdin.
Stdin io.Reader
// If true, duplicates stdout of the command to WriteInfoLog.
LogStdout bool
// Sends the stdout of the command to this Writer, e.g. os.File or bytes.Buffer.
Stdout io.Writer
// If true, duplicates stderr of the command to WriteWarningLog.
LogStderr bool
// Sends the stderr of the command to this Writer, e.g. os.File or bytes.Buffer.
Stderr io.Writer
// Sends the combined stdout and stderr of the command to this Writer, in addition to
// Stdout and Stderr. Only one goroutine will write at a time. Note: the Go runtime seems to
// combine stdout and stderr into one stream as long as LogStdout and LogStderr are false
// and Stdout and Stderr are nil. Otherwise, the stdout and stderr of the command could be
// arbitrarily reordered when written to CombinedOutput.
CombinedOutput io.Writer
// Time limit to wait for the command to finish. No limit if not specified.
Timeout time.Duration
// Whether to log when the command starts.
Verbose Verbosity
// SysProcAttr holds optional, operating system-specific attributes.
// Run passes it to os.StartProcess as the os.ProcAttr's Sys field.
SysProcAttr *syscall.SysProcAttr
}
type Process interface {
Kill() error
}
// Divides commandLine at spaces; treats the first token as the program name and the other tokens
// as arguments. Note: don't expect this function to do anything smart with quotes or escaped
// spaces.
func ParseCommand(commandLine string) Command {
programAndArgs := strings.Split(commandLine, " ")
return Command{Name: programAndArgs[0], Args: programAndArgs[1:]}
}
// Given io.Writers or nils, return a single writer that writes to all, or nil if no non-nil
// writers. Also checks for non-nil io.Writer containing a nil value.
// http://devs.cloudimmunity.com/gotchas-and-common-mistakes-in-go-golang/index.html#nil_in_nil_in_vals
func squashWriters(writers ...io.Writer) io.Writer {
nonNil := []io.Writer{}
for _, writer := range writers {
if writer != nil && !util.IsNil(writer) {
nonNil = append(nonNil, writer)
}
}
switch len(nonNil) {
case 0:
return nil
case 1:
return nonNil[0]
default:
return io.MultiWriter(nonNil...)
}
}
// Returns the Env, Name, and Args of command joined with spaces. Does not perform any quoting.
func DebugString(command *Command) string {
result := ""
result += strings.Join(command.Env, " ")
if len(command.Env) != 0 {
result += " "
}
result += command.Name
if len(command.Args) != 0 {
result += " "
}
result += strings.Join(command.Args, " ")
return result
}
func createCmd(command *Command) *osexec.Cmd {
cmd := osexec.Command(command.Name, command.Args...)
if len(command.Env) != 0 {
cmd.Env = command.Env
if command.InheritEnv {
existing := make(map[string]bool, len(command.Env))
for _, s := range command.Env {
existing[strings.SplitN(s, "=", 2)[0]] = true
}
for _, s := range os.Environ() {
if !existing[strings.SplitN(s, "=", 2)[0]] {
cmd.Env = append(cmd.Env, s)
}
}
} else if command.InheritPath {
cmd.Env = append(cmd.Env, "PATH="+os.Getenv("PATH"))
}
}
cmd.Dir = command.Dir
cmd.Stdin = command.Stdin
var stdoutLog io.Writer
if command.LogStdout {
stdoutLog = WriteInfoLog
}
cmd.Stdout = squashWriters(stdoutLog, command.Stdout, command.CombinedOutput)
var stderrLog io.Writer
if command.LogStderr {
stderrLog = WriteWarningLog
}
cmd.Stderr = squashWriters(stderrLog, command.Stderr, command.CombinedOutput)
if command.SysProcAttr != nil {
cmd.SysProcAttr = command.SysProcAttr
}
return cmd
}
func start(command *Command, cmd *osexec.Cmd) error {
if command.Verbose != Silent {
dirMsg := ""
if cmd.Dir != "" {
dirMsg = " with CWD " + cmd.Dir
}
if command.Verbose == Info {
sklog.Infof("Executing '%s' (where %s is %s)%s", DebugString(command), command.Name, cmd.Path, dirMsg)
} else if command.Verbose == Debug {
sklog.Debugf("Executing '%s' (where %s is %s)%s", DebugString(command), command.Name, cmd.Path, dirMsg)
}
}
err := cmd.Start()
if err != nil {
return fmt.Errorf("Unable to start command %s: %s", DebugString(command), err)
}
return nil
}
func waitSimple(command *Command, cmd *osexec.Cmd) error {
err := cmd.Wait()
if err != nil {
return fmt.Errorf("Command exited with %s: %s", err, DebugString(command))
}
return nil
}
func wait(command *Command, cmd *osexec.Cmd) error {
if command.Timeout == 0 {
return waitSimple(command, cmd)
}
done := make(chan error)
go func() {
done <- cmd.Wait()
}()
select {
case <-time.After(command.Timeout):
if command.Verbose != Silent {
sklog.Debugf("About to kill command '%s'", DebugString(command))
}
if err := cmd.Process.Kill(); err != nil {
return fmt.Errorf("Failed to kill timed out process: %s", err)
}
if command.Verbose != Silent {
sklog.Debugf("Waiting for command to exit after killing '%s'", DebugString(command))
}
<-done // allow goroutine to exit
return fmt.Errorf("%s %f secs", TIMEOUT_ERROR_PREFIX, command.Timeout.Seconds())
case err := <-done:
if err != nil {
return fmt.Errorf("Command exited with %s: %s", err, DebugString(command))
}
return nil
}
}
// IsTimeout returns true if the specified error was raised due to a command
// timing out.
func IsTimeout(err error) bool {
return strings.Contains(err.Error(), TIMEOUT_ERROR_PREFIX)
}
// DefaultRun can be passed to SetRunForTesting to go back to running commands as normal.
func DefaultRun(command *Command) error {
cmd := createCmd(command)
if err := start(command, cmd); err != nil {
return err
}
return wait(command, cmd)
}
// execContext is a struct used for controlling the execution context of Commands.
type execContext struct {
runFn func(*Command) error
}
// NewContext returns a context.Context instance which uses the given function
// to run Commands.
func NewContext(ctx context.Context, runFn func(*Command) error) context.Context {
newCtx := &execContext{
runFn: runFn,
}
return context.WithValue(ctx, contextKey, newCtx)
}
// getCtx retrieves the Context associated with the context.Context.
func getCtx(ctx context.Context) *execContext {
if v := ctx.Value(contextKey); v != nil {
return v.(*execContext)
}
return defaultContext
}
// See documentation for exec.Run.
func (c *execContext) Run(command *Command) error {
return c.runFn(command)
}
// runSimpleCommand executes the given command. Returns the combined stdout and stderr. May also
// return an error if the command exited with a non-zero status or there is any other error.
func (c *execContext) runSimpleCommand(command *Command) (string, error) {
output := bytes.Buffer{}
// We use a ThreadSafeWriter here because command.CombinedOutput may get
// wrapped with an io.MultiWriter if the caller set command.Stdout or
// command.Stderr, or if either command.LogStdout or command.LogStderr
// is true. In that case, the os/exec package will not be able to
// determine that CombinedOutput is shared between stdout and stderr and
// it will spin up an extra goroutine to write to it, causing a data
// race. We could be smarter and only use the ThreadSafeWriter when we
// know that this will be the case, but relies too much on our knowledge
// of the current os/exec implementation.
command.CombinedOutput = util.NewThreadSafeWriter(&output)
// Setting Verbose to Silent to maintain previous behavior.
command.Verbose = Silent
err := c.Run(command)
result := string(output.Bytes())
if err != nil {
return result, fmt.Errorf("%s; Stdout+Stderr:\n%s", err.Error(), result)
}
return result, nil
}
// See documentation for exec.RunSimple.
func (c *execContext) RunSimple(commandLine string) (string, error) {
cmd := ParseCommand(commandLine)
return c.runSimpleCommand(&cmd)
}
// See documentation for exec.RunCommand.
func (c *execContext) RunCommand(command *Command) (string, error) {
return c.runSimpleCommand(command)
}
// See documentation for exec.RunCwd.
func (c *execContext) RunCwd(cwd string, args ...string) (string, error) {
command := &Command{
Name: args[0],
Args: args[1:],
Dir: cwd,
}
return c.runSimpleCommand(command)
}
// Run runs command and waits for it to finish. If any failure, returns non-nil. If a timeout was
// specified, returns an error once the command has exceeded that timeout.
func Run(ctx context.Context, command *Command) error {
return getCtx(ctx).Run(command)
}
// RunSimple executes the given command line string; the command being run is expected to not care
// what its current working directory is. Returns the combined stdout and stderr. May also return
// an error if the command exited with a non-zero status or there is any other error.
func RunSimple(ctx context.Context, commandLine string) (string, error) {
return getCtx(ctx).RunSimple(commandLine)
}
// RunCommand executes the given command and returns the combined stdout and stderr. May also
// return an error if the command exited with a non-zero status or there is any other error.
func RunCommand(ctx context.Context, command *Command) (string, error) {
return getCtx(ctx).runSimpleCommand(command)
}
// RunCwd executes the given command in the given directory. Returns the combined stdout and
// stderr. May also return an error if the command exited with a non-zero status or there is any
// other error.
func RunCwd(ctx context.Context, cwd string, args ...string) (string, error) {
return getCtx(ctx).RunCwd(cwd, args...)
}
// RunIndefinitely starts the command and then returns. Clients can listen for
// the command to end on the returned channel or kill the process manually
// using the Process handle. The timeout param is ignored if it is set. If
// starting the command returns an error, that error is returned.
func RunIndefinitely(command *Command) (Process, <-chan error, error) {
cmd := createCmd(command)
done := make(chan error)
if err := start(command, cmd); err != nil {
close(done)
return nil, done, err
}
go func() {
done <- cmd.Wait()
}()
return cmd.Process, done, nil
}