forked from concourse/concourse
/
stream_handler.go
90 lines (78 loc) · 1.65 KB
/
stream_handler.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
package connection
import (
"encoding/json"
"fmt"
"io"
"sync"
"code.cloudfoundry.org/garden/transport"
"code.cloudfoundry.org/lager"
)
type streamHandler struct {
log lager.Logger
wg *sync.WaitGroup
}
func newStreamHandler(log lager.Logger) *streamHandler {
return &streamHandler{
log: log,
wg: new(sync.WaitGroup),
}
}
func (sh *streamHandler) streamIn(processWriter io.WriteCloser, stdin io.Reader) {
if stdin == nil {
return
}
go func(processInputStream io.WriteCloser, stdin io.Reader, log lager.Logger) {
if _, err := io.Copy(processInputStream, stdin); err == nil {
processInputStream.Close()
} else {
log.Error("streaming-stdin-payload", err)
}
}(processWriter, stdin, sh.log)
}
func (sh *streamHandler) streamOut(streamWriter io.Writer, streamReader io.Reader) {
sh.wg.Add(1)
go func() {
io.Copy(streamWriter, streamReader)
sh.wg.Done()
}()
}
type waitReturn struct {
exitCode int
err error
}
func (sh *streamHandler) wait(decoder *json.Decoder) <-chan waitReturn {
result := make(chan waitReturn)
go func() {
for {
payload := &transport.ProcessPayload{}
err := decoder.Decode(payload)
if err != nil {
sh.wg.Wait()
result <- waitReturn{
0,
fmt.Errorf("connection: decode failed: %s", err),
}
break
}
if payload.Error != nil {
sh.wg.Wait()
result <- waitReturn{
0,
fmt.Errorf("connection: process error: %s", *payload.Error),
}
break
}
if payload.ExitStatus != nil {
sh.wg.Wait()
result <- waitReturn{
int(*payload.ExitStatus),
nil,
}
break
}
// discard other payloads
}
defer close(result)
}()
return result
}