Skip to content

Commit

Permalink
fix: panic when the WebSocket endpoint is under load
Browse files Browse the repository at this point in the history
  • Loading branch information
rodrigorfk committed May 23, 2024
1 parent f350624 commit eba7fe1
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions pkg/api/http/echows.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package http
import (
"net/http"
"strings"
"sync"
"time"

"github.com/gorilla/websocket"
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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
}
}
Expand Down

0 comments on commit eba7fe1

Please sign in to comment.