From eba7fe186ea8464f6ded44c1a855ad704e7e34e9 Mon Sep 17 00:00:00 2001
From: Rodrigo Fior Kuntzer <rodrigo@miro.com>
Date: Thu, 23 May 2024 10:01:01 +0200
Subject: [PATCH] fix: panic when the WebSocket endpoint is under load

---
 pkg/api/http/echows.go | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/pkg/api/http/echows.go b/pkg/api/http/echows.go
index de956ef4..d06dfb73 100644
--- a/pkg/api/http/echows.go
+++ b/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
 		}
 	}