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

cmd/server/main.go

@@ -5,6 +5,7 @@ import (
 	"embed"
 	"io/fs"
 	"log/slog"
+	"net"
 	"net/http"
 	"os"
 	"os/signal"
@@ -58,9 +59,15 @@ func main() {
 
 	srv := api.NewServer(cfg, store, wsHub, staticFS)
 
+	serverCtx, serverCancel := context.WithCancel(context.Background())
+	defer serverCancel()
+
 	httpServer := &http.Server{
 		Addr:    ":" + cfg.Port,
 		Handler: srv.Router(),
+		BaseContext: func(l net.Listener) context.Context {
+			return serverCtx
+		},
 		ReadTimeout:  15 * time.Second,
 		WriteTimeout: 30 * time.Second,
 		IdleTimeout:  60 * time.Second,
@@ -79,7 +86,12 @@ func main() {
 	<-stop
 
 	slog.Info("Shutting down MindSift server gracefully...")
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	serverCancel()
+	wsHub.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 	defer cancel()
-	_ = httpServer.Shutdown(ctx)
+	if err := httpServer.Shutdown(ctx); err != nil {
+		slog.Warn("HTTP server shutdown finished with error", "error", err)
+	}
 }

internal/ws/hub.go

@@ -28,6 +28,29 @@ type Client struct {
 	userID          string
 	send            chan []byte
 	subscribedRooms map[string]bool
+	sendMu          sync.RWMutex
+	closed          bool
+}
+
+func (c *Client) closeSend() {
+	c.sendMu.Lock()
+	defer c.sendMu.Unlock()
+	if !c.closed {
+		c.closed = true
+		close(c.send)
+	}
+}
+
+func (c *Client) trySend(data []byte) {
+	c.sendMu.RLock()
+	defer c.sendMu.RUnlock()
+	if c.closed {
+		return
+	}
+	select {
+	case c.send <- data:
+	default:
+	}
 }
 
 type SSEClient struct {
@@ -36,6 +59,29 @@ type SSEClient struct {
 	roomID          string
 	send            chan []byte
 	subscribedRooms map[string]bool
+	sendMu          sync.RWMutex
+	closed          bool
+}
+
+func (c *SSEClient) closeSend() {
+	c.sendMu.Lock()
+	defer c.sendMu.Unlock()
+	if !c.closed {
+		c.closed = true
+		close(c.send)
+	}
+}
+
+func (c *SSEClient) trySend(data []byte) {
+	c.sendMu.RLock()
+	defer c.sendMu.RUnlock()
+	if c.closed {
+		return
+	}
+	select {
+	case c.send <- data:
+	default:
+	}
 }
 
 type docLockItem struct {
@@ -62,6 +108,8 @@ type Hub struct {
 	locks        map[string]*docLockItem
 	docRooms     map[string]*DocRoom
 	docStore     DocStore
+	stopCleanup  chan struct{}
+	closed       bool
 }
 
 func NewHub() *Hub {
@@ -74,6 +122,7 @@ func NewHub() *Hub {
 		sseRooms:     make(map[string]map[*SSEClient]bool),
 		locks:        make(map[string]*docLockItem),
 		docRooms:     make(map[string]*DocRoom),
+		stopCleanup:  make(chan struct{}),
 	}
 	go h.cleanupLocksLoop()
 	return h
@@ -150,7 +199,7 @@ func (h *Hub) Unregister(c *Client) {
 			userRemaining += len(sseRoom)
 		}
 	}
-	close(c.send)
+	c.closeSend()
 
 	// If this user has no other active client connections, release their document locks
 	var releasedDocIDs []string
@@ -250,7 +299,7 @@ func (h *Hub) UnregisterSSE(c *SSEClient) {
 			userRemaining += len(wsRoom)
 		}
 	}
-	close(c.send)
+	c.closeSend()
 
 	var releasedDocIDs []string
 	if userRemaining == 0 && c.userID != "" {
@@ -414,7 +463,11 @@ func (h *Hub) cleanupLocksLoop() {
 	ticker := time.NewTicker(2 * time.Second)
 	defer ticker.Stop()
 
-	for range ticker.C {
+	for {
+		select {
+		case <-h.stopCleanup:
+			return
+		case <-ticker.C:
 			h.mu.Lock()
 			now := time.Now()
 			var expiredDocIDs []string
@@ -431,6 +484,68 @@ func (h *Hub) cleanupLocksLoop() {
 			}
 		}
 	}
+}
+
+// Close gracefully shuts down the Hub, stops background tickers, flushes dirty document rooms to DB,
+// and immediately disconnects all WebSocket and SSE connections so HTTP shutdown completes in milliseconds.
+func (h *Hub) Close() {
+	h.mu.Lock()
+	if h.closed {
+		h.mu.Unlock()
+		return
+	}
+	h.closed = true
+
+	if h.stopCleanup != nil {
+		select {
+		case <-h.stopCleanup:
+		default:
+			close(h.stopCleanup)
+		}
+	}
+
+	docRooms := make([]*DocRoom, 0, len(h.docRooms))
+	for _, r := range h.docRooms {
+		docRooms = append(docRooms, r)
+	}
+
+	allWsClients := make([]*Client, 0, len(h.allClients))
+	for c := range h.allClients {
+		allWsClients = append(allWsClients, c)
+	}
+
+	allSSEClients := make([]*SSEClient, 0, len(h.sseClients))
+	for c := range h.sseClients {
+		allSSEClients = append(allSSEClients, c)
+	}
+	h.mu.Unlock()
+
+	// 1. Flush any active dirty document rooms and close doc WS connections
+	for _, r := range docRooms {
+		r.Flush()
+		r.mu.Lock()
+		if r.debounceTimer != nil {
+			r.debounceTimer.Stop()
+			r.debounceTimer = nil
+		}
+		for c := range r.wsClients {
+			_ = c.conn.Close()
+		}
+		r.wsClients = make(map[*DocWSClient]bool)
+		r.mu.Unlock()
+	}
+
+	// 2. Close all WebSocket connections
+	for _, c := range allWsClients {
+		_ = c.conn.Close()
+		c.closeSend()
+	}
+
+	// 3. Close all SSE client channels so ServeSSE select loop exits immediately
+	for _, c := range allSSEClients {
+		c.closeSend()
+	}
+}
 
 // Broadcasting Methods
 
@@ -446,19 +561,17 @@ func (h *Hub) BroadcastToAll(event string, payload interface{}) {
 	}
 
 	h.mu.RLock()
+	if h.closed {
+		h.mu.RUnlock()
+		return
+	}
 	defer h.mu.RUnlock()
 
 	for client := range h.allClients {
-		select {
+		client.trySend(data)
-		case client.send <- data:
-		default:
-		}
 	}
 	for sseClient := range h.sseClients {
-		select {
+		sseClient.trySend(data)
-		case sseClient.send <- data:
-		default:
-		}
 	}
 }
 
@@ -474,22 +587,20 @@ func (h *Hub) BroadcastToUser(userID, event string, payload interface{}) {
 	}
 
 	h.mu.RLock()
+	if h.closed {
+		h.mu.RUnlock()
+		return
+	}
 	defer h.mu.RUnlock()
 
 	if room, ok := h.userRooms[userID]; ok {
 		for client := range room {
-			select {
+			client.trySend(data)
-			case client.send <- data:
-			default:
-			}
 		}
 	}
 	if sseRoom, ok := h.sseUserRooms[userID]; ok {
 		for sseClient := range sseRoom {
-			select {
+			sseClient.trySend(data)
-			case sseClient.send <- data:
-			default:
-			}
 		}
 	}
 }
@@ -519,6 +630,10 @@ func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) {
 	}
 
 	h.mu.RLock()
+	if h.closed {
+		h.mu.RUnlock()
+		return
+	}
 	defer h.mu.RUnlock()
 
 	targetRoom := listID
@@ -528,19 +643,12 @@ func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) {
 
 	if room, ok := h.rooms[targetRoom]; ok {
 		for client := range room {
-			select {
+			client.trySend(data)
-			case client.send <- data:
-			default:
-				// Slow client buffer full
-			}
 		}
 	}
 	if sseRoom, ok := h.sseRooms[targetRoom]; ok {
 		for sseClient := range sseRoom {
-			select {
+			sseClient.trySend(data)
-			case sseClient.send <- data:
-			default:
-			}
 		}
 	}
 }
@@ -588,6 +696,10 @@ func (h *Hub) BroadcastRoom(roomID, event string, payload interface{}, excludeCl
 	}
 
 	h.mu.RLock()
+	if h.closed {
+		h.mu.RUnlock()
+		return
+	}
 	defer h.mu.RUnlock()
 
 	if room, ok := h.rooms[roomID]; ok {
@@ -595,18 +707,12 @@ func (h *Hub) BroadcastRoom(roomID, event string, payload interface{}, excludeCl
 			if excludeClient != nil && client == excludeClient {
 				continue
 			}
-			select {
+			client.trySend(data)
-			case client.send <- data:
-			default:
-			}
 		}
 	}
 	if sseRoom, ok := h.sseRooms[roomID]; ok {
 		for sseClient := range sseRoom {
-			select {
+			sseClient.trySend(data)
-			case sseClient.send <- data:
-			default:
-			}
 		}
 	}
 }
@@ -619,10 +725,7 @@ func (c *Client) SendEvent(event string, payload interface{}) {
 	}
 	data, err := json.Marshal(msg)
 	if err == nil {
-		select {
+		c.trySend(data)
-		case c.send <- data:
-		default:
-		}
 	}
 }
 

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