Search & Gracefull Close on Websocket

controllers/general.go

@@ -1,6 +1,8 @@
 package controllers
 
 import (
+	"context"
+	"fmt"
 	"net/http"
 	"oc-catalog/infrastructure"
 
@@ -60,25 +62,47 @@ func (o *GeneralController) GetAll() {
 	o.ServeJSON()
 }
 
-func Websocket(user string, groups []string, dataType int, w http.ResponseWriter, r *http.Request) {
+func Websocket(ctx context.Context, user string, groups []string, dataType int, 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()
-			infrastructure.EmitNATS(user, groups, tools.PropalgationMessage{
+			fmt.Println("CLOSE !")
+			infrastructure.EmitNATS(user, nil, tools.PropalgationMessage{
 				Action:   tools.PB_CLOSE_SEARCH,
 				DataType: dataType,
 			})
 		}()
-		for msg := range infrastructure.SearchStream[user] {
-			if websocket.JSON.Send(ws, msg) != nil {
+		for {
+			select {
+			case msg, ok := <-infrastructure.SearchStream[user]:
+				fmt.Println("FOR", msg, ok)
+				if !ok {
+					continue
+				}
+				if websocket.JSON.Send(ws, msg) != nil {
+					continue
+				}
+			case <-done:
+				return
+			case <-ctx.Done():
 				return
 			}
 		}
-	}).ServeHTTP(w, r)
+	}).ServeHTTP(r, w)
 }