Merge pull request #364 from rodrigorfk/fix-ws-panic

fix: panic when the WebSocket endpoint is under load

pkg/api/http/echows.go

@@ -3,6 +3,7 @@ package http
 import (
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -28,13 +29,19 @@ func (s *Server) echoWsHandler(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
+	var wg sync.WaitGroup
+	wg.Add(1)
+
 	defer c.Close()
 	done := make(chan struct{})
 	defer close(done)
 	in := make(chan interface{})
-	defer close(in)
 	go s.writeWs(c, in)
-	go s.sendHostWs(c, in, done)
+	go s.sendHostWs(c, in, done, &wg)
+	go func() {
+		defer close(in)
+		wg.Wait()
+	}()
 	for {
 		_, message, err := c.ReadMessage()
 		if err != nil {
@@ -54,7 +61,7 @@ func (s *Server) echoWsHandler(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan struct{}) {
+func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan struct{}, wg *sync.WaitGroup) {
 	ticker := time.NewTicker(5 * time.Second)
 	defer ticker.Stop()
 	for {
@@ -70,6 +77,7 @@ func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan s
 			in <- status
 		case <-done:
 			s.logger.Debug("websocket exit")
+			wg.Done()
 			return
 		}
 	}