allow client to disconnect when handlers are done

added a waitgroup on websocketClient to keep track of handler
goroutines specific to the client
This commit is contained in:
Javed Khan 2014-08-01 14:22:35 +05:30
parent 8bc6eee32c
commit 8759d12617

View file

@ -143,6 +143,7 @@ type websocketClient struct {
allRequests chan []byte allRequests chan []byte
responses chan []byte responses chan []byte
quit chan struct{} // closed on disconnect quit chan struct{} // closed on disconnect
wg sync.WaitGroup
} }
func newWebsocketClient(c *websocket.Conn, authenticated bool, remoteAddr string) *websocketClient { func newWebsocketClient(c *websocket.Conn, authenticated bool, remoteAddr string) *websocketClient {
@ -790,7 +791,7 @@ out:
break out break out
} }
f := s.HandlerClosure(raw.Method) f := s.HandlerClosure(raw.Method)
s.wg.Add(1) wsc.wg.Add(1)
go func(request []byte, raw *rawRequest) { go func(request []byte, raw *rawRequest) {
resp := f(request, raw) resp := f(request, raw)
mresp, err := json.Marshal(resp) mresp, err := json.Marshal(resp)
@ -799,7 +800,7 @@ out:
} }
_ = wsc.send(mresp) _ = wsc.send(mresp)
s.wg.Done() wsc.wg.Done()
}(request, &raw) }(request, &raw)
} }
@ -818,6 +819,9 @@ out:
<-s.notificationHandlerQuit <-s.notificationHandlerQuit
} }
// allow client to disconnect after all handler goroutines are done
wsc.wg.Wait()
close(wsc.responses)
s.wg.Done() s.wg.Done()
} }
@ -826,7 +830,11 @@ func (s *rpcServer) WebsocketClientSend(wsc *websocketClient) {
out: out:
for { for {
select { select {
case response := <-wsc.responses: case response, ok := <-wsc.responses:
if !ok {
// client disconnected
break out
}
err := wsc.conn.SetWriteDeadline(time.Now().Add(deadline)) err := wsc.conn.SetWriteDeadline(time.Now().Add(deadline))
if err != nil { if err != nil {
log.Warnf("Cannot set write deadline on "+ log.Warnf("Cannot set write deadline on "+