forked from argoproj/pkg
-
Notifications
You must be signed in to change notification settings - Fork 0
/
exec.go
175 lines (147 loc) · 3.92 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
package exec
import (
"bytes"
"fmt"
"os"
"os/exec"
"runtime"
"strings"
"time"
"github.com/pkg/errors"
log "github.com/sirupsen/logrus"
"github.com/argoproj/pkg/rand"
)
var ErrWaitPIDTimeout = fmt.Errorf("Timed out waiting for PID to complete")
var Unredacted = Redact(nil)
type CmdError struct {
Args string
Stderr string
Cause error
}
func (ce *CmdError) Error() string {
res := fmt.Sprintf("`%v` failed %v", ce.Args, ce.Cause)
if ce.Stderr != "" {
res = fmt.Sprintf("%s: %s", res, ce.Stderr)
}
return res
}
func (ce *CmdError) String() string {
return ce.Error()
}
func newCmdError(args string, cause error, stderr string) *CmdError {
return &CmdError{Args: args, Stderr: stderr, Cause: cause}
}
type CmdOpts struct {
Timeout time.Duration
Redactor func(text string) string
}
var DefaultCmdOpts = CmdOpts{
Timeout: time.Duration(0),
Redactor: Unredacted,
}
func Redact(items []string) func(text string) string {
return func(text string) string {
for _, item := range items {
text = strings.Replace(text, item, "******", -1)
}
return text
}
}
// RunCommandExt is a convenience function to run/log a command and return/log stderr in an error upon
// failure.
func RunCommandExt(cmd *exec.Cmd, opts CmdOpts) (string, error) {
logCtx := log.WithFields(log.Fields{"execID": rand.RandString(5)})
redactor := DefaultCmdOpts.Redactor
if opts.Redactor != nil {
redactor = opts.Redactor
}
// log in a way we can copy-and-paste into a terminal
args := strings.Join(cmd.Args, " ")
logCtx.WithFields(log.Fields{"dir": cmd.Dir}).Info(redactor(args))
var stdout bytes.Buffer
var stderr bytes.Buffer
cmd.Stdout = &stdout
cmd.Stderr = &stderr
start := time.Now()
err := cmd.Start()
if err != nil {
return "", err
}
done := make(chan error)
go func() { done <- cmd.Wait() }()
// Start a timer
timeout := DefaultCmdOpts.Timeout
if opts.Timeout != time.Duration(0) {
timeout = opts.Timeout
}
var timoutCh <-chan time.Time
if timeout != 0 {
timoutCh = time.NewTimer(timeout).C
}
select {
//noinspection ALL
case <-timoutCh:
_ = cmd.Process.Kill()
output := stdout.String()
logCtx.WithFields(log.Fields{"duration": time.Since(start)}).Debug(redactor(output))
err = newCmdError(redactor(args), fmt.Errorf("timeout after %v", timeout), "")
logCtx.Error(err.Error())
return strings.TrimSuffix(output, "\n"), err
case err := <-done:
if err != nil {
output := stdout.String()
logCtx.WithFields(log.Fields{"duration": time.Since(start)}).Debug(redactor(output))
err := newCmdError(redactor(args), errors.New(redactor(err.Error())), strings.TrimSpace(redactor(stderr.String())))
logCtx.Error(err.Error())
return strings.TrimSuffix(output, "\n"), err
}
}
output := stdout.String()
logCtx.WithFields(log.Fields{"duration": time.Since(start)}).Debug(redactor(output))
return strings.TrimSuffix(output, "\n"), nil
}
func RunCommand(name string, opts CmdOpts, arg ...string) (string, error) {
return RunCommandExt(exec.Command(name, arg...), opts)
}
// WaitPIDOpts are options to WaitPID
type WaitPIDOpts struct {
PollInterval time.Duration
Timeout time.Duration
}
// WaitPID waits for a non-child process id to exit
func WaitPID(pid int, opts ...WaitPIDOpts) error {
if runtime.GOOS != "linux" {
return errors.Errorf("Platform '%s' unsupported", runtime.GOOS)
}
var timeout time.Duration
var pollInterval = time.Second
if len(opts) > 0 {
if opts[0].PollInterval != 0 {
pollInterval = opts[0].PollInterval
}
if opts[0].Timeout != 0 {
timeout = opts[0].Timeout
}
}
path := fmt.Sprintf("/proc/%d", pid)
ticker := time.NewTicker(pollInterval)
defer ticker.Stop()
var timoutCh <-chan time.Time
if timeout != 0 {
timoutCh = time.NewTimer(timeout).C
}
for {
select {
case <-ticker.C:
_, err := os.Stat(path)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return errors.WithStack(err)
}
case <-timoutCh:
return ErrWaitPIDTimeout
}
}
}