fix websocket closing again

lib/jsonrpc/client.go

@@ -88,6 +88,7 @@ func NewMergeClient(addr string, namespace string, outs []interface{}, requestHe
 	}
 
 	stop := make(chan struct{})
+	exiting := make(chan struct{})
 	c.requests = make(chan clientRequest)
 
 	handlers := map[string]rpcHandler{}
@@ -96,6 +97,7 @@ func NewMergeClient(addr string, namespace string, outs []interface{}, requestHe
 		handler:  handlers,
 		requests: c.requests,
 		stop:     stop,
+		exiting:  exiting,
 	}).handleWsConn(context.TODO())
 
 	for _, handler := range outs {
@@ -122,6 +124,7 @@ func NewMergeClient(addr string, namespace string, outs []interface{}, requestHe
 
 	return func() {
 		close(stop)
+		<-exiting
 	}, nil
 }
 

lib/jsonrpc/server.go

@@ -52,6 +52,7 @@ func (s *RPCServer) handleWS(ctx context.Context, w http.ResponseWriter, r *http
 	(&wsConn{
 		conn:    c,
 		handler: s.methods,
+		exiting: make(chan struct{}),
 	}).handleWsConn(ctx)
 
 	if err := c.Close(); err != nil {

lib/jsonrpc/websocket.go

@@ -43,6 +43,7 @@ type wsConn struct {
 	handler  handlers
 	requests <-chan clientRequest
 	stop     <-chan struct{}
+	exiting  chan struct{}
 
 	// incoming messages
 	incoming    chan io.Reader
@@ -389,6 +390,7 @@ func (c *wsConn) handleWsConn(ctx context.Context) {
 
 	c.registerCh = make(chan outChanReg)
 	defer close(c.registerCh)
+	defer close(c.exiting)
 
 	// ////