-
Notifications
You must be signed in to change notification settings - Fork 683
/
command.go
176 lines (153 loc) · 3.73 KB
/
command.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
package supervisor
import (
"fmt"
"io"
"os/exec"
"strings"
)
type logger struct {
process *Process
emptyLine bool
}
func (l *logger) Log(prefix, line string) {
if l.emptyLine {
l.process.Log(prefix)
l.emptyLine = false
}
if line == "" {
l.emptyLine = true
} else {
l.process.Logf("%s%s", prefix, line)
l.emptyLine = false
}
}
func (l *logger) LogLines(prefix, str string, err error) {
if strings.HasSuffix(str, "\n") {
str = str[:len(str)-1]
} else {
str += "\\no newline"
}
lines := strings.Split(str, "\n")
for _, line := range lines {
l.Log(prefix, line)
}
if !(err == nil || err == io.EOF) {
l.process.Log(fmt.Sprintf("%v", err))
}
}
type loggingWriter struct {
logger
writer io.Writer
}
func (l *loggingWriter) Write(bytes []byte) (int, error) {
if l.writer == nil {
l.LogLines(" <- ", string(bytes), nil)
return len(bytes), nil
}
n, err := l.writer.Write(bytes)
l.LogLines(" <- ", string(bytes[:n]), err)
return n, err
}
type loggingReader struct {
logger
reader io.Reader
}
func (l *loggingReader) Read(p []byte) (n int, err error) {
n, err = l.reader.Read(p)
l.LogLines(" -> ", string(p[:n]), err)
return n, err
}
// A Cmd is like an os/exec.Cmd, but logs what happens on
// stdin/stdout/stderr, and has a slightly different API.
type Cmd struct {
*exec.Cmd
supervisorProcess *Process
}
func (c *Cmd) pre() {
if c.Stdin != nil {
c.Stdin = &loggingReader{logger: logger{process: c.supervisorProcess}, reader: c.Stdin}
}
c.Stdout = &loggingWriter{logger: logger{process: c.supervisorProcess}, writer: c.Stdout}
c.Stderr = &loggingWriter{logger: logger{process: c.supervisorProcess}, writer: c.Stderr}
c.supervisorProcess.Logf("%s %v", c.Path, c.Args[1:])
}
func (c *Cmd) post(err error) {
if err == nil {
c.supervisorProcess.Logf("%s exited successfully", c.Path)
} else {
if c.ProcessState == nil {
c.supervisorProcess.Logf("%v", err)
} else {
c.supervisorProcess.Logf("%s: %v", c.Path, err)
}
}
}
// Start is like os/exec.Cmd.Start.
func (c *Cmd) Start() error {
c.pre()
return c.Cmd.Start()
}
// Wait is like os/exec.Cmd.Wait.
func (c *Cmd) Wait() error {
err := c.Cmd.Wait()
c.post(err)
return err
}
// Run is like os/exec.Cmd.Run.
func (c *Cmd) Run() error {
c.pre()
err := c.Cmd.Run()
c.post(err)
return err
}
// Command creates a single purpose supervisor and uses it to produce
// and return a *supervisor.Cmd.
func Command(prefix, name string, args ...string) (result *Cmd) {
MustRun(prefix, func(p *Process) error {
result = p.Command(name, args...)
return nil
})
return
}
// Command creates a command that automatically logs inputs, outputs,
// and exit codes to the process logger.
func (p *Process) Command(name string, args ...string) *Cmd {
return &Cmd{exec.Command(name, args...), p}
}
// Capture runs a command with the supplied input and captures the
// output as a string.
func (c *Cmd) Capture(stdin io.Reader) (output string, err error) {
c.Stdin = stdin
out := strings.Builder{}
c.Stdout = &out
err = c.Run()
output = out.String()
return
}
// MustCapture is like Capture, but panics if there is an error.
func (c *Cmd) MustCapture(stdin io.Reader) (output string) {
output, err := c.Capture(stdin)
if err != nil {
panic(err)
}
return output
}
// CaptureErr runs a command with the supplied input and captures
// stdout and stderr as a string.
func (c *Cmd) CaptureErr(stdin io.Reader) (output string, err error) {
c.Stdin = stdin
out := strings.Builder{}
c.Stdout = &out
c.Stderr = &out
err = c.Run()
output = out.String()
return
}
// MustCaptureErr is like CaptureErr, but panics if there is an error.
func (c *Cmd) MustCaptureErr(stdin io.Reader) (output string) {
output, err := c.CaptureErr(stdin)
if err != nil {
panic(err)
}
return output
}