use WithTimeout

handlers.go

@@ -355,7 +355,6 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 	defer s.clientsMu.Unlock()
 	s.clients[conn] = struct{}{}
 	ticker := time.NewTicker(pingPeriod)
-	stop := make(chan struct{})
 
 	ip := conn.RemoteAddr().String()
 	if realIP := r.Header.Get("X-Forwarded-For"); realIP != "" {
@@ -374,12 +373,13 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 		)
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	// reader
 	go func() {
 		defer func() {
+			cancel()
 			ticker.Stop()
-			stop <- struct{}{}
-			close(stop)
 			s.clientsMu.Lock()
 			if _, ok := s.clients[conn]; ok {
 				conn.Close()
@@ -388,8 +388,6 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 			}
 			s.clientsMu.Unlock()
 			s.Log.Infof("disconnected from %s", ip)
-
-			ctx.Done()
 		}()
 
 		conn.SetReadLimit(maxMessageSize)
@@ -432,17 +430,16 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 				continue
 			}
 
-			go s.handleMessage(context.TODO(), ws, message, store)
+			go s.handleMessage(ctx, ws, message, store)
 		}
 	}()
 
 	// writer
 	go func() {
 		defer func() {
+			cancel()
 			ticker.Stop()
 			conn.Close()
-			for range stop {
-			}
 		}()
 
 		for {
@@ -454,7 +451,7 @@ func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 					return
 				}
 				s.Log.Infof("pinging for %s", ip)
-			case <-stop:
+			case <-ctx.Done():
 				return
 			}
 		}

listener.go

@@ -76,6 +76,7 @@ func removeListenerId(ws *WebSocket, id string) {
 func removeListener(ws *WebSocket) {
 	listenersMutex.Lock()
 	defer listenersMutex.Unlock()
+	clear(listeners[ws])
 	delete(listeners, ws)
 }