Websocket client and hub improvements

backend/api/client.go

@@ -123,7 +123,7 @@ func (c *Client) writePump() {
 }
 
 // ServeWs handles websocket requests from the peer.
-func ServeWs(w http.ResponseWriter, r *http.Request) {
+func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		log.Println(err)
@@ -132,8 +132,9 @@ func ServeWs(w http.ResponseWriter, r *http.Request) {
 	client := &Client{
 		conn: conn,
 		send: make(chan []byte, 256),
-		stop: make(chan bool),
+		stop: make(chan bool, 1),
 	}
+	hub.register <- client
 	// Allow collection of memory referenced by the caller by doing all work in
 	// new goroutines.
 	go client.writePump()

backend/api/hub.go

@@ -32,7 +32,7 @@ func NewHub() *Hub {
 	}
 }
 
-func (h *Hub) run() {
+func (h *Hub) Run() {
 	defer func() {
 		for client := range h.clients {
 			client.conn.Close()
@@ -48,6 +48,10 @@ func (h *Hub) run() {
 		case client := <-h.unregister:
 			if _, ok := h.clients[client]; ok {
 				delete(h.clients, client)
+				select {
+				case client.stop <- true:
+				default:
+				}
 			}
 		case message := <-h.broadcast:
 			for client := range h.clients {
@@ -57,7 +61,7 @@ func (h *Hub) run() {
 					delete(h.clients, client)
 				}
 			}
-		case <- h.stop:
+		case <-h.stop:
 			return
 		}
 	}

backend/api/session.go

@@ -124,7 +124,7 @@ func JoinRoom(client *Client, event *Event) GameError {
 func InitalizeRoom(client *Client, newRoomCode string) room {
 	initRoom := NewGame(newRoomCode)
 	initRoom.Hub = NewHub()
-	go initRoom.Hub.run()
+	go initRoom.Hub.Run()
 	go initRoom.gameTimeout()
 	initRoom.Hub.register <- client
 	return initRoom

backend/main.go

@@ -50,8 +50,11 @@ func main() {
 		log.Panicf("Error: unable initalize store: %s", err)
 	}
 
+	hub := api.NewHub()
+	go hub.Run()
+
 	http.HandleFunc("/api/ws", func(httpWriter http.ResponseWriter, httpRequest *http.Request) {
-		api.ServeWs(httpWriter, httpRequest)
+		api.ServeWs(hub, httpWriter, httpRequest)
 	})
 
 	http.HandleFunc("/api/health", func(w http.ResponseWriter, r *http.Request) {