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:
parent
8bc6eee32c
commit
8759d12617
1 changed files with 11 additions and 3 deletions
14
rpcserver.go
14
rpcserver.go
|
@ -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 "+
|
||||||
|
|
Loading…
Reference in a new issue