/
streamer.go
112 lines (100 loc) · 2.83 KB
/
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
package log
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/go-chi/render"
"github.com/gorilla/websocket"
"github.com/ubclaunchpad/inertia/daemon/inertiad/res"
)
// Streamer is a multilogger used by the daemon to pipe output to multiple
// places depending on context.
type Streamer struct {
req *http.Request
httpWriter http.ResponseWriter
httpStream bool
socket SocketWriter
io.Writer
}
// StreamerOptions defines configuration for a daemon streamer
type StreamerOptions struct {
Request *http.Request
Stdout io.Writer
Socket SocketWriter
HTTPWriter http.ResponseWriter
HTTPStream bool
}
// NewStreamer creates a new streamer. It must be closed if a Socket is provided,
// and one of Error() or Success() should be called.
func NewStreamer(opts StreamerOptions) *Streamer {
var w io.Writer
if !opts.HTTPStream {
// Attempt to create a writer with websocket
if opts.Socket != nil {
w = &MultiWriter{
writers: []io.Writer{opts.Stdout, NewWebSocketTextWriter(opts.Socket)},
}
} else {
w = opts.Stdout
}
} else {
// Attempt to create a writer with HTTPWriter
w = &MultiWriter{
writers: []io.Writer{opts.Stdout, opts.HTTPWriter},
}
}
return &Streamer{
req: opts.Request,
httpWriter: opts.HTTPWriter,
httpStream: opts.HTTPStream,
socket: opts.Socket,
Writer: w,
}
}
// GetSocketWriter retrieves the socketwriter as an io.Writer
func (s *Streamer) GetSocketWriter() (io.Writer, error) {
if s.socket != nil {
return NewWebSocketTextWriter(s.socket), nil
}
return nil, errors.New("no websocket active")
}
// Println prints to logger's standard writer
func (s *Streamer) Println(a interface{}) {
fmt.Fprintln(s.Writer, a)
}
// Error directs message and status to http.Error when appropriate
func (s *Streamer) Error(res *res.ErrResponse) {
fmt.Fprintln(s.Writer, res.Error().Error())
if s.socket == nil {
render.Render(s.httpWriter, s.req, res)
} else {
s.Close(CloseOpts{res.Message, res.HTTPStatusCode})
}
}
// Success directs status to Header and sets content type when appropriate
func (s *Streamer) Success(res *res.MsgResponse) {
fmt.Fprintf(s.Writer, "[success %d] %s\n", res.HTTPStatusCode, res.Message)
if s.socket == nil && !s.httpStream {
render.Render(s.httpWriter, s.req, res)
} else {
s.Close(CloseOpts{res.Message, res.HTTPStatusCode})
}
}
// CloseOpts defines options for closing the logger
type CloseOpts struct {
Message string
StatusCode int
}
// Close shuts down the logger
func (s *Streamer) Close(opts ...CloseOpts) error {
if s.socket != nil && !s.httpStream {
if opts != nil && len(opts) > 0 {
return s.socket.CloseHandler()(
websocket.CloseGoingAway,
fmt.Sprintf("status %d: %s", opts[0].StatusCode, opts[0].Message))
}
return s.socket.CloseHandler()(websocket.CloseGoingAway, "connection closed")
}
return nil
}