Keep idle SSH tunnel sessions alive with a websocket ping

[anton-107]

Changes

An idle databricks ssh connect session dies after roughly nine minutes, with no warning and no actionable message — the user sees a raw server-side exception (websocket close 4000, Armeria ClosedStreamException) or simply a dropped connection, and loses whatever was in that shell. The tunnel is healthy the whole time; it dies purely because nothing was sent over it.

Both proxy loops (runSendingLoop / runReceivingLoop) are data-driven, so an idle SSH session puts no frames on the websocket at all and the server side reaps the stream it considers dead. Setting ServerAliveInterval 30 in the SSH client config works around it entirely (the reporting customer verified ~2 hours idle), because SSH-level keepalives are real payload bytes that the sending loop forwards — which is what pinned the diagnosis on the transport rather than on any CLI timer.

The client proxy now pings the websocket every 20 seconds for the life of the connection, as an additional goroutine in the errgroup that already drives the periodic handover tick:

• Client-only by construction. The server never calls RunClientProxy, so there is no flag or conditional that could enable server-side pinging, and the cluster-side server binary needs no redeployment — existing clusters benefit as soon as users update their CLI.
• Reuses the existing serialised write path. Pings go through sendMessage, which already holds the handover mutex, so the serialisation gorilla/websocket requires (it forbids concurrent writers) is inherited rather than newly built. No change to the proxy's write path, handover coordination, or connection swap.
• Handover-safe. A ping that ticks during a handover blocks on the mutex and goes out late. A handover establishes a fresh connection, so the peer's idle clock resets anyway and a subsumed ping is harmless. The handover waits on the receiving loop, not the sending loop, so a blocked ping cannot deadlock it — asserted by a test rather than assumed.
• A failed ping never ends a session. It is logged at debug level and the ticker continues. The receiving loop stays the sole authority on whether the connection is dead and notices within one read; an error returned from this goroutine would cancel the very session the keepalive exists to preserve.
• Keep-warm only. No read or write deadlines are set, and nothing tracks whether the peer answers. Full liveness detection — pong tracking against a read deadline, tearing down a non-responding peer — was deliberately rejected: it converts a missing keepalive into a new way to kill a healthy session.
• Each ping is logged at debug level, so support can confirm from a customer's log whether keepalives were flowing. The second commit adds this after end-to-end verification showed the pong handler alone logs nothing on this transport (see below).

The interval is an unexported 20s constant sitting beside the handover interval in experimental/ssh/cmd/constants.go, injectable at RunClientProxy in the same style as the handover tick (which is already parameterised for testing). There is no user-facing flag: the server's idle threshold is undocumented, so a knob nobody can tune intelligently is worse than a good default. 20s matches the vite bridge's keepalive and sits well under both bounds we have — the ~9 minute observed failure and the verified-sufficient 30s SSH-level keepalive.

Supporting context: the tunnel already rotates its websocket on a schedule via the periodic handover, so hostility to long-lived streams on this transport was already known and designed around here. The keepalive is the missing half of that same story rather than a new concern.

Alternative causes were ruled out and are recorded so nobody re-investigates them: the --shutdown-delay timer is stopped while a connection is registered (connections.go, TryAdd/Remove) so it cannot fire on a connected session; the periodic handover defaults to 30 minutes, not a value near nine; and the only other timer is the 30s initial handshake timeout, which applies to connection setup only. Nothing in the CLI fires near the observed nine minutes.

Rides along: the ssh server --shutdown-delay help text said the server shuts down "after no pings from clients", when no pings existed anywhere in the tunnel. It was inaccurate before this change and becomes actively misleading once real pings exist, so the one-line correction is included.

Note on precedent: libs/apps/vite/bridge.go runs the same 20s ping pattern, but not over the same transport — it connects to a dev-tunnel endpoint on the Databricks Apps domain, whereas the SSH tunnel connects to a driver-proxy path on the cluster. It is good evidence that a 20s ping is a working pattern in this repo against Databricks infrastructure; it is not evidence about the driver proxy's idle threshold or its reaping semantics.

Out of scope, tracked separately: surfacing a websocket close 4000 as an actionable message instead of a raw Armeria exception (plus a dedicated telemetry error category — it currently lands as unknown); any driver-proxy-side idle-timeout change; a configurable keepalive interval.

Resolves DECO-28186.

Why

Idle-session drops are a product defect that currently requires client-side configuration knowledge (ServerAliveInterval) to work around, for the tunnel's most ordinary use case: an open terminal nobody is typing into. A long-lived tunnel should heartbeat its own transport rather than depend on a server's tolerance, so this is the right fix regardless of what the driver proxy does.

Tests

Four unit tests, all asserting externally observable behaviour of the tunnel through the entry point users go through — never the internal shape of the keepalive goroutine:

• TestKeepalivePingReachesServer — a ping actually arrives at the peer on a session that sends no data.
• TestHandover is now run twice, once with keepalive active at a 1ms interval — pings interleaved with thousands of ordered messages across several handovers neither corrupt nor reorder the stream, and never trip gorilla's concurrent-write panic.
• TestKeepalivePingBlockedByHandoverDoesNotDeadlock — a ping blocked on the handover mutex does not deadlock an in-flight handover. This is the assertion that earns its keep: it covers the invariant the design was reasoned to rather than observed. The handover is held mid-flight (blocked on its dial) so the ping is provably waiting on the mutex, then both are required to complete.
• TestKeepalivePingFailureDoesNotEndSession — every ping write fails (a write deadline in the past, set before the connection reaches the proxy, leaves reads healthy) and the session must stay up.

An acceptance test that idles a websocket past a real timeout was deliberately not written: it would be slow and timing-flaky, and could still not prove the driver proxy stops reaping. The burden is split instead — unit tests prove the mechanism, a manual end-to-end run proves the outcome.

End-to-end verification (dogfood serverless)

Every run used an unpatched server binary built from origin/main and uploaded via --releases-dir (checked by inspection that it does not contain the change), so the runs differ in exactly one variable: whether the client pings. Idleness was held by a real SSH session running only date; sleep N; echo MARKER; date, with no SSH-level keepalive (nothing sets ServerAliveInterval, and the ssh-to-ProxyCommand link is a pipe, so TCPKeepAlive cannot apply either).

run	client	idle	result
baseline	unpatched	900s	reproduced the bug — marker due at 08:58:45 never arrived; still hung 48 min in, no error surfaced anywhere, all processes alive
21 min, under the handover	patched	1260s	survived, exit 0
45 min, spanning the handover	patched	2700s	survived, exit 0
forced handovers (--handover-timeout=30s)	patched	180s	survived, exit 0 — ~6 rotations interleaved with pings, no panic, no handover error

Two findings worth knowing when reviewing:

• No pong ever comes back. Zero pong lines across every patched run, with zero ping failures — control frames do not make the round trip here, and the outbound ping alone is what the reaper needs. The spec had bidirectional pinging as the fallback if client-only pings proved insufficient; it is not needed. This is also why each successful ping is now logged: otherwise a customer's log carries no positive evidence of keepalive activity. Verified after the change: ping lines appear at exact 20-second intervals.
• The observed failure was completely silent — no websocket close 4000 and no Armeria exception, on either the data path or the session's own handover tick. The separately-tracked follow-up about surfacing close 4000 as an actionable message would not have helped this case.

Full details and log excerpts are recorded on DECO-28186.

Full unit suite passes; the proxy package is green under -race. go test ./acceptance passes apart from three tests that fail for environment reasons on the machine used (the fips test needs a FIPS-built binary, and two terraform-backed tests hit the 60s script timeout).

This pull request and its description were written by Isaac.

[github-actions]

Waiting for approval

Based on git history, these people are best suited to review:

• @ilia-db -- recent work in experimental/ssh/cmd/, experimental/ssh/internal/client/, experimental/ssh/internal/proxy/

Eligible reviewers: @andrewnester, @denik, @janniklasrose, @lennartkats-db, @pietern, @rclarey, @renaudhartert-db, @rugpanov, @shreyas-goenka, @simonfaltum

_{Suggestions based on git history. See OWNERS for ownership rules.}

[eng-dev-ecosystem-bot]

Integration test report

Commit: bede575

Run: 32741885202

	Env	🔄​flaky	💚​RECOVERED	🙈​SKIP	✅​pass	🙈​skip	Time
💚​	aws linux		1	4	274	1167	9:59
🔄​	aws windows	1	1	4	275	1165	8:40
💚​	azure linux		1	4	273	1167	12:04
💚​	azure windows		1	4	275	1165	10:12
💚​	gcp linux		1	4	274	1167	10:54
🔄​	gcp windows	1	1	4	275	1165	10:07

7 interesting tests: 4 SKIP, 2 flaky, 1 RECOVERED

	Test Name	aws linux	aws windows	azure linux	azure windows	gcp linux	gcp windows
💚​	TestAccept	💚​R	💚​R	💚​R	💚​R	💚​R	💚​R
🙈​	TestAccept/bundle/invariant/no_drift	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/bundle/resources/vector_search_endpoints/drift/recreated_same_name	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/bundle/resources/vector_search_indexes/recreate/embedding_dimension	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/ssh/connection	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🔄​	TestSyncIncrementalSyncFileToPythonNotebook	✅​p	✅​p	✅​p	✅​p	✅​p	🔄​f
🔄​	TestSyncNestedSpacePlusAndHashAreEscapedSync	✅​p	🔄​f	✅​p	✅​p	✅​p	✅​p

Top 18 slowest tests (at least 2 minutes):

duration	env	testname
4:03	aws linux	TestFilerWorkspaceFilesExtensionsReadDir
3:58	gcp windows	TestAccept
3:50	azure windows	TestFilerRecursiveDelete/workspace_files_extensions
3:50	azure windows	TestFilerRecursiveDelete/workspace_files
3:37	gcp windows	TestImportDirDoesNotOverwrite
3:25	azure windows	TestAccept
3:25	azure linux	TestFilerWorkspaceFilesExtensionsReadDir
3:20	gcp windows	TestFilerWorkspaceFilesExtensionsReadDir
3:09	aws windows	TestAccept
2:56	gcp windows	TestFilerRecursiveDelete/workspace_files
2:39	azure linux	TestImportDirWithOverwriteFlag
2:34	azure windows	TestFilerWorkspaceFilesExtensionsStat
2:27	aws windows	TestFilerWorkspaceFilesExtensionsReadDir
2:23	gcp linux	TestFilerWorkspaceFilesExtensionsStat
2:20	gcp linux	TestFilerWorkspaceFilesExtensionsReadDir
2:08	azure windows	TestWorkspaceFilesExtensions_ExportFormatIsPreserved/source_python
2:07	gcp windows	TestImportDirWithOverwriteFlag
2:06	gcp linux	TestImportDirDoesNotOverwrite

[anton-107]

Should-fix: the keepalive ping can block shutdown/handover for up to ~15 min on a stalled connection

The keepalive ping is sent via sendMessage(...) → WriteMessage under handoverMutex with no write deadline. Because close() also routes through that same mutex (sendMessage(CloseMessage, …)) and there is no path that closes the underlying socket outside the mutex, a ping that parks in WriteMessage on a stalled/half-open TCP connection (full send buffer) holds the lock that both close() and handover need to acquire. Context cancellation cannot interrupt a goroutine parked in a blocking write, so shutdown/handover is stuck until the kernel TCP retransmission timeout fires (~15 min on default Linux tcp_retries2=15).

Severity: this is self-healing (the write eventually errors, the mutex releases, g.Wait() returns) — a multi-minute hang-on-exit, not a permanent zombie. Probability is low for a purely idle session (0-byte control frames into a near-empty buffer), but non-trivial in the realistic path: session pushes data → peer vanishes mid-transfer (buffer holds unacked data) → session goes idle → next ping parks in a full buffer. That idle-with-a-silently-dead-connection case is exactly the scenario this feature is meant to handle, so the ping converts a latent corner of the sending loop into a probable one in the feature's own domain.

Suggested fix (small, and arguably the cleaner design): send pings via conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(deadline)) on the loaded conn, without taking handoverMutex. gorilla explicitly permits WriteControl concurrently with WriteMessage, it takes a deadline, and it's already used for pongs in the keepalive test's fake server. This removes both the unbounded park and the ping↔close/handover coupling entirely. Note this also dissolves the stated rationale for routing pings through sendMessage — the "single concurrent writer" rule applies to WriteMessage, not WriteControl. If you'd rather keep the mutex path, the minimum fix is a bounded SetWriteDeadline before the control write.

Missing test: the current suite proves "handover blocks ping" and "write fails fast" (deadline set in the past → immediate error), but not the ordering that actually deadlocks: a ping parked in a blocking write while close()/handover waits on the mutex. Add a case that parks the ping write and asserts close()/handover still completes.

Everything else looks good — this is the one item worth addressing before merge (or tracking as a fast follow-up if bounded shutdown latency isn't a hard requirement for databricks ssh connect).

---

This comment was generated with GitHub MCP.

[anton-107]

Good catch — fixed in bede575. I verified both halves of the claim against gorilla's source before acting, and the diagnosis holds: WriteMessage sets no deadline (c.writeDeadline is never set, so write() passes the zero value to SetWriteDeadline), close() goes through the same sendMessage → handoverMutex, and nothing closes the socket outside that mutex, so a parked ping does hold up the closing handshake until the kernel gives up.

Pings now use conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout)) on the loaded conn, with no handover mutex. proxyPingWriteTimeout is 5s, matching wsWriteTimeout in libs/apps/vite/bridge.go.

Two refinements to the framing, for the record:

It bounds the coupling rather than removing it. WriteControl still acquires gorilla's internal write lock c.mu, which WriteMessage also needs — so a parked control write can still delay close(). What changes is that both the lock wait and the socket write are bounded by the deadline (select on <-c.mu vs a timer, then c.conn.SetWriteDeadline(deadline)), so the worst case is ~5s instead of ~15 min. The handoverMutex coupling is gone entirely, so the client handover path is fully decoupled. Also worth noting the deadline doesn't leak into data writes: gorilla's data path calls SetWriteDeadline(c.writeDeadline) (zero) before every frame, and c.mu prevents interleaving.

Two of the existing tests were asserting the old path, so they were reworked rather than added to:

• TestKeepalivePingBlockedByHandoverDoesNotDeadlock drove sendMessage directly, so it kept passing while testing something production no longer does. It is now TestKeepalivePingDuringHandoverDoesNotDisruptIt and asserts the opposite, stronger property: a ping issued while a handover is held mid-dial completes without waiting for it, the handover still completes, and the tunnel still carries data on the connection the handover installed.
• TestKeepalivePingFailureDoesNotEndSession induced failure with a past write deadline on the conn, which WriteControl now overrides with its own — so pings were succeeding and the test passed for the wrong reason. Failure is now induced at the socket.

On the new parked-write test. Wrote it as you suggested and it exposed something worth knowing: asserting on the session's shutdown doesn't isolate the ping. g.Wait() also waits on the receiving loop, which unblocks only when the peer reacts to the close frame — and a peer that has silently gone away never does, so shutdown hangs with or without a keepalive. That's pre-existing on the data path and out of scope here (the spec explicitly rules out reworking the shutdown coordination). So TestKeepalivePingParkedInWriteDoesNotStallClose is scoped to the ping's own contribution: it parks a ping in the socket write via a net.Conn wrapper (injected through Dialer.NetDialContext, so gorilla's real locking and deadline handling stay in the test) and requires the closing handshake to complete within the ping's deadline. The wrapper parks only for as long as the write's deadline and fails at once when the caller set none, which keeps the pre-existing unbounded park out of the measurement. That limitation is documented on the test.

If the silently-dead-peer shutdown hang is worth chasing, it wants its own issue — it needs an out-of-band conn.Close() on the teardown path, which is exactly the write-path rework this change was scoped to avoid. Happy to file it.

Re-verified end to end on dogfood after the change: 10 pings at exact 20-second intervals across a 200-second idle session, zero failures, session intact. Proxy package green under -race.