fix: avoid holding previous graphMuxes when reloading

[endigma]

Fixes a nasty GC bug:

Context:

• hot reload calls newGraphServer (new caches, parsed schema, executor, etc)
• SwapGraphServer installs it then calls Shutdown() on the old one
• Old server should now be GCed, BUT:

The problem was in buildMultiGraphHandler:

func (s *graphServer) buildMultiGraphHandler(opts...) {
	...
	return func(w http.ResponseWriter, r *http.Request) {
		// Extract the feature flag and run the corresponding mux
		// 1. From the request header
		// 2. From the cookie

		ff := strings.TrimSpace(r.Header.Get(featureFlagHeader))
		if ff == "" {
			cookie, err := r.Cookie(featureFlagCookie)
			if err == nil && cookie != nil {
				ff = strings.TrimSpace(cookie.Value)
			}
		}

		if mux, ok := featureFlagToMux[ff]; ok {
			w.Header().Set(featureFlagHeader, ff)
			mux.ServeHTTP(w, r)
			return
		}

		opts.baseMux.ServeHTTP(w, r) <- BIG PROBLEM
	}, reused, nil
}

This holds a reference to opts inside this closure since #2445, at the time of merge buildMultiGraphHandlerOptions struct didn't have much in it so this went undetected.

THEN: #2847 added a lot of big things to this type, including one currentGraphMuxes map[string]*graphMux which contains all of the stuff created by newGraphServer as detailed above.

Each reload now captures opts and by extension currentGraphMuxes, which grows with each reload because it folds in on itself via buildGraphMux mounting middleware that captures its own graphServer.

The fix is 2 lines:

baseMux := opts.baseMux // capture just *chi.Mux the handler needs
return func(w http.ResponseWriter, r *http.Request) {
    ...
    baseMux.ServeHTTP(w, r)
}

Summary by CodeRabbit

• New Features
  • Added additional demo subgraph services and updated demo routing and schema feature-flag entries.
  • Enabled config file watching in the debug configuration.
  • Added a hot-reload profiling workflow to report heap growth across reloads.
• Bug Fixes
  • Improved handler and pub/sub lifecycle handling during shutdown and reuse.
  • Improved runtime request accounting and deterministic feature-flag logging.
• Tests
  • Added coverage for multi-graph handler stability and shutdown semantics.
• Chores
  • Updated demo configuration, Dockerfiles, and compose files (quoting/port adjustments).

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

This PR updates graph-server request handling and pubsub shutdown ownership, adjusts router context internals and a reload comment, changes router/demo configuration and demo service wiring, and adds a hot-reload profiling script.

Changes

Graph server lifecycle and request handling

Layer / File(s)	Summary
Handler capture and retention tests router/core/graph_server.go, router/core/graph_server_test.go	buildMultiGraphHandler captures baseMux locally and the fallback uses that local mux; tests verify stale mux collection and reused mux shutdown behavior, with toKeys using cmp.Ordered.
Mux pubsub ownership and cache setup router/core/graph_server.go	graphMux stores pubsub providers, starts them from the provider slice returned by ecb.Build, shuts them down in graphMux.Shutdown, and drops explicit ristretto.NewCache type parameters.
Server counters and logging router/core/graph_server.go	graphServer.inFlightRequests changes to atomic.Int64, deferred completion uses Add(-1), and enabled feature flags are logged in sorted order.

Context and reload comment adjustments

Layer / File(s)	Summary
Context and reload comment edits router/core/context.go, router/core/reload_persistent_state.go	SubgraphHeadersBuilder and buildRequestContext bodies are removed from the shown diff, VariablesView() uses receiver o, and the reload comment names OnRouterConfigReload directly.

Router config, demo services, and Docker updates

Layer / File(s)	Summary
Schema and router config changes router/__schemas/graph.yaml, router/demo.config.yaml, router/debug.config.yaml	graph.yaml adds feature flags and new subgraph routing, demo.config.yaml changes scalar quoting, and debug.config.yaml enables execution-config file watching.
Demo Dockerfiles and compose services demo/docker/*.Dockerfile, docker-compose.full.yml	The demo Dockerfiles standardize stage alias casing, the countries and test1 images change exposed ports, and docker-compose.full.yml adds matching services plus quoting updates.

Hot-reload profiling script

Layer / File(s)	Summary
Reload profiling workflow mise-tasks/profile-reloads	The script handles path resolution, reload monitoring, router startup and cleanup, heap snapshot collection, and final diff generation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• wundergraph/cosmo#2823: Both PRs modify graph-mux lifecycle and shutdown behavior in router/core/graph_server.go.
• wundergraph/cosmo#2838: Both PRs change shutdown ownership for reused mux-related resources in the same router core path.
• wundergraph/cosmo#2844: Both PRs touch mux reuse, shutdown semantics, and related tests in router/core/graph_server.go.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main fix: preventing old graphMux instances from being retained during reloads.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

Router-nonroot image scan passed

✅ No security vulnerabilities found in image:

ghcr.io/wundergraph/cosmo/router:sha-c4659f72d44e0946148cfec133aa8c912bdc3be1-nonroot

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 61.38%. Comparing base (7715d70) to head (7d0d0f1).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3035      +/-   ##
==========================================
- Coverage   61.39%   61.38%   -0.01%     
==========================================
  Files         259      259              
  Lines       30053    30054       +1     
==========================================
- Hits        18451    18449       -2     
- Misses      10116    10119       +3     
  Partials     1486     1486              

Files with missing lines	Coverage Δ
router/core/context.go	74.85% <100.00%> (ø)
router/core/graph_server.go	85.26% <100.00%> (+0.01%)	⬆️
router/core/reload_persistent_state.go	93.42% <ø> (ø)

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

router/core/graph_server.go (1)

2225-2249: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Use a broadcast timeout instead of one shared timer channel.

All goroutines select on the same timer.C; when it fires, only one goroutine receives the tick. If another provider action hangs, providersGroup.Wait() can still block indefinitely.

Proposed fix

 func (s *graphServer) providersActionWithTimeout(ctx context.Context, providers []datasource.Provider, action func(ctx context.Context, provider datasource.Provider) error, timeout time.Duration, timeoutMessage string) error {
-	cancellableCtx, cancel := context.WithCancel(ctx)
+	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
-	timer := time.NewTimer(timeout)
-	defer timer.Stop()
-
 	providersGroup := new(errgroup.Group)
 	for _, provider := range providers {
 		providersGroup.Go(func() error {
 			actionDone := make(chan error, 1)
 			go func() {
-				actionDone <- action(cancellableCtx, provider)
+				actionDone <- action(timeoutCtx, provider)
 			}()
 			select {
 			case err := <-actionDone:
 				return err
-			case <-timer.C:
+			case <-timeoutCtx.Done():
+				if errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) {
+					return errors.New(timeoutMessage)
+				}
+				return timeoutCtx.Err()
-				return errors.New(timeoutMessage)
 			}
 		})
 	}

🤖 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/core/graph_server.go` around lines 2225 - 2249, The timeout handling
in providersActionWithTimeout is using one shared timer channel for all provider
goroutines, so only one can observe the timeout and the rest may hang. Update
the function to use a broadcast-style timeout signal for each provider action,
such as deriving a per-goroutine context from cancellableCtx or a shared timeout
context that all selects can observe, and make sure each action goroutine can
exit when the timeout fires. Keep the existing providersGroup and actionDone
flow, but replace the single timer.C select dependency with a cancellation
mechanism that all provider workers can detect.

🧹 Nitpick comments (3)

demo/graph-no-edg.yaml (1)

10-30: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Duplicate feature flag definitions.

myff2, myff3, and myff4 are byte-for-byte copies of myff (same subgraph, routing URL, schema). If the intent is just to have multiple named flags for reload/profiling stress testing, this is fine as-is; otherwise consider differentiating them or extracting a YAML anchor to avoid drift.

🤖 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 `@demo/graph-no-edg.yaml` around lines 10 - 30, The feature flag entries in the
YAML are duplicated across myff2, myff3, and myff4, so update the demo config to
either make each feature_graphs entry intentionally distinct or consolidate the
repeated products_fg definition using a YAML anchor/alias. Use the myff2, myff3,
and myff4 blocks as the target locations for the cleanup.

mise-tasks/profile-reloads (2)

88-88: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Interactive-inspect hint can reference a non-existent binary path.

$bin is only produced when --run-router builds it (line 117); otherwise bin="$outdir/router" (line 30) never gets created. The final echoed suggestion at line 185 still includes $bin in the example command even when it doesn't exist, which could confuse users running the script in attach-only mode.

Also applies to: 183-185

🤖 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 `@mise-tasks/profile-reloads` at line 88, The interactive-inspect hint is
building example args from $bin even when that binary is never created in
attach-only mode. Update the hint generation around the args assembly and final
echo in the profile-reloads script so it only references a path when it actually
exists, using the same $bin / $outdir-router logic already used by the
--run-router flow. Make the suggestion conditional on the binary being present
so the printed example command cannot point to a nonexistent file.

---

84-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

final is leaked as an implicit global instead of being returned explicitly.

report_diff() assigns final via printf -v final ... (line 85) without local, while diff/args are correctly scoped as local. The script then relies on this side-effect leak to reference $final at line 185, outside the function. This works today, but it's fragile — a future edit that adds local to line 85 (matching the pattern of the other variables) would silently make line 185 reference an unset/empty variable.

♻️ Proposed fix: compute and use `final` explicitly at the call site

 report_diff() {
-    printf -v final "%s/heap-%03d.pb.gz" "$outdir" "$reloads"
+    local final="$outdir/heap-$(printf '%03d' "$reloads").pb.gz"
     local diff="$outdir/diff-inuse_space.txt"
     local args=(-base "$baseline" -inuse_space -cum -top)
     [[ -f $bin ]] && args+=("$bin")
     go tool pprof "${args[@]}" "$final" >"$diff" 2>&1
     echo "baseline: $baseline"
     echo "final:    $final"
     echo "diff:     $diff"
     echo
     echo "Top inuse_space growth (baseline → final):"
     head -n 25 "$diff"
 }

 if [[ $profile == true ]]; then
     echo "Done — $reloads reloads profiled"
     report_diff
+    printf -v final "%s/heap-%03d.pb.gz" "$outdir" "$reloads"
     echo
     echo "Inspect interactively with:"
     echo "  go tool pprof -base $baseline -inuse_space -cum $bin $final"

Also applies to: 180-185

🤖 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 `@mise-tasks/profile-reloads` around lines 84 - 96, report_diff() is leaking
final as an implicit global, and later code is depending on that side effect
instead of using an explicit value. Update report_diff() to avoid setting final
implicitly with printf -v, and make the call site that uses $final compute or
capture the heap file path directly before invoking go tool pprof. Keep the
existing scoped handling for diff and args, and ensure any references to the
heap snapshot use the explicit path rather than relying on report_diff() state.

🤖 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/core/graph_server.go`:
- Around line 1503-1507: buildGraphMux currently starts pub/sub providers via
startupPubSubProviders and can still fail later, leaving partially initialized
providers running. Update the buildGraphMux error paths so that any failure
after providers start explicitly shuts them down before returning, using the
existing providers and shutdown flow around s.pubSubProviders. Also make sure
providersStarted is set to false immediately before the successful return gm,
nil so cleanup state stays correct.
- Around line 1506-1507: The pubsub shutdown bookkeeping in buildGraphMux only
runs for newly created muxes, so providers from reused muxes are missed. Update
the reuse path in commitReusedMuxes, or otherwise ensure reused mux providers
are appended to s.pubSubProviders before shutdown collection runs, so all
providers referenced by a mux are tracked consistently.

---

Outside diff comments:
In `@router/core/graph_server.go`:
- Around line 2225-2249: The timeout handling in providersActionWithTimeout is
using one shared timer channel for all provider goroutines, so only one can
observe the timeout and the rest may hang. Update the function to use a
broadcast-style timeout signal for each provider action, such as deriving a
per-goroutine context from cancellableCtx or a shared timeout context that all
selects can observe, and make sure each action goroutine can exit when the
timeout fires. Keep the existing providersGroup and actionDone flow, but replace
the single timer.C select dependency with a cancellation mechanism that all
provider workers can detect.

---

Nitpick comments:
In `@demo/graph-no-edg.yaml`:
- Around line 10-30: The feature flag entries in the YAML are duplicated across
myff2, myff3, and myff4, so update the demo config to either make each
feature_graphs entry intentionally distinct or consolidate the repeated
products_fg definition using a YAML anchor/alias. Use the myff2, myff3, and
myff4 blocks as the target locations for the cleanup.

In `@mise-tasks/profile-reloads`:
- Line 88: The interactive-inspect hint is building example args from $bin even
when that binary is never created in attach-only mode. Update the hint
generation around the args assembly and final echo in the profile-reloads script
so it only references a path when it actually exists, using the same $bin /
$outdir-router logic already used by the --run-router flow. Make the suggestion
conditional on the binary being present so the printed example command cannot
point to a nonexistent file.
- Around line 84-96: report_diff() is leaking final as an implicit global, and
later code is depending on that side effect instead of using an explicit value.
Update report_diff() to avoid setting final implicitly with printf -v, and make
the call site that uses $final compute or capture the heap file path directly
before invoking go tool pprof. Keep the existing scoped handling for diff and
args, and ensure any references to the heap snapshot use the explicit path
rather than relying on report_diff() state.

🪄 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: d6765bf9-8243-4e2f-9716-6db9104f87f9

📥 Commits

Reviewing files that changed from the base of the PR and between f14ae91 and b2dd666.

📒 Files selected for processing (5)

• demo/graph-no-edg.yaml
• mise-tasks/profile-reloads
• router/core/graph_server.go
• router/debug.config.yaml
• router/demo.config.yaml

💤 Files with no reviewable changes (1)

• router/demo.config.yaml

[coderabbitai]

🧹 Nitpick comments (1)

router/__schemas/graph.yaml (1)

2-22: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Duplicate feature flag definitions.

myff, myff2, myff3, and myff4 are byte-for-byte identical (same feature_graphs entry, same routing_url). If this is intentional demo/test fixture data for exercising multiple simultaneous flags, this is fine; otherwise consider whether distinct routing/subgraph targets were intended.

🤖 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/__schemas/graph.yaml` around lines 2 - 22, The feature flag entries in
graph.yaml are duplicated across myff, myff2, myff3, and myff4 with the same
feature_graphs and routing_url, so update the graph definitions to use distinct
targets if that was intended or consolidate/remove the duplicates if they are
not needed. Check the feature_flags list in the graph schema and ensure each
flag name maps to the correct unique subgraph_name and routing_url, or keep the
repeated entries only if this file is intentionally a fixture for testing
multiple identical flags.

🤖 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.

Nitpick comments:
In `@router/__schemas/graph.yaml`:
- Around line 2-22: The feature flag entries in graph.yaml are duplicated across
myff, myff2, myff3, and myff4 with the same feature_graphs and routing_url, so
update the graph definitions to use distinct targets if that was intended or
consolidate/remove the duplicates if they are not needed. Check the
feature_flags list in the graph schema and ensure each flag name maps to the
correct unique subgraph_name and routing_url, or keep the repeated entries only
if this file is intentionally a fixture for testing multiple identical flags.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: a1368db6-7947-4c8c-98df-5ba29cce37de

📥 Commits

Reviewing files that changed from the base of the PR and between b2dd666 and 0b0addd.

📒 Files selected for processing (15)

• demo/docker/availability.Dockerfile
• demo/docker/countries.Dockerfile
• demo/docker/employees.Dockerfile
• demo/docker/family.Dockerfile
• demo/docker/hobbies.Dockerfile
• demo/docker/mood.Dockerfile
• demo/docker/products.Dockerfile
• demo/docker/products_fg.Dockerfile
• demo/docker/test1.Dockerfile
• docker-compose.full.yml
• router/__schemas/graph.yaml
• router/core/context.go
• router/core/graph_server.go
• router/core/graph_server_test.go
• router/core/reload_persistent_state.go

✅ Files skipped from review due to trivial changes (1)

• router/core/reload_persistent_state.go

🚧 Files skipped from review as they are similar to previous changes (1)

• router/core/graph_server_test.go

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

router/core/graph_server.go (1)

2225-2248: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Broadcast the timeout to every provider wrapper

timer.C is single-consumer here: if multiple providers hang, only one goroutine can observe the timeout and the rest can stay blocked in select, leaving providersGroup.Wait() stuck. Use a shared timeout context (or a separate timer per provider) so every wrapper exits on the deadline.

🤖 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/core/graph_server.go` around lines 2225 - 2248, The
providersActionWithTimeout helper uses a single timer.C shared across all
provider goroutines, so only one wrapper may observe the timeout while others
can remain blocked and keep providersGroup.Wait() from returning. Update
providersActionWithTimeout to use a shared timeout context created from ctx (or
give each provider its own timer) and have each provider wrapper select on that
timeout signal alongside actionDone, referencing providersActionWithTimeout and
providersGroup.Go so every provider exits when the deadline is hit.

🤖 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.

Outside diff comments:
In `@router/core/graph_server.go`:
- Around line 2225-2248: The providersActionWithTimeout helper uses a single
timer.C shared across all provider goroutines, so only one wrapper may observe
the timeout while others can remain blocked and keep providersGroup.Wait() from
returning. Update providersActionWithTimeout to use a shared timeout context
created from ctx (or give each provider its own timer) and have each provider
wrapper select on that timeout signal alongside actionDone, referencing
providersActionWithTimeout and providersGroup.Go so every provider exits when
the deadline is hit.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 8f6ded79-27dd-4a33-a96f-b95679476c7d

📥 Commits

Reviewing files that changed from the base of the PR and between 0b0addd and dd49593.

📒 Files selected for processing (2)

• router/core/graph_server.go
• router/core/graph_server_test.go