Skip to content

fix: drain tunnel gateway routes and sessions on shutdown - #5935

Merged
ThomasRooney merged 1 commit into
mainfrom
tunnel-gateway-graceful-drain
Sep 2, 2026
Merged

fix: drain tunnel gateway routes and sessions on shutdown#5935
ThomasRooney merged 1 commit into
mainfrom
tunnel-gateway-graceful-drain

Conversation

@ThomasRooney

@ThomasRooney ThomasRooney commented Sep 1, 2026

Copy link
Copy Markdown
Member

Summary

Introduces a single-writer routeReconciler: one per-pod goroutine owns every mutating
route-store call (publish, connection snapshots, unpublish, global revoke). Handlers mark a
tunnel dirty and nudge; the reconciler reads live registry state at write time, batch-refreshes
all live tunnels on one ticker (replacing the per-session refresher goroutines), gates
publishes on IsActive, and prunes tracking for cleanly-removed tunnels.

Because all writes are serialized through one goroutine, graceful shutdown becomes the
reconciler's final act: on SIGTERM the registry stops admitting agent connects (checked again
under the lock at registration), the reconciler performs one bounded batched cleanup of every
route this pod may own, the forward server drains in-flight requests, agent sessions close
concurrently, and main blocks until the sequence completes — all within the existing 25s
budget. Publish-after-cleanup races are impossible by construction rather than coordinated
away.

Tests are black-box at the gateway boundary (real WebSocket agent + yamux sessions against an
in-memory store): post-drain store emptiness with no writes after drain returns, 503 on
connect during drain, in-flight forwards completing through shutdown, TTL refresh cadence,
revocation across multiple owners, and a concurrent-churn convergence property under the race
detector.

Motivation

The gateway exited without unpublishing its routes, so every rollout left pod-IP routes in
Redis for up to the 30s TTL while the pod was already gone. gram-server dialed the dead IP and
each affected tunneled request hung ~30s before surfacing a 502, producing an error burst on
every deploy. Draining removes the stale-route window; agents re-home to surviving pods in
~0.5–2s using their existing reconnect backoff. Complements #5888 (fast dial timeout) and
#5930 (dead-route eviction on dial failure), which cover unplanned pod death.

closes AIM-166

@ThomasRooney
ThomasRooney requested a review from a team as a code owner September 1, 2026 14:39
@ThomasRooney ThomasRooney added bug Something isn't working go Pull requests that update go code labels Sep 1, 2026
@changeset-bot

changeset-bot Bot commented Sep 1, 2026

Copy link
Copy Markdown

⚠️ No Changeset found

Latest commit: c67226b

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

@cubic-dev-ai

cubic-dev-ai Bot commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Running ultrareview automatically — This drain-on-shutdown PR adds substantial concurrency (publisher waits, bounded cleanup, registry drain state) across gateway and registry; a subtle race could stall shutdown or leak stale routes.. I'll post findings when complete.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ultrareview completed in 28m 28s

All reported issues were addressed across 4 files

Reply with feedback, questions, or to request a fix.

Fix all with cubic | Re-trigger cubic

Comment thread tunnel/gateway/registry.go Outdated
Comment thread tunnel/gateway/gateway.go Outdated
Comment thread tunnel/gateway/registry.go Outdated
Comment thread tunnel/gateway/gateway.go Outdated

@bflad bflad left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving, the drain sequencing and the test coverage here are solid and this should eliminate the deploy 502 bursts in the common case. I left inline comments on findings worth discussing: three edge cases that can undermine the drain (a late publish landing after the sweep, one 25s budget shared by all shutdown phases, and /healthz staying green while draining) and three smaller ones (retained ownership never retried, cleanupSessionState budget growth plus drain redundancy, and the MarkConnected drain re-check being best-effort). None of these block merging from my side and they could all be followups.

Comment thread tunnel/gateway/gateway.go Outdated
Comment thread tunnel/cmd/tunnel-gateway/main.go Outdated
Comment thread tunnel/cmd/tunnel-gateway/main.go
Comment thread tunnel/gateway/gateway.go Outdated
Comment thread tunnel/gateway/gateway.go Outdated
Comment thread tunnel/gateway/gateway.go Outdated
@ThomasRooney
ThomasRooney force-pushed the tunnel-gateway-graceful-drain branch from 4af0381 to a7a106a Compare September 1, 2026 16:40

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 7 files (changes from recent commits).

Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.

Fix all with cubic | Re-trigger cubic

Comment thread tunnel/cmd/tunnel-gateway/main.go Outdated
Comment thread tunnel/gateway/registry.go Outdated
Comment thread tunnel/gateway/registry.go
Comment thread tunnel/e2e/tunnel_test.go
@ThomasRooney
ThomasRooney force-pushed the tunnel-gateway-graceful-drain branch from a7a106a to 2f176a4 Compare September 1, 2026 16:49

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 7 files (changes from recent commits).

Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Tip: Review your code locally with the cubic CLI to iterate faster.

Fix all with cubic | Re-trigger cubic

Comment thread tunnel/gateway/reconciler.go
Comment thread tunnel/gateway/reconciler.go Outdated
Comment thread tunnel/e2e/tunnel_test.go
Comment thread tunnel/gateway/reconciler.go Outdated
@ThomasRooney
ThomasRooney force-pushed the tunnel-gateway-graceful-drain branch from 2f176a4 to c943cd7 Compare September 1, 2026 17:13

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 7 files (changes from recent commits).

Tip: Review your code locally with the cubic CLI to iterate faster.

Fix all with cubic | Re-trigger cubic

Comment thread tunnel/gateway/gateway.go
Comment thread tunnel/gateway/gateway_test.go Outdated
@ThomasRooney
ThomasRooney force-pushed the tunnel-gateway-graceful-drain branch from c943cd7 to c67226b Compare September 1, 2026 17:30

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 7 files (changes from recent commits).

Tip: Review your code locally with the cubic CLI to iterate faster.

Fix all with cubic | Re-trigger cubic

Comment thread tunnel/gateway/reconciler.go
Comment thread tunnel/e2e/tunnel_test.go
@linear-code

linear-code Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

AIM-166

@ThomasRooney
ThomasRooney added this pull request to the merge queue Sep 2, 2026
Merged via the queue into main with commit 24334ec Sep 2, 2026
55 checks passed
@ThomasRooney
ThomasRooney deleted the tunnel-gateway-graceful-drain branch September 2, 2026 13:00
@github-actions github-actions Bot locked and limited conversation to collaborators Sep 2, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

bug Something isn't working go Pull requests that update go code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants