Use a channel to notify when the connection is closed

pubsub.go

@@ -53,6 +53,7 @@ type TwitchPubSub struct {
 	onceHandlers map[string][]*eventHandlerInstance
 
 	// Responses
+	listening chan interface{}
 	responseCh chan *twitchMessage
 
 	SubscribedTopics []string
@@ -78,8 +79,10 @@ func (t *TwitchPubSub) Open() error {
 
 	t.wsConn = c
 
-	go t.reader(c)
-	go t.pinger()
+	t.listening = make(chan interface{})
+
+	go t.reader(t.wsConn, t.listening)
+	go t.pinger(t.wsConn, t.listening)
 
 	if len(t.SubscribedTopics) > 0 {
 		return t.listen(t.SubscribedTopics)
@@ -91,6 +94,11 @@ func (t *TwitchPubSub) Open() error {
 }
 
 func (t *TwitchPubSub) Close() error {
+	if t.listening != nil {
+		close(t.listening)
+		t.listening = nil
+	}
+
 	if t.wsConn != nil {
 		// To cleanly close a connection, a client should send a close
 		// frame and wait for the server to close the connection.
@@ -197,7 +205,7 @@ func (t *TwitchPubSub) Unlisten(topics []string) error {
 	return nil
 }
 
-func (t *TwitchPubSub) reader(wsConn *websocket.Conn) {
+func (t *TwitchPubSub) reader(wsConn *websocket.Conn, listening <-chan interface{}) {
 	for {
 		var message twitchMessage
 		err := t.wsConn.ReadJSON(&message)
@@ -214,6 +222,10 @@ func (t *TwitchPubSub) reader(wsConn *websocket.Conn) {
 			return
 		}
 
+		select {
+		case <-listening:
+			return
+		default:
 			if message.Type == Pong {
 				// PONG!
 			} else if message.Type == Message {
@@ -245,20 +257,26 @@ func (t *TwitchPubSub) reader(wsConn *websocket.Conn) {
 			}
 		}
 	}
+}
 
-func (t *TwitchPubSub) pinger() {
+func (t *TwitchPubSub) pinger(wsConn *websocket.Conn, listening <-chan interface{}) {
 	ticker := time.NewTicker(150 * time.Second)
 	defer ticker.Stop()
 
 	for {
-		<- ticker.C
-
 		t.wsMutex.Lock()
-		err := t.wsConn.WriteJSON(&twitchMessage{Type: Ping})
+		err := wsConn.WriteJSON(&twitchMessage{Type: Ping})
 		t.wsMutex.Unlock()
 
 		if err != nil {
 			return
 		}
+
+		select {
+		case <-ticker.C:
+		// continue loop and send heartbeat
+		case <-listening:
+			return
+		}
 	}
 }