/
exec_streamer.go
145 lines (117 loc) · 3 KB
/
exec_streamer.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
package execstreamer
import (
"bufio"
"fmt"
"io"
"net/http"
"os/exec"
"runtime/debug"
"strings"
"sync"
"time"
)
//ExecStreamer is the streamer interface (built by the ExecStreamerBuilder)
type ExecStreamer interface {
StartExec() (*exec.Cmd, error)
ExecAndWait() error
}
type execStreamer struct {
stdOutAndErrWaitGroup *sync.WaitGroup
ExecutorName string
Exe string
Args []string
Dir string
Env []string
StdoutWriter io.Writer
StdoutPrefix string
StderrWriter io.Writer
StderrPrefix string
AutoFlush bool
DebugInfo string
OnStarted func(startedDetails *StartedDetails)
}
func (e *execStreamer) recoverPanic(description string) {
if r := recover(); r != nil {
defer recover()
fmt.Println(fmt.Sprintf("Exec-Stream-Recovery (%s - debug info: %s): %T %+v. Stack: %s\n------END STACK---------\n", description, e.DebugInfo, r, r, strings.Replace(string(debug.Stack()), "\n", "\\n", -1)))
}
}
func (e *execStreamer) flushIfEnabled(writer io.Writer) {
if e.AutoFlush && writer != nil {
if flusher, ok := writer.(http.Flusher); ok {
defer e.recoverPanic("flushIfEnabled")
if flusher != nil {
flusher.Flush()
}
}
}
}
func (e *execStreamer) handleStdout(stdoutScanner *bufio.Scanner) {
defer e.recoverPanic("handleStdout")
defer e.stdOutAndErrWaitGroup.Done()
for stdoutScanner.Scan() {
fmt.Fprintf(e.StdoutWriter, "%s%s\n", e.StdoutPrefix, stdoutScanner.Text())
e.flushIfEnabled(e.StdoutWriter)
}
}
func (e *execStreamer) handleStderr(stderrScanner *bufio.Scanner) {
defer e.recoverPanic("handleStderr")
defer e.stdOutAndErrWaitGroup.Done()
for stderrScanner.Scan() {
fmt.Fprintf(e.StderrWriter, "%s%s\n", e.StderrPrefix, stderrScanner.Text())
e.flushIfEnabled(e.StderrWriter)
}
}
//StartExec will execute the command using the given executor and return (without waiting for completion) with the exec.Cmd
func (e *execStreamer) StartExec() (*exec.Cmd, error) {
x, err := NewExecutorFromName(e.ExecutorName)
if err != nil {
return nil, err
}
cmd := x.GetCommand(e.Exe, e.Args...)
if e.Dir != "" {
cmd.Dir = e.Dir
}
if len(e.Env) > 0 {
cmd.Env = e.Env
}
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, err
}
stderr, err := cmd.StderrPipe()
if err != nil {
return nil, err
}
err = cmd.Start()
if err != nil {
return nil, err
}
if e.OnStarted != nil {
e.OnStarted(&StartedDetails{
Pid: cmd.Process.Pid,
Time: time.Now(),
})
}
e.stdOutAndErrWaitGroup = &sync.WaitGroup{}
e.stdOutAndErrWaitGroup.Add(1)
stdoutScanner := bufio.NewScanner(stdout)
go e.handleStdout(stdoutScanner)
e.stdOutAndErrWaitGroup.Add(1)
stderrScanner := bufio.NewScanner(stderr)
go e.handleStderr(stderrScanner)
return cmd, nil
}
//ExecAndWait will execute the command using the given executor and wait until completion
func (e *execStreamer) ExecAndWait() error {
cmd, err := e.StartExec()
if err != nil {
return err
}
e.stdOutAndErrWaitGroup.Wait()
err = cmd.Wait()
if err != nil {
return err
}
return nil
}