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 } }