From 7ca16dfe70799e73422dd28d2a37ebb5ae06757d Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@conformal.com>
Date: Mon, 14 Jul 2014 09:24:41 -0500
Subject: [PATCH] Synchronize notifications and client gr shutdown.

The responses chan for a websocket client was being closed by one of
the websocket goroutines, but it was not the only sender to this
channel.  There was also the notification handler, run by the server
to handle notifications to all websocket clients.  It was possible to
hit cases where sends to this channel would still occur (the select
statement doesn't guarantee that the picked channel operation won't
panic, even if there's another that won't).  To fix this, wait on the
client being removed from the notification group, or if the server is
already shutting down, wait on the notification handler completely
closing, to ensure that no more sends to the channel will occur,
before closing the channel.

Fixes #110.
---
 rpcserver.go | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/rpcserver.go b/rpcserver.go
index 171c386..05a0687 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -218,9 +218,10 @@ type rpcServer struct {
 
 	requests chan handlerJob
 
-	addWSClient    chan *websocketClient
-	removeWSClient chan *websocketClient
-	broadcasts     chan []byte
+	addWSClient             chan *websocketClient
+	removeWSClient          chan *websocketClient
+	broadcasts              chan []byte
+	notificationHandlerQuit chan struct{}
 
 	quit chan struct{}
 }
@@ -239,11 +240,12 @@ func newRPCServer(listenAddrs []string, maxClients, maxWebsockets int64) (*rpcSe
 			// Allow all origins.
 			CheckOrigin: func(r *http.Request) bool { return true },
 		},
-		requests:       make(chan handlerJob),
-		addWSClient:    make(chan *websocketClient),
-		removeWSClient: make(chan *websocketClient),
-		broadcasts:     make(chan []byte),
-		quit:           make(chan struct{}),
+		requests:                make(chan handlerJob),
+		addWSClient:             make(chan *websocketClient),
+		removeWSClient:          make(chan *websocketClient),
+		broadcasts:              make(chan []byte),
+		notificationHandlerQuit: make(chan struct{}),
+		quit: make(chan struct{}),
 	}
 
 	// Check for existence of cert file and key file
@@ -742,6 +744,17 @@ out:
 			break out
 		}
 	}
+
+	// Remove websocket client from notification group, or if the server is
+	// shutting down, wait until the notification handler has finished
+	// running.  This is needed to ensure that no more notifications will be
+	// sent to the client's responses chan before it's closed below.
+	select {
+	case s.removeWSClient <- wsc:
+	case <-s.quit:
+		<-s.notificationHandlerQuit
+	}
+
 	close(wsc.responses)
 	s.wg.Done()
 }
@@ -939,6 +952,7 @@ out:
 			break out
 		}
 	}
+	close(s.notificationHandlerQuit)
 	s.wg.Done()
 }