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.
• Bounded, and off the handover mutex. Pings go out with conn.WriteControl(websocket.PingMessage, nil, now+proxyPingWriteTimeout) (5s). The value is borrowed from vite's wsWriteTimeout, but note vite bounds a channel handoff there and sets no socket write deadline at all, so it is precedent for the number, not for the mechanism. gorilla permits WriteControl concurrently with the data writes, and its deadline bounds both the wait for the connection's write lock and the socket write itself. Routing pings through sendMessage instead — the first cut of this change — put an unbounded write on the shared write path: WriteMessage sets no deadline and close() takes the same handover mutex, so a ping parked on a half-open socket (full send buffer, no RST) held up the closing handshake until the kernel abandoned its retransmits, ~15 min with Linux defaults. Thanks @rugpanov for catching that. No change to the proxy's write path, handover coordination, or connection swap.
• Handover-safe. A ping that ticks during a handover goes to the connection being replaced and may simply fail, which is already non-fatal; the next tick uses the new connection. It never waits on the handover, and the handover never waits on it.
• A failed ping does not end the session, but it is not free. gorilla puts a connection into a permanent write-error state after any failed write, so one timed-out ping stops every later write — data frames and the close frame alike. The session is left to end the way it would without a keepalive: the next write fails and the sending loop reports it. It is logged at warn, since it means the tunnel can no longer send. Ending the session on the failed ping itself was considered and rejected: reads are unaffected, and a long job that prints nothing for minutes and then produces output is one of the cases this feature exists for — cutting that off would be worse than letting it finish.
• Teardown closes the connection. pc.close only sends a close message, which needs a live peer to react to it, and it cannot go out at all on a poisoned write path — so the receiving loop stayed blocked in ReadMessage and the session hung instead of exiting. The teardown goroutine's comment already said it closed the connection; now it does. This also fixes the same hang for a peer that silently goes away, keepalive or not.
• Keep-warm only. No read deadline, no pong tracking, no dead-peer detection. Full liveness detection — tracking pongs against a read deadline and 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. Added 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

Five keepalive tests plus one reused, 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.
• TestKeepalivePingDuringHandoverDoesNotDisruptIt — 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.
• TestKeepalivePingParkedInWriteDoesNotStallClose — a ping parked in the socket write must not hold up the closing handshake for longer than its deadline allows. The park is induced at the socket, through a net.Conn wrapper injected via Dialer.NetDialContext, so gorilla's real locking and deadline handling stay in the test.
• TestKeepalivePingFailureDoesNotEndSession — every ping write fails (induced at the socket, leaving reads healthy) and the session must stay up.
• TestKeepalivePingFailureDoesNotHangTheSession — a session whose write path a failed ping has poisoned must end promptly with an error, not hang.
• 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.

Every one of these was mutation-tested — production code broken deliberately, to check the suite notices. All five kill their mutants: neutering the ping, making a failed ping fatal, reverting sendPing to the unbounded sendMessage path, putting the ping back on the handover mutex, and removing the teardown close. Two of them did not, before an independent review caught it: the parked-write fake modelled a deadline-less write as an instant failure (the opposite of an unbounded park), and the failure test was vacuous because the fake's broken close handshake meant the session could not end for reasons unrelated to the keepalive. Nothing asserts logPongs; on a transport that returns no pongs there is nothing to assert.

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
re-verify after the WriteControl change	patched	200s	survived, exit 0 — 10 pings at exact 20s intervals, no failures
re-verify after the teardown change	patched	70s	survived, exit 0 — 3 pings, no warnings, no spurious disconnect message

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

The hang that surfaced while testing — a peer that silently goes away leaves g.Wait() waiting on the receiving loop forever — turned out to be reachable through the keepalive, so it is fixed here rather than deferred: one failed ping poisons the write path, which is enough to keep the close message from going out. That is the teardown close above.

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.

[eng-dev-ecosystem-bot]

Integration test report

Commit: f2086b3

Run: 32826292844

	Env	🔄​flaky	💚​RECOVERED	🙈​SKIP	✅​pass	🙈​skip	Time
🔄​	aws linux	2	1	4	272	1180	4:12
💚​	aws windows		1	4	276	1178	4:22
💚​	azure linux		1	4	273	1180	4:03
💚​	azure windows		1	4	275	1178	4:02
💚​	gcp linux		1	4	274	1180	4:07
💚​	gcp windows		1	4	276	1178	3:12

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
🔄​	TestFilerWorkspaceNotebook	🔄​f	✅​p	✅​p	✅​p	✅​p	✅​p
🔄​	TestFilerWorkspaceNotebook/sqlNb.sql	🔄​f	✅​p	✅​p	✅​p	✅​p	✅​p

Top 3 slowest tests (at least 2 minutes):

duration	env	testname
4:17	aws windows	TestAccept
3:58	azure windows	TestAccept
3:08	gcp windows	TestAccept

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

[rugpanov]

Nit (non-blocking): the PR description is stale relative to the shipped implementation and is worth refreshing before merge so reviewers reading top-down aren't misled.

It still describes the superseded mutex-based design — "pings go through `sendMessage`, which already holds the handover mutex" — and references tests that no longer exist (`TestKeepalivePingBlockedByHandoverDoesNotDeadlock`, "TestHandover is now run twice").

The shipped code (commit `bede5759e`) instead sends pings via `WriteControl` off the handover mutex with a 5s write deadline, and the actual tests are `TestKeepalivePing{ReachesServer,FailureDoesNotEndSession,ParkedInWriteDoesNotStallClose,DuringHandoverDoesNotDisruptIt}`.

The code and commit messages are consistent — only the PR description drifted.

[anton-107]

Fixed — description rewritten to match what shipped.

Corrected: the mutex-based mechanism bullet now describes WriteControl off the handover mutex with the 5s deadline (and records why the first cut was wrong, crediting your earlier catch); the handover-interaction bullet no longer claims a ping blocks and goes out late; "no read or write deadlines are set" is now "no read deadline, no pong tracking, no dead-peer detection", since the ping write does carry one; and the test list is the real one — TestKeepalivePing{ReachesServer,DuringHandoverDoesNotDisruptIt,ParkedInWriteDoesNotStallClose,FailureDoesNotEndSession}, with the failure-induction description fixed too (it's induced at the socket now, not with a past write deadline).

Also added the WriteControl re-verification run to the end-to-end table, and a note about the one pre-existing hazard the parked-write test surfaced but this change does not address.

You checked and said code and commit messages were consistent — nearly. One test comment had drifted the same way ("pings share the proxy's serialised write path", which is no longer where they go), fixed in 88dcba3. TestHandover's two subtests were real, so that part of the description stayed.

[anton-107]

Defect confirmed and fixed in f2086b34b, along with all three test-power findings. Thanks — the mutation testing in particular found things a passing suite could never have told me.

The defect. The diagnosis is exactly right: "a failed ping never ends a session" was true of the errgroup and false of gorilla's connection state. writeFatal sets a connection-level writeErr that every later write checks first, so one timed-out ping stops the data frames and the close frame, and the session then died through the data path — silently, by hanging.

Fixed by taking the second of the two directions: the teardown goroutine now closes the websocket, not just sends a close message. Worth noting it already claimed to — its comment reads "we close the connection and the source ... to unblock them" — so this implements documented intent that was never written. Both loops unblock and the session exits with the sending loop's error. That also fixes the same hang for a peer that silently goes away, keepalive or not, which I had flagged as wanting its own issue; it turned out to be reachable through the keepalive, so it belongs here.

Why the ping failure is still not fatal. I considered the first direction and rejected it, on a case the spec cares about: a poisoned write path leaves reads working, and "a long job that prints nothing for minutes, then produces output" is one of this feature's own user stories. Ending the session on a failed ping would cut off data still arriving fine. There is also a non-poisoning failure — WriteControl returns a timeout when it cannot take the write lock within the deadline, with no side effects — and it is indistinguishable from a poisoning failure without matching on error strings, which this repo forbids. So the session ends the way it would with no keepalive at all (next write fails, sending loop reports it), and the ping failure is now logged at warn rather than debug, since it means the tunnel can no longer send. If you think a doomed idle session should die immediately instead of at the user's next keystroke, the discriminator exists — a poisoned connection fails every ping forever, a lock-wait timeout does not — and I'll take it as a follow-up.

Test power. All five assertions now kill their mutants; I re-ran your method to check:

mutation	before	now
neuter the ping send	caught	caught
make a failed ping fatal	survived	caught
revert sendPing to sendMessage	survived (0.006s)	caught
ping back on the handover mutex	caught, but as a 120s package panic	caught, clean, 5.00s
remove the teardown close	—	caught (new test)

• The parked-write fake modelled a deadline-less write as an instant failure — the opposite of an unbounded park. It now parks such a write for longer than any assertion, exactly the fix you verified.
• TestKeepalivePingFailureDoesNotEndSession was vacuous for precisely the reason given. Closing the connection on teardown restores its power rather than needing a rewrite: the session can now end, so the fatal-ping mutant fails it. Its stated property also survives the defect fix, since a failed ping still doesn't end a session.
• The handover test's wedge took two attempts. Deferring the release before client.Cleanup didn't work — cleanup runs first and waits on proxy loops that can't finish until the release, so release() never ran. Deferred after cleanup so it runs before it.
• New TestKeepalivePingFailureDoesNotHangTheSession asserts the defect directly.

Minor notes. Vite correction taken — the description no longer cites vite for the write deadline, only for the 20s interval. The 2s unconditional wait is now 1s. Left as-is with reasons: the zero-interval ticker panic (a panic at session start is the right failure mode for a programming error, and no live path reaches it), and logPongs (nothing to assert on a transport that returns no pongs — now stated in the description rather than left as an apparent gap).

Re-verified end to end on dogfood, since the teardown change is on the path every session exit takes: normal exit still clean at exit 0, 70s true idle, 3 pings at 20s spacing, no warnings, no spurious disconnect message. Unit suite green, -race green over repeated runs, lint clean.

[rclarey]

We should fix this. I fixed the driver proxy <-> server side, but not the client <-> driver proxy side