fix: prevent goroutine leaks in SSE implementation #152
Conversation
WalkthroughThis change refactors the server's SSE (Server-Sent Events) connection handler to use explicit context cancellation and goroutine synchronization. It introduces a cancellable context for each SSE session, ensuring that all spawned goroutines are properly signaled to exit when the session ends. A sync.WaitGroup is used to track and wait for goroutine completion. The handler now checks for write errors to the client and returns early if the client disconnects. In message handling, event queuing is protected by a context with a timeout to prevent goroutine leaks. Formatting changes to struct declarations are also included. Changes
Possibly related PRs
Tip ⚡💬 Agentic Chat (Pro Plan, General Availability)
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
server/sse.go (2)
236-238: Redundantcancel()call could be clarified.
You callcancel()here and also within thecleanup()function (line 321). Although canceling a context more than once is safe, you might consider removing the second call for clarity.func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) { ... - ctx, cancel := context.WithCancel(r.Context()) - defer cancel() + ctx, outerCancel := context.WithCancel(r.Context()) + // Defer outerCancel only once if desired, or rely purely on cleanup(). ... }
288-312: Keep-alive goroutine likewise uses WaitGroup and handles cancellation gracefully.
One suggestion: sending pings intosession.eventQueuecan block if the channel reaches capacity, potentially causing a deadlock. Consider using a timeout or selective send, similar to the message handler, to avoid blocking the keep-alive loop if the queue is full.for { select { case <-ticker.C: - session.eventQueue <- fmt.Sprintf(":ping - %s\n\n", time.Now().Format(time.RFC3339)) + ctxPing, cancelPing := context.WithTimeout(ctx, 2*time.Second) + select { + case session.eventQueue <- fmt.Sprintf(":ping - %s\n\n", time.Now().Format(time.RFC3339)): + // Ping sent + case <-ctxPing.Done(): + // Avoid blocking indefinitely + case <-session.done: + cancelPing() + return + } + cancelPing() case <-session.done: return case <-ctx.Done(): return } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/sse.go(6 hunks)
🔇 Additional comments (8)
server/sse.go (8)
56-64: New fields in SSEServer struct look good for configuration and concurrency.
All added fields facilitate flexible configuration of the SSE server. Usingsync.Mapfor storing sessions is valid for concurrent read/writes. No immediate concerns here.
161-167: Initialization of default settings is appropriate.
Defining sensible defaults forsseEndpoint,messageEndpoint, andkeepAliveIntervalis a good practice. This helps ensure that the server functions with minimal user configuration.
252-255: Session registration error handling is correct.
Early return on registration failure prevents partially initialized sessions from persisting. No immediate concurrency concerns here.
258-259: Introduction of a WaitGroup for concurrency synchronization is beneficial.
This allows proper coordination of goroutines to avoid leaks. Great choice to avoid leftover goroutines when the connection ends.
261-285: Notification handler goroutine cancellation logic is well-structured.
The select statement cleanly handlessession.doneandctx.Done(). This ensures each branch properly terminates the goroutine without blocking.
318-325: Cleanup function efficiently terminates all goroutines.
Callingclose(session.done),cancel(), thenwg.Wait()ensures no goroutine remains. Just watch out for the duplicatecancel()call at line 237.
331-337: Immediate return on write failure prevents stuck writes.
Exiting onerr != nilis a good approach to handle client disconnections promptly. This prevents wasted effort sending data to a dead connection.
394-405: Timed context for event queue insertion prevents deadlocks.
Applying a 5-second timeout here is a solid strategy to avoid blocking the goroutine if the queue is full.
The SSE implementation had potential goroutine leaks when clients disconnect unexpectedly. This commit addresses these issues by: 1. Using context.WithCancel to properly control goroutine lifecycle 2. Adding a WaitGroup to ensure all goroutines complete before returning 3. Implementing a cleanup function that closes all resources properly 4. Adding error checking when writing to client connections 5. Adding timeout when sending to eventQueue channels This prevents resource exhaustion in long-running servers with many client connections.
- Move cancel() to cleanup function as final step - Add clear comments explaining cleanup order - Move cleanup defer immediately after context creation - Ensure context cancellation happens on all return paths - Remove redundant context cancellation
Napat
force-pushed
the
fix/sse-goroutine-leaks
branch
from
08ed317 to
78c4c7b
Compare
The SSE implementation had potential goroutine leaks when clients disconnect
unexpectedly. This commit addresses these issues by:
This prevents resource exhaustion in long-running servers with many
client connections.
Summary by CodeRabbit
Bug Fixes
Refactor