@@ -86,24 +86,23 @@ type ErrorCallback func(err string)
86
86
// it must be created using the NewServer function.
87
87
type Server struct {
88
88
impl Discovery
89
- outputChan chan * message
90
- outputWaiter sync.WaitGroup
91
89
userAgent string
92
90
reqProtocolVersion int
93
91
initialized bool
94
92
started bool
95
93
syncStarted bool
96
94
cachedPorts map [string ]* Port
97
95
cachedErr string
96
+ output io.Writer
97
+ outputMutex sync.Mutex
98
98
}
99
99
100
100
// NewServer creates a new discovery server backed by the
101
101
// provided pluggable discovery implementation. To start the server
102
102
// use the Run method.
103
103
func NewServer (impl Discovery ) * Server {
104
104
return & Server {
105
- impl : impl ,
106
- outputChan : make (chan * message ),
105
+ impl : impl ,
107
106
}
108
107
}
109
108
@@ -113,20 +112,20 @@ func NewServer(impl Discovery) *Server {
113
112
// the input stream is closed. In case of IO error the error is
114
113
// returned.
115
114
func (d * Server ) Run (in io.Reader , out io.Writer ) error {
116
- d .startOutputProcessor ( out )
115
+ d .output = out
117
116
reader := bufio .NewReader (in )
118
117
for {
119
118
fullCmd , err := reader .ReadString ('\n' )
120
119
if err != nil {
121
- d .outputChan <- messageError ("command_error" , err .Error ())
120
+ d .send ( messageError ("command_error" , err .Error () ))
122
121
return err
123
122
}
124
123
fullCmd = strings .TrimSpace (fullCmd )
125
124
split := strings .Split (fullCmd , " " )
126
125
cmd := strings .ToUpper (split [0 ])
127
126
128
127
if ! d .initialized && cmd != "HELLO" && cmd != "QUIT" {
129
- d .outputChan <- messageError ("command_error" , fmt .Sprintf ("First command must be HELLO, but got '%s'" , cmd ))
128
+ d .send ( messageError ("command_error" , fmt .Sprintf ("First command must be HELLO, but got '%s'" , cmd ) ))
130
129
continue
131
130
}
132
131
@@ -142,61 +141,62 @@ func (d *Server) Run(in io.Reader, out io.Writer) error {
142
141
case "STOP" :
143
142
d .stop ()
144
143
case "QUIT" :
145
- d .quit ()
144
+ d .impl .Quit ()
145
+ d .send (messageOk ("quit" ))
146
146
return nil
147
147
default :
148
- d .outputChan <- messageError ("command_error" , fmt .Sprintf ("Command %s not supported" , cmd ))
148
+ d .send ( messageError ("command_error" , fmt .Sprintf ("Command %s not supported" , cmd ) ))
149
149
}
150
150
}
151
151
}
152
152
153
153
func (d * Server ) hello (cmd string ) {
154
154
if d .initialized {
155
- d .outputChan <- messageError ("hello" , "HELLO already called" )
155
+ d .send ( messageError ("hello" , "HELLO already called" ) )
156
156
return
157
157
}
158
158
re := regexp .MustCompile (`^(\d+) "([^"]+)"$` )
159
159
matches := re .FindStringSubmatch (cmd )
160
160
if len (matches ) != 3 {
161
- d .outputChan <- messageError ("hello" , "Invalid HELLO command" )
161
+ d .send ( messageError ("hello" , "Invalid HELLO command" ) )
162
162
return
163
163
}
164
164
d .userAgent = matches [2 ]
165
165
v , err := strconv .ParseInt (matches [1 ], 10 , 64 )
166
166
if err != nil {
167
- d .outputChan <- messageError ("hello" , "Invalid protocol version: " + matches [2 ])
167
+ d .send ( messageError ("hello" , "Invalid protocol version: " + matches [2 ]) )
168
168
return
169
169
}
170
170
d .reqProtocolVersion = int (v )
171
171
if err := d .impl .Hello (d .userAgent , 1 ); err != nil {
172
- d .outputChan <- messageError ("hello" , err .Error ())
172
+ d .send ( messageError ("hello" , err .Error () ))
173
173
return
174
174
}
175
- d .outputChan <- & message {
175
+ d .send ( & message {
176
176
EventType : "hello" ,
177
177
ProtocolVersion : 1 , // Protocol version 1 is the only supported for now...
178
178
Message : "OK" ,
179
- }
179
+ })
180
180
d .initialized = true
181
181
}
182
182
183
183
func (d * Server ) start () {
184
184
if d .started {
185
- d .outputChan <- messageError ("start" , "Discovery already STARTed" )
185
+ d .send ( messageError ("start" , "Discovery already STARTed" ) )
186
186
return
187
187
}
188
188
if d .syncStarted {
189
- d .outputChan <- messageError ("start" , "Discovery already START_SYNCed, cannot START" )
189
+ d .send ( messageError ("start" , "Discovery already START_SYNCed, cannot START" ) )
190
190
return
191
191
}
192
192
d .cachedPorts = map [string ]* Port {}
193
193
d .cachedErr = ""
194
194
if err := d .impl .StartSync (d .eventCallback , d .errorCallback ); err != nil {
195
- d .outputChan <- messageError ("start" , "Cannot START: " + err .Error ())
195
+ d .send ( messageError ("start" , "Cannot START: " + err .Error () ))
196
196
return
197
197
}
198
198
d .started = true
199
- d .outputChan <- messageOk ("start" )
199
+ d .send ( messageOk ("start" ) )
200
200
}
201
201
202
202
func (d * Server ) eventCallback (event string , port * Port ) {
@@ -215,99 +215,84 @@ func (d *Server) errorCallback(msg string) {
215
215
216
216
func (d * Server ) list () {
217
217
if ! d .started {
218
- d .outputChan <- messageError ("list" , "Discovery not STARTed" )
218
+ d .send ( messageError ("list" , "Discovery not STARTed" ) )
219
219
return
220
220
}
221
221
if d .syncStarted {
222
- d .outputChan <- messageError ("list" , "discovery already START_SYNCed, LIST not allowed" )
222
+ d .send ( messageError ("list" , "discovery already START_SYNCed, LIST not allowed" ) )
223
223
return
224
224
}
225
225
if d .cachedErr != "" {
226
- d .outputChan <- messageError ("list" , d .cachedErr )
226
+ d .send ( messageError ("list" , d .cachedErr ) )
227
227
return
228
228
}
229
229
ports := []* Port {}
230
230
for _ , port := range d .cachedPorts {
231
231
ports = append (ports , port )
232
232
}
233
- d .outputChan <- & message {
233
+ d .send ( & message {
234
234
EventType : "list" ,
235
235
Ports : & ports ,
236
- }
236
+ })
237
237
}
238
238
239
239
func (d * Server ) startSync () {
240
240
if d .syncStarted {
241
- d .outputChan <- messageError ("start_sync" , "Discovery already START_SYNCed" )
241
+ d .send ( messageError ("start_sync" , "Discovery already START_SYNCed" ) )
242
242
return
243
243
}
244
244
if d .started {
245
- d .outputChan <- messageError ("start_sync" , "Discovery already STARTed, cannot START_SYNC" )
245
+ d .send ( messageError ("start_sync" , "Discovery already STARTed, cannot START_SYNC" ) )
246
246
return
247
247
}
248
248
if err := d .impl .StartSync (d .syncEvent , d .errorEvent ); err != nil {
249
- d .outputChan <- messageError ("start_sync" , "Cannot START_SYNC: " + err .Error ())
249
+ d .send ( messageError ("start_sync" , "Cannot START_SYNC: " + err .Error () ))
250
250
return
251
251
}
252
252
d .syncStarted = true
253
- d .outputChan <- messageOk ("start_sync" )
253
+ d .send ( messageOk ("start_sync" ) )
254
254
}
255
255
256
256
func (d * Server ) stop () {
257
257
if ! d .syncStarted && ! d .started {
258
- d .outputChan <- messageError ("stop" , "Discovery already STOPped" )
258
+ d .send ( messageError ("stop" , "Discovery already STOPped" ) )
259
259
return
260
260
}
261
261
if err := d .impl .Stop (); err != nil {
262
- d .outputChan <- messageError ("stop" , "Cannot STOP: " + err .Error ())
262
+ d .send ( messageError ("stop" , "Cannot STOP: " + err .Error () ))
263
263
return
264
264
}
265
265
d .started = false
266
266
if d .syncStarted {
267
267
d .syncStarted = false
268
268
}
269
- d .outputChan <- messageOk ("stop" )
269
+ d .send ( messageOk ("stop" ) )
270
270
}
271
271
272
272
func (d * Server ) syncEvent (event string , port * Port ) {
273
- d .outputChan <- & message {
273
+ d .send ( & message {
274
274
EventType : event ,
275
275
Port : port ,
276
- }
277
- }
278
-
279
- func (d * Server ) quit () {
280
- d .impl .Quit ()
281
- d .outputChan <- messageOk ("quit" )
282
- close (d .outputChan )
283
- // If we don't wait for all messages
284
- // to be consumed by the output processor
285
- // we risk not printing the "quit" message.
286
- // This may cause issues to consumers of
287
- // the discovery since they expect a message
288
- // that is never sent.
289
- d .outputWaiter .Wait ()
276
+ })
290
277
}
291
278
292
279
func (d * Server ) errorEvent (msg string ) {
293
- d .outputChan <- messageError ("start_sync" , msg )
280
+ d .send ( messageError ("start_sync" , msg ) )
294
281
}
295
282
296
- func (d * Server ) startOutputProcessor (outWriter io.Writer ) {
297
- // Start go routine to serialize messages printing
298
- d .outputWaiter .Add (1 )
299
- go func () {
300
- for msg := range d .outputChan {
301
- data , err := json .MarshalIndent (msg , "" , " " )
302
- if err != nil {
303
- // We are certain that this will be marshalled correctly
304
- // so we don't handle the error
305
- data , _ = json .MarshalIndent (messageError ("command_error" , err .Error ()), "" , " " )
306
- }
307
- fmt .Fprintln (outWriter , string (data ))
308
- }
309
- // We finished consuming all messages, now
310
- // we can exit for real
311
- d .outputWaiter .Done ()
312
- }()
283
+ func (d * Server ) send (msg * message ) {
284
+ data , err := json .MarshalIndent (msg , "" , " " )
285
+ if err != nil {
286
+ // We are certain that this will be marshalled correctly
287
+ // so we don't handle the error
288
+ data , _ = json .MarshalIndent (messageError ("command_error" , err .Error ()), "" , " " )
289
+ }
290
+ data = append (data , '\n' )
291
+
292
+ d .outputMutex .Lock ()
293
+ defer d .outputMutex .Unlock ()
294
+ n , err := d .output .Write (data )
295
+ if n != len (data ) || err != nil {
296
+ panic ("ERROR" )
297
+ }
313
298
}
0 commit comments