Gracefull stop on websocket

controllers/peer.go

@@ -237,21 +237,42 @@ func (o *PeerController) DeleteState() {
 
 func Websocket(ctx ctx.Context, user string, r http.ResponseWriter, w *http.Request) {
 	websocket.Handler(func(ws *websocket.Conn) {
+		done := make(chan struct{})
+		go func() {
+			var discard interface{}
+			for {
+				if websocket.JSON.Receive(ws, &discard) != nil {
+					close(done)
+					return
+				}
+			}
+		}()
 		defer func() {
 			ws.Close()
-			infrastructure.SearchMu.Lock()
 			if ch, ok := infrastructure.SearchStream[user]; ok {
 				close(ch)
+				infrastructure.SearchMu.Lock()
 				delete(infrastructure.SearchStream, user)
+				infrastructure.SearchMu.Unlock()
 			}
-			infrastructure.SearchMu.Unlock()
+			fmt.Println("CLOSE !")
 			infrastructure.EmitNATS(user, nil, tools.PropalgationMessage{
 				Action:   tools.PB_CLOSE_SEARCH,
 				DataType: tools.PEER.EnumIndex(),
 			})
 		}()
-		for msg := range infrastructure.SearchStream[user] {
-			if websocket.JSON.Send(ws, msg) != nil {
+		for {
+			select {
+			case msg, ok := <-infrastructure.SearchStream[user]:
+				if !ok {
+					continue
+				}
+				if websocket.JSON.Send(ws, msg) != nil {
+					continue
+				}
+			case <-done:
+				return
+			case <-ctx.Done():
 				return
 			}
 		}