From 37833043bd946bd9741a9c05f3c40d422af75b9e Mon Sep 17 00:00:00 2001
From: Kane York <rikingcoding@gmail.com>
Date: Thu, 29 Oct 2015 01:30:25 -0700
Subject: [PATCH] Make sure that the reader thread finishes when main thread

calls CloseConnection
---
 socketserver/internal/server/handlecore.go | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/socketserver/internal/server/handlecore.go b/socketserver/internal/server/handlecore.go
index c2d05ce8..f22bf842 100644
--- a/socketserver/internal/server/handlecore.go
+++ b/socketserver/internal/server/handlecore.go
@@ -147,12 +147,13 @@ func HandleSocketConnection(conn *websocket.Conn) {
 	_clientChan := make(chan ClientMessage)
 	_serverMessageChan := make(chan ClientMessage)
 	_errorChan := make(chan error)
+	stoppedChan := make(chan struct{})
 
 	var client ClientInfo
 	client.MessageChannel = _serverMessageChan
 
 	// Launch receiver goroutine
-	go func(errorChan chan<- error, clientChan chan<- ClientMessage) {
+	go func(errorChan chan<- error, clientChan chan<- ClientMessage, stoppedChan <-chan struct{}) {
 		var msg ClientMessage
 		var messageType int
 		var packet []byte
@@ -171,7 +172,13 @@ func HandleSocketConnection(conn *websocket.Conn) {
 			if msg.MessageID == 0 {
 				continue
 			}
-			clientChan <- msg
+			select {
+			case clientChan <- msg:
+			case <-stoppedChan:
+				close(errorChan)
+				close(clientChan)
+				return
+			}
 		}
 
 		_, isClose := err.(*websocket.CloseError)
@@ -182,7 +189,7 @@ func HandleSocketConnection(conn *websocket.Conn) {
 		close(errorChan)
 		close(clientChan)
 		// exit
-	}(_errorChan, _clientChan)
+	}(_errorChan, _clientChan, stoppedChan)
 
 	conn.SetPongHandler(func(pongBody string) error {
 		client.pingCount = 0
@@ -244,6 +251,8 @@ RunLoop:
 		}
 	}()
 
+	close(stoppedChan)
+
 	// Stop getting messages...
 	UnsubscribeAll(&client)