diff --git a/cmd/server/main.go b/cmd/server/main.go index 9f034c9..6ef54bd 100644 --- a/cmd/server/main.go +++ b/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(), + 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) + } } diff --git a/internal/ws/hub.go b/internal/ws/hub.go index 899f967..3b6d737 100644 --- a/internal/ws/hub.go +++ b/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,22 +463,88 @@ func (h *Hub) cleanupLocksLoop() { ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() - for range ticker.C { - h.mu.Lock() - now := time.Now() - var expiredDocIDs []string - for docID, l := range h.locks { - if l.expiresAt.Before(now) { - expiredDocIDs = append(expiredDocIDs, docID) - delete(h.locks, docID) + for { + select { + case <-h.stopCleanup: + return + case <-ticker.C: + h.mu.Lock() + now := time.Now() + var expiredDocIDs []string + for docID, l := range h.locks { + if l.expiresAt.Before(now) { + expiredDocIDs = append(expiredDocIDs, docID) + delete(h.locks, docID) + } + } + h.mu.Unlock() + + for _, docID := range expiredDocIDs { + h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } } - h.mu.Unlock() + } +} - for _, docID := range expiredDocIDs { - h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) +// 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 { - case client.send <- data: - default: - } + client.trySend(data) } for sseClient := range h.sseClients { - select { - case sseClient.send <- data: - default: - } + sseClient.trySend(data) } } @@ -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 { - case client.send <- data: - default: - } + client.trySend(data) } } if sseRoom, ok := h.sseUserRooms[userID]; ok { for sseClient := range sseRoom { - select { - case sseClient.send <- data: - default: - } + sseClient.trySend(data) } } } @@ -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 { - case client.send <- data: - default: - // Slow client buffer full - } + client.trySend(data) } } if sseRoom, ok := h.sseRooms[targetRoom]; ok { for sseClient := range sseRoom { - select { - case sseClient.send <- data: - default: - } + sseClient.trySend(data) } } } @@ -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 { - case client.send <- data: - default: - } + client.trySend(data) } } if sseRoom, ok := h.sseRooms[roomID]; ok { for sseClient := range sseRoom { - select { - case sseClient.send <- data: - default: - } + sseClient.trySend(data) } } } @@ -619,10 +725,7 @@ func (c *Client) SendEvent(event string, payload interface{}) { } data, err := json.Marshal(msg) if err == nil { - select { - case c.send <- data: - default: - } + c.trySend(data) } } diff --git a/internal/ws/hub_test.go b/internal/ws/hub_test.go new file mode 100644 index 0000000..eb47855 --- /dev/null +++ b/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: + } + } +}