fix(router): close Kafka consumer client when subscription ends

[mwisner]

Summary by CodeRabbit

• Bug Fixes
  • Kafka subscriptions now shut down promptly when canceled or when polling ends.
  • Improved resource cleanup prevents consumer connections from remaining open.
  • Idle subscriptions are reliably unblocked during cancellation, helping avoid lingering background activity.
  • Subscription shutdown is handled safely to prevent duplicate cleanup and improve connection stability.
  • Polling and subscription metrics now stop consistently when a subscription is canceled.

Description

kafka.ProviderAdapter.Subscribe creates a dedicated kgo.Client per subscription and never closes it. When a subscription ends, the poller goroutine returns and the client is simply dropped — Close() is never called. Because a franz-go client owns background goroutines that keep it reachable, the client is never garbage collected, so every ended subscription permanently leaks a kgo.Client together with its broker connections and buffered fetches. Heap usage and goroutine count grow with the cumulative (not concurrent) number of subscriptions until the process is OOM-killed. The only consumer teardown in the adapter today is for the shared producer (writeClient), closed in Shutdown.

A second, compounding issue: topicPoller blocks in PollRecords on the adapter (application) context, not the per-subscription context. A subscription cancelled while its topic is idle therefore never unblocks the poller, so neither the goroutine nor its client is reclaimed even at the loop level.

How to reproduce

• Point a Kafka EDFS subscription at a topic and churn subscriptions (repeated connect/disconnect, or many short-lived subscribers).
• inuse_space heap and goroutine count climb monotonically with the total number of subscriptions ever created and never fall — independent of how many subscribers are currently active.
• goroutine/heap profiles show accumulating franz-go client stacks (broker read/write loops and PollRecords).

Fix

In the Subscribe poller goroutine:

• Close the consumer client on every exit path via a sync.Once-guarded defer (Client.Close is not safe to call twice).
• Register context.AfterFunc(ctx, closeClient) on the subscription context, so that when the subscription is cancelled an in-flight PollRecords returns IsClientClosed and the poller exits promptly — even on an otherwise idle topic.

As a side benefit, Shutdown now also reclaims consumer clients through the existing closeWg.Wait(). The change is minimal (no new imports, no public API change) and preserves the stateless-consumer model.

Checklist

• I have discussed my proposed changes in an issue and have received approval to proceed.
• I have followed the coding standards of the project.
• Tests or benchmarks have been added or updated.
• If applicable, I have provided instructions for reviewers for manually validating my changes.
• Documentation has been updated on https://github.com/wundergraph/docs-website.
• I have read the Contributors Guide.

Open Source AI Manifesto

This project follows the principles of the Open Source AI Manifesto. Please ensure your contribution aligns with its principles.

---

Notes for reviewers

Opening as a draft for discussion. This is a bugfix rather than a feature, so no prior issue was opened — happy to file one if you'd prefer. I'd also be glad to add a regression test asserting the consumer client is closed on subscription-context cancellation (e.g. that PollRecords returns IsClientClosed and the poller goroutine exits) before this is marked ready.

Prepared with assistance from Claude (per the Open Source AI Manifesto's transparency principle).

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c5158468-fc9c-4de2-9a77-7332b8acef5a

📥 Commits

Reviewing files that changed from the base of the PR and between 4c8226d and fd6fdb2.

📒 Files selected for processing (1)

• router/pkg/pubsub/kafka/adapter.go

---

Walkthrough

Kafka subscription pollers now use the subscription context for polling and metrics. The adapter derives a polling context that combines subscription and adapter cancellation. Pollers run with WaitGroup.Go and defer consumer-client cleanup.

Changes

Kafka subscription context handling

Layer / File(s)	Summary
Poller context and cleanup router/pkg/pubsub/kafka/adapter.go	Polling and metrics use the subscription context. The adapter derives a polling context, cancels it when the adapter ends, runs the poller with WaitGroup.Go, and defers consumer-client cleanup.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the primary change: closing the Kafka consumer client when the subscription ends.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

🧹 Nitpick comments (1)

router/pkg/pubsub/kafka/adapter.go (1)

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

Use closeWg.Go for this goroutine.

The router module targets Go 1.25.0, so replace p.closeWg.Add(1), the goroutine, and deferred p.closeWg.Done() with p.closeWg.Go to keep lifecycle accounting in one API call.

Proposed refactor

-	p.closeWg.Add(1)
-
-	go func() {
-		defer p.closeWg.Done()
+	p.closeWg.Go(func() {
 		// existing poller cleanup and polling logic
-	}()
+	})

🤖 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/pkg/pubsub/kafka/adapter.go` around lines 169 - 184, Update the
goroutine lifecycle management in the surrounding poller code to use
p.closeWg.Go instead of manually calling p.closeWg.Add(1), launching a
goroutine, and deferring p.closeWg.Done(). Preserve the existing goroutine body
and ensure the closeWg.Go callback contains the current cleanup logic.

Source: Learnings

🤖 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/pkg/pubsub/kafka/adapter.go`:
- Around line 169-184: Update the goroutine lifecycle management in the
surrounding poller code to use p.closeWg.Go instead of manually calling
p.closeWg.Add(1), launching a goroutine, and deferring p.closeWg.Done().
Preserve the existing goroutine body and ensure the closeWg.Go callback contains
the current cleanup logic.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: da2bb928-0e66-4949-937a-cec54aca0353

📥 Commits

Reviewing files that changed from the base of the PR and between e6c2b07 and 55b18d6.

📒 Files selected for processing (1)

• router/pkg/pubsub/kafka/adapter.go

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 54.58%. Comparing base (7a68887) to head (fd6fdb2).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #3133       +/-   ##
===========================================
+ Coverage   48.94%   54.58%    +5.63%     
===========================================
  Files        1130      247      -883     
  Lines      157571    30493   -127078     
  Branches    10883        0    -10883     
===========================================
- Hits        77130    16644    -60486     
+ Misses      78588    12242    -66346     
+ Partials     1853     1607      -246     

Files with missing lines	Coverage Δ
router/pkg/pubsub/kafka/adapter.go	65.67% <100.00%> (-0.83%)	⬇️

... and 1006 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.

[dkorittki]

Hey @mwisner, thanks for the finding! I had a look at this and was able to reproduce the issue. Are you open to a few adjustments in this PR? Your fix works but ideally I would like to solve it differently: Make the topic poller return immidiately by using their contexts correctly and close the kafka client alongside it.

Currently the topic poller mixes up using p.ctx and ctx and its blocking for longer than it should be. You could create a new context, which cancels when either ctx or p.ctx is cancelled. You could do this right after the go routine is started:

p.closeWg.Go(func() {
	defer client.Close()

	pollerCtx, cancel := context.WithCancel(ctx) // cancels pollerCtx when ctx is cancelled
	defer cancel()
	stop := context.AfterFunc(p.ctx, cancel) // cancels pollerCtx when p.ctx is cancelled, too
	defer stop()
	
	err := p.topicPoller(pollerCtx, ...)
	// [...]
}

Inside p.topicPoller get rid of all direct p.ctx references like case <-p.ctx.Done() and client.PollRecords(p.ctx, 10_000), etc. The only context being used in there is via its ctx argument.

This will lead to the poller goroutine exiting immidiately when either a trigger closes or when the router shuts down or hot reloads. When that happens the defer client.Close() ensures the kafka client is closed as well.

If you want I can also close this PR, cherry-pick your commits here into a new PR with these changes. Just let me know.

[mwisner]

Thanks for reproducing it and for the guidance — that approach is cleaner, applied in the latest push:

• topicPoller no longer references p.ctx at all: PollRecords and the streamMetricStore.Consume call both use the passed ctx.
• In Subscribe, the goroutine derives pollerCtx, cancel := context.WithCancel(ctx) and adds context.AfterFunc(p.ctx, cancel), so the poller context is cancelled when either the subscription context or the application context is cancelled.
• A single defer client.Close() reclaims the client once the poller returns, so the sync.Once guard is gone.

Result: the poller returns immediately on a trigger close, router shutdown, or hot reload, and the client is closed alongside it. go build/vet/test ./pkg/pubsub/kafka/ are green.

Happy to keep it in this PR (no need to re-open a new one), and I can add a regression test asserting the poller exits and the client is closed on subscription-context cancellation if you'd like it here.