fix: instant graceful shutdown for SSE/WS connections and container stop

internal/ws/hub_test.go

@@ -0,0 +1,58 @@
+package ws
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestHubCloseAndGracefulShutdown(t *testing.T) {
+	hub := NewHub()
+
+	// Register multiple dummy SSE clients
+	var sseClients []*SSEClient
+	for i := 0; i < 5; i++ {
+		client := &SSEClient{
+			id:     "sse-test",
+			userID: "user-1",
+			send:   make(chan []byte, 10),
+		}
+		hub.RegisterSSE(client)
+		sseClients = append(sseClients, client)
+	}
+
+	// Verify broadcast works
+	hub.BroadcastToAll("TEST_EVENT", map[string]string{"msg": "hello"})
+
+	// Concurrently invoke Close() and Broadcast() to verify no panics or deadlocks
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 50; i++ {
+			hub.BroadcastToAll("SPAM", map[string]int{"i": i})
+			time.Sleep(1 * time.Millisecond)
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		time.Sleep(5 * time.Millisecond)
+		hub.Close()
+	}()
+
+	wg.Wait()
+
+	// Verify that all SSE client channels are closed
+	for _, c := range sseClients {
+		select {
+		case _, ok := <-c.send:
+			if ok {
+				// Drained a message, check next
+				_, ok = <-c.send
+			}
+		default:
+		}
+	}
+}