fix(router): keep websockets + resolver alive for reused graph muxes - #2838
Conversation
Each graphMux now owns its context derived from the router context rather than the graph server context. The executor (resolver event loop), websocket middleware, subscription clients, and cache warmup closures all use the per-mux context. Graph-server-scoped resources (connector, pubsub providers) continue to use the graph server context, because they are graph server resources. On graph server shutdown, reused muxes are skipped entirely. Their context stays alive so existing websocket connections and active subscriptions are not interrupted. Recreated muxes have their context cancelled in graphMux.Shutdown(), which terminates the resolver, websocket handler, and upstream subscription connections.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThis PR refactors the graph server's context management from using shared borrowed contexts to a hierarchical model with server-scoped and per-mux contexts. The refactoring enables mux-level context cancellation during hot-reload. A new integration test validates that updating a single feature-flag mux correctly closes only its connections while preserving others. ChangesContext Hierarchy Refactoring and Hot-Reload Test
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
Router image scan passed✅ No security vulnerabilities found in image: |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## feat/split-configs #2838 +/- ##
=====================================================
Coverage ? 65.80%
=====================================================
Files ? 256
Lines ? 26859
Branches ? 0
=====================================================
Hits ? 17674
Misses ? 7757
Partials ? 1428
🚀 New features to boost your workflow:
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
dkorittki
changed the title
dkorittki
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@router-tests/protocol/config_hot_reload_test.go`:
- Around line 625-630: Replace the direct websocket write call
conn.WriteJSON(...) used for the subscription handshake with the test helper
testenv.WSWriteJSON so the test benefits from the built-in retry/timeout
behavior; locate the call that constructs a testenv.WebSocketMessage (ID "1",
Type "subscribe", Payload ...) and change the invocation to
testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{...}) while keeping the
same message fields and error assertions intact (remove the manual
require.NoError on conn.WriteJSON since WSWriteJSON already handles errors).
In `@router/core/graph_server.go`:
- Around line 185-189: The constructors currently create contexts and partial
resources (graphMuxCtx/graphMuxCancel in buildGraphMux and
graphServerCtx/graphServerCancel in newGraphServer) but return on error without
cleaning them up; add deferred rollback logic in both functions that, on non-nil
error return, cancels the corresponding context(s) and invokes any partial
resource cleanup (e.g., call mux.Shutdown()/ShutdownMux, close caches, metric
stores, pub/sub providers, connectors) so no mux-scoped work lingers; implement
the defer immediately after creating graphMuxCtx/graphMuxCancel and
graphServerCtx/graphServerCancel (and after any resource that needs explicit
Close) that checks if err != nil and then cancels and closes those resources
before returning.
- Around line 877-879: The Shutdown method on graphMux currently defers
s.cancel(), which delays cancelling the mux context until after
cache/metric/exporter shutdown; change it to call s.cancel() immediately at the
start of graphMux.Shutdown (i.e., invoke s.cancel() before shutting down caches,
metrics, exporters and websocket/resolver goroutines) so that subscriptions and
goroutines tied to the mux context are stopped as soon as shutdown begins.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: bcbbfb2b-66bc-40fb-92dc-1f2deeb05fb1
📒 Files selected for processing (2)
router-tests/protocol/config_hot_reload_test.gorouter/core/graph_server.go
| err := conn.WriteJSON(&testenv.WebSocketMessage{ | ||
| ID: "1", | ||
| Type: "subscribe", | ||
| Payload: []byte(`{"query":"subscription { currentTime { unixTime timeStamp }}"}`), | ||
| }) | ||
| require.NoError(t, err) |
There was a problem hiding this comment.
Use testenv.WSWriteJSON for the subscribe handshake.
This is still the success path, so the new hot-reload test should use the websocket helper rather than conn.WriteJSON directly. That keeps the test aligned with the repo’s retry/deadline behavior and reduces flakes.
As per coding guidelines: "Use testenv.WSReadJSON and testenv.WSWriteJSON for WebSocket reads and writes in tests instead of conn.ReadJSON and conn.WriteJSON, as these helpers include retry logic with 2-second deadlines and exponential backoff."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router-tests/protocol/config_hot_reload_test.go` around lines 625 - 630,
Replace the direct websocket write call conn.WriteJSON(...) used for the
subscription handshake with the test helper testenv.WSWriteJSON so the test
benefits from the built-in retry/timeout behavior; locate the call that
constructs a testenv.WebSocketMessage (ID "1", Type "subscribe", Payload ...)
and change the invocation to testenv.WSWriteJSON(t, conn,
&testenv.WebSocketMessage{...}) while keeping the same message fields and error
assertions intact (remove the manual require.NoError on conn.WriteJSON since
WSWriteJSON already handles errors).
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.
Tip: disable this comment in your organization's Code Review settings.
| }) | ||
| } | ||
|
|
||
| // TestConfigHotReloadReusedFeatureFlagMuxKeepsWebSocketsAlive verifies that after a |
2 participants
Since we are now able to reuse feature flag and base graph muxes if their config parts haven't changed there is no reason to kill their websocket connections during a graph server swap. Unfortunately their contexts, which are used to actually close them, is tied to the graph server. So whenever the graph server context is canceled during graph server swap, all websocket connections across all muxes get closed regardless of wether their mux does not get cancelled.
The fix is to inherit the graph mux context from the router instead of the graph server. What that does is that websocket connections only get closed on graph mux shutdown, not on graph server shutdown. When a graph server shuts down it shuts down only those muxes, which actually need to be closed. This in turn means that only those websocket connections get closed, which need to get closed.
One important note: This pull request should not have an impact on customers not using the split config feature in the token. But it does have changes on their code path, since I changed the context inheritance. Just something to keep in mind during review.
Summary by CodeRabbit
Checklist
Open Source AI Manifesto
This project follows the principles of the Open Source AI Manifesto. Please ensure your contribution aligns with its principles.