From 1ec34c9a795b25a50a64f0d973e95437626d1d2a Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Mon, 24 Aug 2026 08:24:09 +0000 Subject: [PATCH 1/6] Keep idle SSH tunnel sessions alive with a websocket ping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An idle `databricks ssh connect` session dies after roughly nine minutes. Both proxy loops are purely data-driven, so a session nobody is typing into puts no frames on the websocket at all, and the server side reaps the stream it then considers dead (websocket close 4000, Armeria ClosedStreamException). Setting `ServerAliveInterval` in the SSH client config works around it entirely, because SSH-level keepalives are real payload bytes that the sending loop forwards — which is what pinned the diagnosis on the transport. 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. That placement makes the keepalive client-only by construction: the server never calls RunClientProxy, so no flag can enable server-side pinging, and the cluster-side binary needs no redeployment. Pings take the proxy's existing serialised write path (sendMessage), which already holds the handover mutex, so the serialisation gorilla/websocket requires is inherited rather than newly built. A ping that ticks during a handover blocks and goes out late; a handover establishes a fresh connection, so the peer's idle clock resets anyway. A failed ping is logged at debug level and never returned: the receiving loop stays the sole authority on whether the connection is dead, and an error here would cancel the session the keepalive exists to preserve. Liveness posture is keep-warm only — no read or write deadlines, and the pong handler only logs. 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 story. Also corrects the `ssh server --shutdown-delay` help text, which claimed the server shuts down "after no pings from clients" when no pings existed anywhere in the tunnel — inaccurate today, and actively misleading once real pings exist. Co-authored-by: Isaac --- .nextchanges/cli/ssh-tunnel-keepalive.md | 1 + experimental/ssh/cmd/connect.go | 1 + experimental/ssh/cmd/constants.go | 13 +- experimental/ssh/cmd/server.go | 2 +- experimental/ssh/internal/client/client.go | 4 +- experimental/ssh/internal/proxy/client.go | 48 ++++- .../ssh/internal/proxy/client_server_test.go | 39 ++-- .../ssh/internal/proxy/keepalive_test.go | 173 ++++++++++++++++++ experimental/ssh/internal/proxy/proxy_test.go | 14 ++ 9 files changed, 271 insertions(+), 24 deletions(-) create mode 100644 .nextchanges/cli/ssh-tunnel-keepalive.md create mode 100644 experimental/ssh/internal/proxy/keepalive_test.go diff --git a/.nextchanges/cli/ssh-tunnel-keepalive.md b/.nextchanges/cli/ssh-tunnel-keepalive.md new file mode 100644 index 00000000000..71802b97adb --- /dev/null +++ b/.nextchanges/cli/ssh-tunnel-keepalive.md @@ -0,0 +1 @@ +Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes. The tunnel now sends a websocket keepalive every 20 seconds, so a session nobody is typing into stays connected without setting `ServerAliveInterval` in the SSH client config. diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 2c50b871902..7502d70f560 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -119,6 +119,7 @@ Connect to a dedicated cluster: ShutdownDelay: shutdownDelay, MaxClients: maxClients, HandoverTimeout: handoverTimeout, + KeepaliveInterval: defaultKeepaliveInterval, ReleasesDir: releasesDir, ServerTimeout: max(serverTimeout, shutdownDelay), TaskStartupTimeout: startupTimeout, diff --git a/experimental/ssh/cmd/constants.go b/experimental/ssh/cmd/constants.go index 64c99b5bd48..dd5d3b2fdf2 100644 --- a/experimental/ssh/cmd/constants.go +++ b/experimental/ssh/cmd/constants.go @@ -3,10 +3,15 @@ package ssh import "time" const ( - defaultServerPort = 7772 - defaultMaxClients = 10 - defaultShutdownDelay = 10 * time.Minute - defaultHandoverTimeout = 30 * time.Minute + defaultServerPort = 7772 + defaultMaxClients = 10 + defaultShutdownDelay = 10 * time.Minute + defaultHandoverTimeout = 30 * time.Minute + // How often the client pings the tunnel websocket so an idle SSH session keeps the transport + // alive. Matches the keepalive interval of the vite bridge (libs/apps/vite/bridge.go), and sits + // well under both the ~9 minutes after which idle sessions were observed to drop and the + // 30 second SSH-level keepalive that was verified to prevent it. + defaultKeepaliveInterval = 20 * time.Second defaultEnvironmentVersion = 4 serverTimeout = 24 * time.Hour diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 21c651b2365..3675a4a7fe8 100644 --- a/experimental/ssh/cmd/server.go +++ b/experimental/ssh/cmd/server.go @@ -41,7 +41,7 @@ and proxies them to local SSH daemon processes.`, cmd.MarkFlagRequired("authorized-key-secret-name") cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients") - cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down after no pings from clients") + cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server when there are no active connections") cmd.Flags().StringVar(&version, "version", "", "Client version of the Databricks CLI") cmd.Flags().BoolVar(&serverless, "serverless", false, "Enable serverless mode for Jupyter initialization") cmd.Flags().StringVar(&usagePolicyID, "usage-policy-id", "", "Usage policy ID the job was submitted with") diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 90e99f6b590..b2a8c9d0b7f 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -86,6 +86,8 @@ type ClientOptions struct { ServerMetadata string // How often the CLI should reconnect to the server with new auth. HandoverTimeout time.Duration + // How often the CLI pings the tunnel websocket to keep an idle session alive. + KeepaliveInterval time.Duration // Max amount of time the server process is allowed to live ServerTimeout time.Duration // Max amount of time to wait for the SSH server task to reach RUNNING state @@ -895,7 +897,7 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server requestHandoverTick := func() <-chan time.Time { return time.After(opts.HandoverTimeout) } - return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, createConn) + return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, createConn) } // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI. diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index a1e8389e7ff..47c4cb09b6d 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -9,6 +9,7 @@ import ( "time" "github.com/databricks/cli/libs/log" + "github.com/gorilla/websocket" "golang.org/x/sync/errgroup" ) @@ -37,8 +38,25 @@ func (f *firstByteWriter) Write(p []byte) (int, error) { return f.w.Write(p) } -func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error { - proxy := newProxyConnection(createConn) +// logPongs wraps a connection factory so every connection it creates — the initial one and each +// one a handover creates — logs the pongs coming back for our keepalive pings. Debug visibility only: +// the receiving loop stays the only judge of whether a connection is alive. +func logPongs(ctx context.Context, createConn createWebsocketConnectionFunc) createWebsocketConnectionFunc { + return func(connCtx context.Context, connID string) (*websocket.Conn, error) { + conn, err := createConn(connCtx, connID) + if err != nil { + return nil, err + } + conn.SetPongHandler(func(string) error { + log.Debugf(ctx, "Received websocket keepalive pong") + return nil + }) + return conn, nil + } +} + +func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, createConn createWebsocketConnectionFunc) error { + proxy := newProxyConnection(logPongs(ctx, createConn)) log.Infof(ctx, "Establishing SSH proxy connection...") ctx, cancel := context.WithCancel(ctx) defer cancel() @@ -68,9 +86,33 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque } } }) + g.Go(func() error { + // Keep the websocket carrying traffic while the SSH session is idle. Both proxy loops + // are data-driven, so an idle session puts no frames on the connection at all and the + // server side reaps the stream it then considers dead (websocket close 4000). + ticker := time.NewTicker(keepaliveInterval) + defer ticker.Stop() + for { + select { + case <-gCtx.Done(): + return gCtx.Err() + case <-ticker.C: + // Pings take the same serialised write path as data — gorilla forbids + // concurrent writers — so a ping that ticks during a handover blocks until it + // finishes and then goes out late. Harmless: a handover establishes a fresh + // connection, which resets the peer's idle clock anyway. + if err := proxy.sendMessage(websocket.PingMessage, nil); err != nil { + // Never fatal. A failed ping knows nothing the data loops don't, and an + // error returned here would cancel the session it exists to preserve. + // The receiving loop notices a genuinely dead connection within one read. + log.Debugf(gCtx, "Failed to send websocket keepalive ping: %v", err) + } + } + } + }) g.Go(func() error { // When proxy.start returns (EOF from ssh, or the server closing the connection), - // cancel so the handover goroutine stops too and g.Wait can return. + // cancel so the handover and keepalive goroutines stop too and g.Wait can return. defer cancel() return proxy.start(gCtx, src, wrappedDst) }) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 1915cc07c88..a52bc1a919d 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -38,7 +38,7 @@ type testClient struct { Cleanup func() } -func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, errChan chan error) *testClient { +func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, errChan chan error) *testClient { ctx := cmdio.MockDiscard(t.Context()) clientInput, clientInputWriter := io.Pipe() clientOutput := newTestBuffer(t) @@ -49,13 +49,11 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() return conn, err } if requestHandoverTick == nil { - requestHandoverTick = func() <-chan time.Time { - return time.After(time.Hour) - } + requestHandoverTick = neverTick } wg := sync.WaitGroup{} wg.Go(func() { - err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, createConn) + err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, createConn) if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, io.ErrClosedPipe) { if errChan != nil { errChan <- err @@ -78,7 +76,7 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() func TestClientServerEcho(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - client := createTestClient(t, server.URL, nil, nil) + client := createTestClient(t, server.URL, nil, time.Hour, nil) defer client.Cleanup() testMsg1 := []byte("test message 1\n") @@ -100,9 +98,9 @@ func TestClientServerEcho(t *testing.T) { func TestMultipleClients(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - client1 := createTestClient(t, server.URL, nil, nil) + client1 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client1.Cleanup() - client2 := createTestClient(t, server.URL, nil, nil) + client2 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client2.Cleanup() messageCount := 10 @@ -131,9 +129,9 @@ func TestMaxClients(t *testing.T) { maxClients := 2 server := createTestServer(t, maxClients, time.Hour) defer server.Close() - client1 := createTestClient(t, server.URL, nil, nil) + client1 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client1.Cleanup() - client2 := createTestClient(t, server.URL, nil, nil) + client2 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client2.Cleanup() testMsg1 := []byte("test message 1\n") @@ -147,7 +145,7 @@ func TestMaxClients(t *testing.T) { require.NoError(t, err) errChan := make(chan error, 1) - client3 := createTestClient(t, server.URL, nil, errChan) + client3 := createTestClient(t, server.URL, nil, time.Hour, errChan) defer client3.Cleanup() select { case err = <-errChan: @@ -158,6 +156,17 @@ func TestMaxClients(t *testing.T) { } func TestHandover(t *testing.T) { + t.Run("without keepalive", func(t *testing.T) { + runHandoverExchange(t, time.Hour) + }) + // Pings share the proxy's serialised write path with the data stream: they must not corrupt or + // reorder it, nor trip gorilla's concurrent-write panic. + t.Run("with keepalive", func(t *testing.T) { + runHandoverExchange(t, time.Millisecond) + }) +} + +func runHandoverExchange(t *testing.T, keepaliveInterval time.Duration) { server := createTestServer(t, 2, time.Hour) defer server.Close() @@ -165,7 +174,7 @@ func TestHandover(t *testing.T) { requestHandoverTick := func() <-chan time.Time { return handoverChan } - client := createTestClient(t, server.URL, requestHandoverTick, nil) + client := createTestClient(t, server.URL, requestHandoverTick, keepaliveInterval, nil) defer client.Cleanup() var expectedOutput []byte @@ -204,7 +213,7 @@ func TestQuickHandover(t *testing.T) { requestHandoverTick := func() <-chan time.Time { return handoverChan } - client := createTestClient(t, server.URL, requestHandoverTick, nil) + client := createTestClient(t, server.URL, requestHandoverTick, time.Hour, nil) defer client.Cleanup() var expectedOutput []byte @@ -256,7 +265,7 @@ func TestClientExitsWhenServerCommandFails(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) }() select { @@ -303,7 +312,7 @@ func TestClientTimesOutWhenServerSendsNothing(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) }() select { diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go new file mode 100644 index 00000000000..d1d779d4bff --- /dev/null +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -0,0 +1,173 @@ +package proxy + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +// startKeepaliveTestServer stands up a websocket peer that behaves like the SSH server side of the +// tunnel for an idle session: it sends the first bytes (which the client waits for before it +// considers the session established), then only reads. Pings it receives are reported on the +// returned channel and answered with a pong, the same way gorilla's default ping handler does. +func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{}) { + pings := make(chan struct{}, 1) + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + conn.SetPingHandler(func(appData string) error { + select { + case pings <- struct{}{}: + default: + } + return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) + }) + if err := conn.WriteMessage(websocket.BinaryMessage, []byte("SSH-2.0-test\r\n")); err != nil { + return + } + // Ping handlers only run while a read is in progress, so keep reading until the client goes away. + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + return server, pings +} + +func keepaliveTestDialer(serverURL string, onConn func(*websocket.Conn)) createWebsocketConnectionFunc { + wsURL := "ws" + serverURL[4:] + return func(ctx context.Context, connID string) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + if err != nil { + return nil, err + } + if onConn != nil { + onConn(conn) + } + return conn, nil + } +} + +// TestKeepalivePingReachesServer covers the fix itself: an idle session sends no data, so without +// the keepalive nothing at all crosses the websocket and the server side eventually reaps the +// stream it considers dead. +func TestKeepalivePingReachesServer(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, pings := startKeepaliveTestServer(t) + defer server.Close() + + // Never written to: the session stays idle for the whole test. + src, _ := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, nil)) + }() + + select { + case <-pings: + case err := <-done: + t.Fatalf("session ended before a keepalive ping arrived: %v", err) + case <-time.After(10 * time.Second): + t.Fatal("no keepalive ping arrived at the server") + } +} + +// TestKeepalivePingFailureDoesNotEndSession asserts a keepalive can never be the thing that ends a +// session: it has no information the data loops lack, and it runs in the errgroup that drives them, +// so a returned error would tear the session down. +func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + failWrites := func(conn *websocket.Conn) { + // A deadline in the past fails every write on this connection, so every ping fails. + // Set before the connection is handed to the proxy, so no writer can be in flight. + // Reads are unaffected: the connection is otherwise healthy and the session must survive. + conn.SetWriteDeadline(time.Now().Add(-time.Second)) // nolint:errcheck + } + + src, _ := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, failWrites)) + }() + + // Long enough for many pings to be attempted and fail. + select { + case err := <-done: + t.Fatalf("session ended after a failed keepalive ping: %v", err) + case <-time.After(2 * time.Second): + } +} + +// TestKeepalivePingBlockedByHandoverDoesNotDeadlock asserts the invariant the keepalive design +// rests on: a ping sent through the proxy's serialised write path blocks for the duration of a +// handover, and a handover waits on the receiving loop rather than that write path, so the two +// cannot deadlock each other. +func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) { + ctx := t.Context() + + server := setupTestServer(ctx, t) + defer server.Cleanup() + + // Holds the handover open at the point where it has taken the write path but has not yet + // completed, which is when a ping must block rather than break the handover. + handoverDialing := make(chan struct{}) + releaseHandover := make(chan struct{}) + var dials atomic.Int32 + + client := setupTestClientWithDialHook(ctx, t, server.URL, func() { + if dials.Add(1) > 1 { + close(handoverDialing) + <-releaseHandover + } + }) + defer client.Cleanup() + + handoverDone := make(chan error, 1) + go func() { + handoverDone <- client.Proxy.initiateHandover(ctx) + }() + <-handoverDialing + + pingDone := make(chan error, 1) + go func() { + pingDone <- client.Proxy.sendMessage(websocket.PingMessage, nil) + }() + + select { + case err := <-pingDone: + t.Fatalf("ping was written while a handover held the write path: %v", err) + case <-time.After(100 * time.Millisecond): + } + + close(releaseHandover) + + select { + case err := <-handoverDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("handover deadlocked while a keepalive ping waited on the write path") + } + select { + case err := <-pingDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("keepalive ping never completed after the handover finished") + } +} diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index 0e1db9021e2..1e3af34f40d 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -139,12 +139,26 @@ func createTestWebsocketConnection(url string) (*websocket.Conn, error) { return conn, err } +// neverTick is a tick channel that never fires, for tests that don't exercise a periodic behaviour. +func neverTick() <-chan time.Time { + return time.After(time.Hour) +} + func setupTestClient(ctx context.Context, t *testing.T, serverURL string) *TestProxy { + return setupTestClientWithDialHook(ctx, t, serverURL, nil) +} + +// setupTestClientWithDialHook is setupTestClient with a hook called on every websocket dial: the +// initial connection and each one a handover creates. +func setupTestClientWithDialHook(ctx context.Context, t *testing.T, serverURL string, onDial func()) *TestProxy { ctx = log.NewContext(ctx, log.GetLogger(ctx).With("Client", true)) clientInput, clientInputWriter := io.Pipe() clientOutput := newTestBuffer(t) wsURL := "ws" + serverURL[4:] clientProxy := newProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) { + if onDial != nil { + onDial() + } return createTestWebsocketConnection(wsURL) }) err := clientProxy.connect(ctx) From 81d83b890e32b79f58b6a14cefc37793324d69ee Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:42:00 +0000 Subject: [PATCH 2/6] Log each keepalive ping, not just failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit End-to-end verification against dogfood showed the far end never returns a pong: across 21-minute, 45-minute and forced-handover runs the pong handler logged nothing, while ping writes never failed. Control frames do not make the round trip on this transport, and the outbound ping alone is what keeps the stream from being reaped. That leaves a support engineer reading a customer's debug log with no positive evidence that keepalives were flowing — only the absence of failures, which is indistinguishable from a build without the keepalive. Log each successful ping instead, at debug level: three lines a minute on a transport whose debug log already carries full HTTP bodies. Verified end to end: ping lines appear at exactly 20-second intervals on an idle session, pong lines remain absent. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 47c4cb09b6d..a71d53c692b 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -106,6 +106,10 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque // error returned here would cancel the session it exists to preserve. // The receiving loop notices a genuinely dead connection within one read. log.Debugf(gCtx, "Failed to send websocket keepalive ping: %v", err) + } else { + // The driver proxy does not return pongs (verified end to end), so this + // line is the only evidence in a customer's log that pings were flowing. + log.Debugf(gCtx, "Sent websocket keepalive ping") } } } From bede5759e2b088364391887a98fdb10445507e44 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:56:08 +0000 Subject: [PATCH 3/6] Send keepalive pings with WriteControl, off the handover mutex MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of #6358 pointed out that routing pings through sendMessage puts an unbounded write on the connection's shared write path. WriteMessage takes the handover mutex and sets no deadline, and close() needs that same mutex, so a ping that parks on a stalled or half-open socket — a full send buffer, no RST — holds up the closing handshake until the kernel abandons its retransmits, roughly 15 minutes with Linux defaults. Context cancellation cannot interrupt a goroutine parked in a blocking write. The probability is low for a purely idle session, whose 8-byte control frames go into a near-empty buffer, but not for the path this feature exists to serve: data flows, the peer vanishes mid-transfer leaving unacked bytes in the buffer, the session goes idle, and the next ping parks. That is the feature's own domain, so the exposure belongs to this change even though the hazard predates it on the data path. Pings now go out with WriteControl on the loaded connection, taking no handover mutex. gorilla explicitly 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, so a stalled ping can hold that lock for at most proxyPingWriteTimeout instead of minutes. The handover path is fully decoupled: a ping that ticks during a rotation goes to the connection being replaced and may simply fail, which is already non-fatal. This drops the "single concurrent writer" rationale for the mutex, which applies to WriteMessage and not to WriteControl, and adds a write deadline the original design ruled out. The rule it was protecting — a keepalive must never end a session — is untouched: a timed-out ping is logged at debug and the ticker continues. Tests: the two that asserted the mutex path were reworked, since one drove sendMessage directly and the other's past write deadline is now overridden by WriteControl's own. A ping is now asserted to complete during an in-flight handover rather than to block on it, its failure is induced at the socket, and a new case parks a ping in the socket write and requires the closing handshake to finish within the ping's deadline. Verified end to end: 10 pings at exact 20-second intervals across a 200-second idle session on dogfood, no failures, session intact. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 9 +- .../ssh/internal/proxy/keepalive_test.go | 175 ++++++++++++++---- experimental/ssh/internal/proxy/proxy.go | 13 ++ 3 files changed, 158 insertions(+), 39 deletions(-) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index a71d53c692b..59e13baf979 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -97,11 +97,10 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque case <-gCtx.Done(): return gCtx.Err() case <-ticker.C: - // Pings take the same serialised write path as data — gorilla forbids - // concurrent writers — so a ping that ticks during a handover blocks until it - // finishes and then goes out late. Harmless: a handover establishes a fresh - // connection, which resets the peer's idle clock anyway. - if err := proxy.sendMessage(websocket.PingMessage, nil); err != nil { + // A ping that ticks during a handover goes to the connection being replaced and + // may simply fail. Harmless: a handover establishes a fresh connection, which + // resets the peer's idle clock anyway, and the next tick uses the new one. + if err := proxy.sendPing(); err != nil { // Never fatal. A failed ping knows nothing the data loops don't, and an // error returned here would cancel the session it exists to preserve. // The receiving loop notices a genuinely dead connection within one read. diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go index d1d779d4bff..a7fc0c1bf10 100644 --- a/experimental/ssh/internal/proxy/keepalive_test.go +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -2,10 +2,14 @@ package proxy import ( "context" + "errors" "fmt" "io" + "net" "net/http" "net/http/httptest" + "os" + "sync" "sync/atomic" "testing" "time" @@ -15,6 +19,55 @@ import ( "github.com/stretchr/testify/require" ) +// Write modes for pausableConn. +const ( + connWriteOK = iota + connWriteFail + connWritePark +) + +var errTestWriteFailed = errors.New("test: socket write failed") + +// pausableConn emulates the socket conditions a keepalive meets on a stalled peer: writes can be +// made to fail outright, or to park the way a full send buffer does — blocking until the write +// deadline expires. Wrapping the socket rather than the websocket keeps the production write path +// (gorilla's own locking and deadline handling) in the test. +type pausableConn struct { + net.Conn + mode atomic.Int32 + deadline atomic.Pointer[time.Time] + parked chan struct{} + signalParked func() +} + +func newPausableConn(conn net.Conn) *pausableConn { + parked := make(chan struct{}) + return &pausableConn{ + Conn: conn, + parked: parked, + signalParked: sync.OnceFunc(func() { close(parked) }), + } +} + +func (c *pausableConn) SetWriteDeadline(t time.Time) error { + c.deadline.Store(&t) + return c.Conn.SetWriteDeadline(t) +} + +func (c *pausableConn) Write(p []byte) (int, error) { + switch c.mode.Load() { + case connWriteFail: + return 0, errTestWriteFailed + case connWritePark: + c.signalParked() + if d := c.deadline.Load(); d != nil && !d.IsZero() { + time.Sleep(time.Until(*d)) + } + return 0, os.ErrDeadlineExceeded + } + return c.Conn.Write(p) +} + // startKeepaliveTestServer stands up a websocket peer that behaves like the SSH server side of the // tunnel for an idle session: it sends the first bytes (which the client waits for before it // considers the session established), then only reads. Pings it receives are reported on the @@ -48,17 +101,26 @@ func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{}) return server, pings } -func keepaliveTestDialer(serverURL string, onConn func(*websocket.Conn)) createWebsocketConnectionFunc { +// keepaliveTestDialer returns a connection factory for RunClientProxy. onNetConn, when set, receives +// each connection's underlying socket so a test can control how its writes behave. +func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) createWebsocketConnectionFunc { wsURL := "ws" + serverURL[4:] + dialer := websocket.Dialer{ + NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := net.Dial(network, addr) + if err != nil { + return nil, err + } + wrapped := newPausableConn(conn) + if onNetConn != nil { + onNetConn(wrapped) + } + return wrapped, nil + }, + } return func(ctx context.Context, connID string) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose - if err != nil { - return nil, err - } - if onConn != nil { - onConn(conn) - } - return conn, nil + conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + return conn, err } } @@ -94,19 +156,19 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { server, _ := startKeepaliveTestServer(t) defer server.Close() - failWrites := func(conn *websocket.Conn) { - // A deadline in the past fails every write on this connection, so every ping fails. - // Set before the connection is handed to the proxy, so no writer can be in flight. - // Reads are unaffected: the connection is otherwise healthy and the session must survive. - conn.SetWriteDeadline(time.Now().Add(-time.Second)) // nolint:errcheck - } - + var socket atomic.Pointer[pausableConn] src, _ := io.Pipe() done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, failWrites)) + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) }() + // Fail every write once the connection is up, leaving reads healthy: the connection is otherwise + // fine and the session must survive the pings that then fail. + require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond) + socket.Load().mode.Store(connWriteFail) + // Long enough for many pings to be attempted and fail. select { case err := <-done: @@ -115,18 +177,63 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { } } -// TestKeepalivePingBlockedByHandoverDoesNotDeadlock asserts the invariant the keepalive design -// rests on: a ping sent through the proxy's serialised write path blocks for the duration of a -// handover, and a handover waits on the receiving loop rather than that write path, so the two -// cannot deadlock each other. -func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) { +// TestKeepalivePingParkedInWriteDoesNotStallClose covers the hazard of writing from a third +// goroutine: a ping on a stalled or half-open connection parks in the socket write while holding the +// websocket's write lock, which the closing handshake also needs. Its deadline is what keeps that +// from lasting until the kernel abandons its retransmits, minutes later. +// +// Scoped to the ping's own contribution. pausableConn parks only for as long as the write's deadline, +// and a write whose caller set none fails at once, so the pre-existing unbounded park on the data +// path is out of the picture. The session's own shutdown cannot be measured here either: it 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 — with or without a keepalive. +func TestKeepalivePingParkedInWriteDoesNotStallClose(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + var socket atomic.Pointer[pausableConn] + proxy := newProxyConnection(keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) + require.NoError(t, proxy.connect(ctx)) + require.NotNil(t, socket.Load()) + + socket.Load().mode.Store(connWritePark) + pingDone := make(chan error, 1) + go func() { + pingDone <- proxy.sendPing() + }() + select { + case <-socket.Load().parked: + case <-time.After(10 * time.Second): + t.Fatal("the keepalive ping did not park in the socket write") + } + + start := time.Now() + closeErr := proxy.close() + require.Less(t, time.Since(start), proxyPingWriteTimeout+5*time.Second, + "the closing handshake waited on the parked keepalive ping for longer than its write deadline allows") + // The write itself fails, which close() reports; the point is that it was not held indefinitely. + require.Error(t, closeErr) + + select { + case err := <-pingDone: + require.Error(t, err, "a parked ping write must end in an error, not succeed") + case <-time.After(10 * time.Second): + t.Fatal("the keepalive ping never returned from its parked write") + } +} + +// TestKeepalivePingDuringHandoverDoesNotDisruptIt asserts the two periodic behaviours of the tunnel +// stay independent: a ping sent while a handover is in flight neither waits for the handover nor +// breaks it. The keepalive writes from a third goroutine, so nothing else guarantees this. +func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) { ctx := t.Context() server := setupTestServer(ctx, t) defer server.Cleanup() - // Holds the handover open at the point where it has taken the write path but has not yet - // completed, which is when a ping must block rather than break the handover. + // Holds the handover open at the point where it has taken the handover mutex but has not yet + // swapped the connection, which is when a ping must neither block nor interfere. handoverDialing := make(chan struct{}) releaseHandover := make(chan struct{}) var dials atomic.Int32 @@ -147,13 +254,14 @@ func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) { pingDone := make(chan error, 1) go func() { - pingDone <- client.Proxy.sendMessage(websocket.PingMessage, nil) + pingDone <- client.Proxy.sendPing() }() select { case err := <-pingDone: - t.Fatalf("ping was written while a handover held the write path: %v", err) - case <-time.After(100 * time.Millisecond): + require.NoError(t, err) + case <-time.After(proxyPingWriteTimeout): + t.Fatal("keepalive ping waited on the in-flight handover") } close(releaseHandover) @@ -162,12 +270,11 @@ func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) { case err := <-handoverDone: require.NoError(t, err) case <-time.After(10 * time.Second): - t.Fatal("handover deadlocked while a keepalive ping waited on the write path") - } - select { - case err := <-pingDone: - require.NoError(t, err) - case <-time.After(10 * time.Second): - t.Fatal("keepalive ping never completed after the handover finished") + t.Fatal("handover did not complete after a concurrent keepalive ping") } + + // The tunnel still carries data on the connection the handover installed. + _, err := client.Input.Write(createTestMessage("client", 1)) + require.NoError(t, err) + require.NoError(t, server.Output.WaitForWrite(createTestMessage("client", 1))) } diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index e761cbf54c4..90526e17ea9 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -29,6 +29,10 @@ const ( proxyHandoverInitTimeout = 30 * time.Second // Timeout for the handover process, when accepted by the server. proxyHandoverAcceptTimeout = 25 * time.Second + // Bounds how long a keepalive ping may hold the websocket's write lock. A stalled or half-open + // connection parks a write until the kernel gives up retransmitting (~15 minutes with Linux + // defaults), and close() and the sending loop need that same lock, so the ping caps its wait. + proxyPingWriteTimeout = 5 * time.Second ) // handoverCoordination holds the context and channels used to coordinate a single handover operation @@ -204,6 +208,15 @@ func (pc *proxyConnection) sendMessage(mt int, data []byte) error { return conn.WriteMessage(mt, data) } +// sendPing writes a keepalive ping on the current connection. Unlike sendMessage it takes neither +// the handover mutex nor an unbounded wait: gorilla permits WriteControl concurrently with the data +// writes, and its deadline bounds how long a stalled socket holds the connection's write lock, which +// close() and the sending loop also need. +func (pc *proxyConnection) sendPing() error { + conn := pc.conn.Load() + return conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout)) +} + func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) error { for { if ctx.Err() != nil { From 88dcba3403f93d119a28bbbe7e8ad6c225c9c703 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 25 Aug 2026 07:49:08 +0000 Subject: [PATCH 4/6] Fix a test comment left stale by the WriteControl change The keepalive no longer goes through sendMessage, so pings and the data stream share the connection's write lock rather than the proxy's serialised write path. The property the subtest protects is unchanged. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client_server_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index a52bc1a919d..7aa45d35e81 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -159,8 +159,8 @@ func TestHandover(t *testing.T) { t.Run("without keepalive", func(t *testing.T) { runHandoverExchange(t, time.Hour) }) - // Pings share the proxy's serialised write path with the data stream: they must not corrupt or - // reorder it, nor trip gorilla's concurrent-write panic. + // Pings and the data stream share the connection's write lock: they must not corrupt or reorder + // the stream, nor trip gorilla's concurrent-write panic. t.Run("with keepalive", func(t *testing.T) { runHandoverExchange(t, time.Millisecond) }) From f2086b34b2b21e1f9dd11950dcd239ea94e9fed3 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 25 Aug 2026 08:20:04 +0000 Subject: [PATCH 5/6] Close the websocket on teardown so a dead write path cannot hang the session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent verification (recorded on DECO-28186) found that the PR's claim "a failed ping never ends a session" was only true of the errgroup, not of gorilla's connection state. Any failed write puts a gorilla connection into a permanent write-error state, so one timed-out keepalive ping stops every later write — data frames and the close frame alike. The session then died through the data path instead: data written afterwards never arrived, and because the close message could not go out, the peer never closed the connection, so the receiving loop stayed blocked in ReadMessage and g.Wait never returned. A silent hang, which is the very symptom this change exists to remove. The teardown goroutine already documents the fix as its intent — "we close the connection and the source ... to unblock them" — but pc.close only sends a close message, which needs a live peer to act on it. It now closes the connection too, so both loops unblock and the session exits with the sending loop's error instead of hanging. The ping failure itself stays non-fatal, and deliberately so: reads are unaffected by a poisoned write path and may still be delivering output the user is waiting on — a long job that prints nothing for minutes is one of the cases this feature was written for. Ending the session there would cut off data that is still arriving. It is logged at warn rather than debug, since it now means the tunnel can no longer send anything. Tests, after mutation-testing every assertion: - The parked-write fake modelled a deadline-less write as an instant failure, the opposite of an unbounded park, so the test could not tell the bounded write path from the one it replaced. It now parks such a write for longer than any assertion, and the reverted-to-sendMessage mutant fails. - TestKeepalivePingFailureDoesNotEndSession was vacuous: the fake's broken close handshake meant the session could not end for reasons unrelated to the keepalive. Closing the connection on teardown restores its power — a mutant that makes a failed ping fatal now fails it. - A t.Fatal in the handover test left the dial hook parked, and cleanup, which waits on the proxy loops, then deadlocked and took the package down with a timeout panic. The release is now deferred after cleanup so it runs first; the same mutant is caught in 5s as one clean failure. - New: TestKeepalivePingFailureDoesNotHangTheSession asserts a session whose write path has been poisoned by a failed ping ends promptly with an error. All five keepalive assertions now kill their mutants. Verified end to end on dogfood that a normal exit is still clean after the teardown change: exit 0, 70s idle, 3 pings, no warnings, no spurious disconnect message. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 10 +-- .../ssh/internal/proxy/keepalive_test.go | 67 ++++++++++++++++--- experimental/ssh/internal/proxy/proxy.go | 16 ++++- 3 files changed, 77 insertions(+), 16 deletions(-) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 59e13baf979..9d40cac5169 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -101,10 +101,12 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque // may simply fail. Harmless: a handover establishes a fresh connection, which // resets the peer's idle clock anyway, and the next tick uses the new one. if err := proxy.sendPing(); err != nil { - // Never fatal. A failed ping knows nothing the data loops don't, and an - // error returned here would cancel the session it exists to preserve. - // The receiving loop notices a genuinely dead connection within one read. - log.Debugf(gCtx, "Failed to send websocket keepalive ping: %v", err) + // Not fatal, but not harmless either: gorilla puts the connection into a + // permanent write-error state after any failed write, so nothing more can + // be sent on it. Reads are unaffected and may still be delivering output + // the user is waiting on, so the session is left to end the way it would + // anyway — the next write fails and the sending loop reports it. + log.Warnf(gCtx, "Failed to send websocket keepalive ping, the connection can no longer send: %v", err) } else { // The driver proxy does not return pongs (verified end to end), so this // line is the only evidence in a customer's log that pings were flowing. diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go index a7fc0c1bf10..c972624f6f1 100644 --- a/experimental/ssh/internal/proxy/keepalive_test.go +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -28,10 +28,15 @@ const ( var errTestWriteFailed = errors.New("test: socket write failed") +// unboundedParkDuration stands in for "parks until the kernel gives up", which is what a write with +// no deadline does on a stalled socket. It must outlast the assertions of any test that parks a +// write, so that a write path which forgets to bound itself is measured as a stall, not as a failure. +const unboundedParkDuration = 30 * time.Second + // pausableConn emulates the socket conditions a keepalive meets on a stalled peer: writes can be -// made to fail outright, or to park the way a full send buffer does — blocking until the write -// deadline expires. Wrapping the socket rather than the websocket keeps the production write path -// (gorilla's own locking and deadline handling) in the test. +// made to fail outright, or to park the way a full send buffer does — until the write's deadline +// expires, or effectively forever if it has none. Wrapping the socket rather than the websocket keeps +// the production write path (gorilla's own locking and deadline handling) in the test. type pausableConn struct { net.Conn mode atomic.Int32 @@ -62,6 +67,8 @@ func (c *pausableConn) Write(p []byte) (int, error) { c.signalParked() if d := c.deadline.Load(); d != nil && !d.IsZero() { time.Sleep(time.Until(*d)) + } else { + time.Sleep(unboundedParkDuration) } return 0, os.ErrDeadlineExceeded } @@ -169,11 +176,11 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond) socket.Load().mode.Store(connWriteFail) - // Long enough for many pings to be attempted and fail. + // Long enough for tens of pings to be attempted and fail at the 20ms interval above. select { case err := <-done: t.Fatalf("session ended after a failed keepalive ping: %v", err) - case <-time.After(2 * time.Second): + case <-time.After(time.Second): } } @@ -182,11 +189,8 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { // websocket's write lock, which the closing handshake also needs. Its deadline is what keeps that // from lasting until the kernel abandons its retransmits, minutes later. // -// Scoped to the ping's own contribution. pausableConn parks only for as long as the write's deadline, -// and a write whose caller set none fails at once, so the pre-existing unbounded park on the data -// path is out of the picture. The session's own shutdown cannot be measured here either: it 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 — with or without a keepalive. +// A ping write with no deadline parks for unboundedParkDuration here, so this fails if the ping ever +// goes back to a write path that does not bound itself. func TestKeepalivePingParkedInWriteDoesNotStallClose(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) server, _ := startKeepaliveTestServer(t) @@ -223,6 +227,42 @@ func TestKeepalivePingParkedInWriteDoesNotStallClose(t *testing.T) { } } +// TestKeepalivePingFailureDoesNotHangTheSession covers what a failed ping actually costs. gorilla +// puts the connection into a permanent write-error state after any failed write, so one timed-out +// keepalive stops the close message going out too — and a session whose close message never reaches +// the peer used to wait forever for the peer to close the connection, which is the same silent +// black hole this feature exists to remove. The session must end, promptly and with an error. +func TestKeepalivePingFailureDoesNotHangTheSession(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + var socket atomic.Pointer[pausableConn] + src, srcWriter := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) + }() + + require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond) + socket.Load().mode.Store(connWriteFail) + // Let a ping fail, which is what poisons the connection. + time.Sleep(100 * time.Millisecond) + + // The session ends the way it would without a keepalive at all: the next write fails and the + // sending loop reports it. Before, the teardown could not close the connection and this hung. + _, err := srcWriter.Write([]byte("keystroke")) + require.NoError(t, err) + + select { + case err := <-done: + require.Error(t, err, "a session that can no longer send must end with an error, not silently") + case <-time.After(30 * time.Second): + t.Fatal("session hung after a failed keepalive ping poisoned the connection") + } +} + // TestKeepalivePingDuringHandoverDoesNotDisruptIt asserts the two periodic behaviours of the tunnel // stay independent: a ping sent while a handover is in flight neither waits for the handover nor // breaks it. The keepalive writes from a third goroutine, so nothing else guarantees this. @@ -236,6 +276,7 @@ func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) { // swapped the connection, which is when a ping must neither block nor interfere. handoverDialing := make(chan struct{}) releaseHandover := make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseHandover) }) var dials atomic.Int32 client := setupTestClientWithDialHook(ctx, t, server.URL, func() { @@ -245,6 +286,10 @@ func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) { } }) defer client.Cleanup() + // Deferred after client.Cleanup so it runs before it: a t.Fatal below would otherwise leave the + // handover parked in the dial hook, and cleanup waits on proxy loops that cannot finish until it + // is released — which wedges the whole package instead of failing one test. + defer release() handoverDone := make(chan error, 1) go func() { @@ -264,7 +309,7 @@ func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) { t.Fatal("keepalive ping waited on the in-flight handover") } - close(releaseHandover) + release() select { case err := <-handoverDone: diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 90526e17ea9..685b02a9b46 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "os" "sync" @@ -138,7 +139,7 @@ func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io. // Both loops can still be stuck on conn.ReadMessage or src.Read and won't notice context cancellation, // so we close the connection and the source (sshd stdout pipe or ssh client stdio) to unblock them. <-gCtx.Done() - return errors.Join(pc.close(), pc.closeSource(src)) + return errors.Join(pc.close(), pc.closeConnection(), pc.closeSource(src)) }) err := g.Wait() if err == nil || isNormalClosure(err) { @@ -274,6 +275,19 @@ func (pc *proxyConnection) close() error { return nil } +// closeConnection closes the underlying websocket. The close message pc.close sends only ends the +// session if the peer is still there to react to it by closing the connection, and it does not even +// go out once a failed write has put the connection into gorilla's permanent write-error state (one +// timed-out keepalive ping is enough). Without this the receiving loop stays blocked in ReadMessage +// and the session hangs instead of exiting. +func (pc *proxyConnection) closeConnection() error { + err := pc.conn.Load().Close() + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + func (pc *proxyConnection) closeSource(src io.ReadCloser) error { err := src.Close() if err != nil && (errors.Is(err, os.ErrClosed) || errors.Is(err, io.ErrClosedPipe)) { From 5fbfcedad7fef76ad8c4d4b4d721d604ffb188e6 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:28:19 +0200 Subject: [PATCH 6/6] Keep idle SSH tunnel sessions alive on dedicated clusters (#6382) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Stacked on #6358: the base branch is that PR's branch, so this diff is only the change on top of it. GitHub retargets to `main` when #6358 merges. ## Changes #6358 fixes idle-session drops on serverless and does nothing for a dedicated cluster. Its websocket ping is a *control* frame: the driver proxy terminates it itself, so it never becomes payload on the leg past it, and that leg is reaped after ~8m20s of carrying nothing. Zero pongs came back in any run on either compute type — the frame does not traverse. Serverless sits behind a different reaper that the ping does reset, which is why it fixes serverless and only serverless. The fix is to generate traffic the tunnel forwards end to end. An SSH keepalive reply is a real SSH packet, so the proxy loops carry it as a websocket *data* frame across every hop — which is why setting `ServerAliveInterval` by hand is a total workaround. The CLI now asks for it itself, from both ends: - **`ServerAliveInterval 30`** in `buildSSHArgs` (the `connect` path) and `GenerateHostConfig` (the `setup` and `--ide` paths). Needs nothing from the compute, so it also fixes a cluster already running a server binary from an older CLI. - **`ClientAliveInterval 30`** in the sshd config the tunnel's server writes. Covers clients the CLI never configures — a hand-written `ProxyCommand` block, or an IDE supplying its own ssh options. Either half alone is sufficient (measured below). Both are in because they close *different* gaps and neither subsumes the other. The websocket ping stays: it is what fixes serverless, and it keeps the client↔control-plane hop warm regardless of the user's ssh config. **One behaviour change, deliberately.** These options bring in `ServerAliveCountMax` / `ClientAliveCountMax` (OpenSSH default 3), so a tunnel that stops responding is now torn down after ~90s with ssh's own "server not responding" message instead of hanging. That is the trade we want — today a reaped idle session is a silent black hole, which is the original report — and 90s is well clear of the longest legitimate pause on a healthy tunnel, the up to 30s a handover can hold the sending loop (`proxyHandoverInitTimeout`). `sshdConfigContent` is `prepareSSHDConfig`'s config string moved verbatim into a function plus the one new line, so the content is testable without mocking three secret reads. #6358's changelog fragment is updated rather than joined by a second, contradicting one. An out-of-band keepalive frame of our own was rejected: `runReceivingLoop` writes every binary message straight into the SSH byte stream, so it needs a framing change on both ends and only helps once the uploaded server binary is new enough. Resolves DECO-28186 for dedicated clusters; #6358 resolves it for serverless. ## Why Dedicated compute is what the README requires for remote development in an IDE, so the compute type most affected was the one #6358 left broken. Shipping only the websocket ping would close the ticket with the reported failure still in place for those users, and add a log line that makes an idle session look healthy while it is dying. ## Tests Three unit tests, one per site the directive has to appear in, each also bounding the interval against the reap window and against the handover pause. All six mutants killed: dropping the option from each of the three sites, setting either constant to `0`, and pushing the interval past the reap window. `./task test-exp-ssh`, `./task lint` (0 issues, all three modules), `./task fmt`, and the whitespace, deadcode and changelog checks are all clean; the four touched packages are green under `-race`. ### End-to-end (dogfood, dedicated clusters) One dedicated single-node cluster **per build** — two tunnels on one cluster collide on port 7772 and take each other down. DBR 17.3 LTS, m5d.xlarge, SINGLE_USER; distinct pinned version per build, since `uploadReleases` skips the upload when the versioned workspace path already exists. Every session ran on a real PTY, was left **completely untyped** for the whole window, then sent exactly one command. Idle is measured remote-clock to remote-clock. | run | build | cli `ServerAliveInterval` | sshd `ClientAliveInterval` | idle | pings/pongs | verdict | | --- | --- | --- | --- | --- | --- | --- | | `dc-base` | #6358 head | none | none | died at +8m20s | 30 / 0 | **FAIL** | | `dc-fix` | this PR | 30 | 30 | 613 s | 32 / 0 | **PASS** | | `dc-sshd-only` | this PR | **0 (forced off)** | 30 | 655 s | 34 / 0 | **PASS** | | `dc-fix-long` | this PR | 30 | 30 | 1194 s | 61 / 0 | **PASS** | - `dc-base` reproduces the bug on #6358's own head, as reported: alive-looking for the whole window, killed by the first keystroke. The server lost its half at `18:04:46` — 8m20s after the session went quiet — with `websocket: close 1006 (abnormal closure): unexpected EOF`. The client had sent 30 pings with zero failures by then; its next two failed, after the server was already gone. Ended `close 4000: Handler crashed: ...ClosedStreamException`, `exit status 255`. - `dc-sshd-only` isolates the server half: same build, but ssh invoked by hand with `-o ServerAliveInterval=0`, so sshd drives the only SSH-level keepalive. A command-line `-o` beats any config file in OpenSSH, and `ClientAliveInterval 30` was read back off the cluster's generated `sshd_config`. The complementary control for the client half — old server binary, client-side option only — passed earlier on the same setup. - `dc-fix-long` is past two reap windows, so the PASS is a real fix rather than a delayed drop. Full measurements and log excerpts are on DECO-28186. The test clusters have been terminated. _This PR and its description were written by Isaac._ --------- Co-authored-by: Isaac Co-authored-by: Russell Clarey --- .nextchanges/cli/ssh-tunnel-keepalive.md | 2 +- experimental/ssh/internal/client/client.go | 1 + .../internal/client/client_internal_test.go | 14 ++++++++ experimental/ssh/internal/server/sshd.go | 35 ++++++++++++++----- experimental/ssh/internal/server/sshd_test.go | 16 +++++++++ .../ssh/internal/sshconfig/sshconfig.go | 15 +++++++- .../ssh/internal/sshconfig/sshconfig_test.go | 16 +++++++++ 7 files changed, 88 insertions(+), 11 deletions(-) diff --git a/.nextchanges/cli/ssh-tunnel-keepalive.md b/.nextchanges/cli/ssh-tunnel-keepalive.md index 71802b97adb..cf1d33a18e5 100644 --- a/.nextchanges/cli/ssh-tunnel-keepalive.md +++ b/.nextchanges/cli/ssh-tunnel-keepalive.md @@ -1 +1 @@ -Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes. The tunnel now sends a websocket keepalive every 20 seconds, so a session nobody is typing into stays connected without setting `ServerAliveInterval` in the SSH client config. +Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes, on dedicated clusters and on serverless. The tunnel now keeps itself warm: the SSH client and the SSH server on the compute exchange keepalives every 30 seconds, and the CLI's proxy pings the tunnel's websocket every 20 seconds. A session nobody is typing into stays connected, with no need to set `ServerAliveInterval` by hand. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index b2a8c9d0b7f..031bbd4d7af 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -820,6 +820,7 @@ func buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome strin "-o", "IdentitiesOnly=yes", "-o", "StrictHostKeyChecking=accept-new", "-o", "ConnectTimeout=360", + "-o", "ServerAliveInterval=" + strconv.Itoa(sshconfig.ServerAliveIntervalSeconds), "-o", "ProxyCommand=" + proxyCommand, } if opts.UserKnownHostsFile != "" { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 1e3accbaf46..cb47cd71bb8 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -5,10 +5,13 @@ import ( "encoding/json" "errors" "fmt" + "slices" + "strconv" "strings" "testing" "time" + "github.com/databricks/cli/experimental/ssh/internal/sshconfig" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go/experimental/mocks" @@ -373,6 +376,17 @@ func TestBuildRemoteShellArgs(t *testing.T) { }) } +func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { + args := buildSSHArgs("user", "/key", "proxy command", "myhost", "", ClientOptions{}) + + // ssh stops parsing options at the destination, so an option placed after the host would be + // treated as part of the remote command rather than as an ssh option. + optIdx := slices.Index(args, "ServerAliveInterval="+strconv.Itoa(sshconfig.ServerAliveIntervalSeconds)) + require.NotEqual(t, -1, optIdx, "ssh must be asked to send keepalives") + require.Equal(t, "-o", args[optIdx-1]) + assert.Less(t, optIdx, slices.Index(args, "myhost"), "the option must precede the destination host") +} + func TestBuildSSHArgsPTYPlacement(t *testing.T) { indexOf := func(args []string, want string) int { for i, a := range args { diff --git a/experimental/ssh/internal/server/sshd.go b/experimental/ssh/internal/server/sshd.go index bfafbbe5212..01fd0e3cdb0 100644 --- a/experimental/ssh/internal/server/sshd.go +++ b/experimental/ssh/internal/server/sshd.go @@ -9,6 +9,7 @@ import ( "os/exec" "path" "path/filepath" + "strconv" "strings" "github.com/databricks/cli/experimental/ssh/internal/keys" @@ -17,6 +18,18 @@ import ( "github.com/databricks/databricks-sdk-go" ) +// clientAliveIntervalSeconds is how often sshd asks the client to confirm it is still there. It +// drives the keepalive from the server end of the tunnel; sshconfig.ServerAliveIntervalSeconds +// documents why an SSH keepalive is what keeps the leg past the driver proxy from being reaped. +// Configuring it here covers clients the CLI does not configure — a hand-written ProxyCommand host +// block, or an IDE that supplies its own ssh options — where nothing sets ServerAliveInterval. +// The two intervals are deliberately independent: neither package imports the other, and each end +// keeps its own leg warm, so they need not track a single shared value. +// +// It also brings in ClientAliveCountMax (OpenSSH default 3), so sshd reclaims a session whose +// client has gone away after ~90s instead of holding it open until the server's shutdown delay. +const clientAliveIntervalSeconds = 30 + func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOptions) (string, error) { clientPublicKey, err := keys.GetSecret(ctx, client, opts.SecretScopeName, opts.AuthorizedKeySecretName) if err != nil { @@ -76,15 +89,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, } setEnv := setEnvBuf.String() - sshdConfigContent := "PubkeyAuthentication yes\n" + - "PasswordAuthentication no\n" + - "ChallengeResponseAuthentication no\n" + - "Subsystem sftp internal-sftp\n" + - "HostKey " + keyPath + "\n" + - "AuthorizedKeysFile " + authKeysPath + "\n" + - setEnv + "\n" - - if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent), 0o600); err != nil { + if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent(keyPath, authKeysPath, setEnv)), 0o600); err != nil { return "", err } @@ -97,6 +102,18 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, return sshdConfig, nil } +// sshdConfigContent assembles the configuration the tunnel's sshd runs with. +func sshdConfigContent(hostKeyPath, authorizedKeysPath, setEnv string) string { + return "PubkeyAuthentication yes\n" + + "PasswordAuthentication no\n" + + "ChallengeResponseAuthentication no\n" + + "ClientAliveInterval " + strconv.Itoa(clientAliveIntervalSeconds) + "\n" + + "Subsystem sftp internal-sftp\n" + + "HostKey " + hostKeyPath + "\n" + + "AuthorizedKeysFile " + authorizedKeysPath + "\n" + + setEnv + "\n" +} + func createSSHDProcess(ctx context.Context, configPath string) *exec.Cmd { return exec.CommandContext(ctx, "/usr/sbin/sshd", "-f", configPath, "-i") } diff --git a/experimental/ssh/internal/server/sshd_test.go b/experimental/ssh/internal/server/sshd_test.go index a453d987a00..4887f9895d8 100644 --- a/experimental/ssh/internal/server/sshd_test.go +++ b/experimental/ssh/internal/server/sshd_test.go @@ -1,6 +1,7 @@ package server import ( + "strconv" "testing" "github.com/stretchr/testify/assert" @@ -71,3 +72,18 @@ func TestEscapeEnvValue(t *testing.T) { }) } } + +func TestSSHDConfigSetsClientAliveInterval(t *testing.T) { + config := sshdConfigContent("/keys/server-private-key", "/keys/authorized_keys", `SetEnv FOO="bar"`) + + // This is the half of the keepalive that reaches clients the CLI never configures, so sshd has + // to drive it: without ClientAliveInterval sshd sends nothing on an idle session. + assert.Contains(t, config, "\nClientAliveInterval "+strconv.Itoa(clientAliveIntervalSeconds)+"\n") + + // The interval has to fire well inside the ~8 minute reap window, and ClientAliveCountMax + // (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy + // tunnel — the up to 30s a handover can hold the sending loop. The 30 below is + // proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced. + assert.Less(t, clientAliveIntervalSeconds, 8*60) + assert.Greater(t, 3*clientAliveIntervalSeconds, 30) +} diff --git a/experimental/ssh/internal/sshconfig/sshconfig.go b/experimental/ssh/internal/sshconfig/sshconfig.go index ad8ca0ee2a7..fbbf3b42fee 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig.go +++ b/experimental/ssh/internal/sshconfig/sshconfig.go @@ -19,6 +19,18 @@ const ( configDirName = ".databricks/ssh-tunnel-configs" ) +// ServerAliveIntervalSeconds is how often the ssh client asks the SSH server to confirm it is +// still there. The reply is a real SSH packet, so the keepalive puts payload bytes on every hop +// of the tunnel — and payload is what an idle session needs. The driver proxy terminates +// websocket control frames itself, so the proxy's own websocket ping never becomes payload on +// the leg past it, and that leg is reaped after ~8 minutes without any. +// +// It also brings in ServerAliveCountMax (OpenSSH default 3), so a tunnel that stops responding +// is torn down after ~90s with ssh's own "server not responding" message instead of hanging. +// That is well clear of the up to 30s a handover can hold the sending loop +// (proxyHandoverInitTimeout), the longest legitimate pause on a healthy tunnel. +const ServerAliveIntervalSeconds = 30 + func GetConfigDir(ctx context.Context) (string, error) { homeDir, err := env.UserHomeDir(ctx) if err != nil { @@ -206,9 +218,10 @@ func GenerateHostConfig(hostName, userName, identityFile, proxyCommand string) s Host %s User %s ConnectTimeout 360 + ServerAliveInterval %d StrictHostKeyChecking accept-new IdentitiesOnly yes IdentityFile %q ProxyCommand %s -`, hostName, userName, identityFile, proxyCommand) +`, hostName, userName, ServerAliveIntervalSeconds, identityFile, proxyCommand) } diff --git a/experimental/ssh/internal/sshconfig/sshconfig_test.go b/experimental/ssh/internal/sshconfig/sshconfig_test.go index 6c453910cdc..23abcbc34d7 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig_test.go +++ b/experimental/ssh/internal/sshconfig/sshconfig_test.go @@ -1,6 +1,7 @@ package sshconfig import ( + "fmt" "os" "path/filepath" "testing" @@ -10,6 +11,21 @@ import ( "github.com/stretchr/testify/require" ) +func TestGenerateHostConfigSetsServerAliveInterval(t *testing.T) { + config := GenerateHostConfig("myhost", "root", "/keys/myhost", "databricks ssh connect --proxy") + + // `ssh setup` and `--ide` reach ssh through this block and nothing else, so the option has to + // be in it. + assert.Contains(t, config, fmt.Sprintf("\n ServerAliveInterval %d\n", ServerAliveIntervalSeconds)) + + // The interval has to fire well inside the ~8 minute reap window, and ServerAliveCountMax + // (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy + // tunnel — the up to 30s a handover can hold the sending loop. The 30 below is + // proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced. + assert.Less(t, ServerAliveIntervalSeconds, 8*60) + assert.Greater(t, 3*ServerAliveIntervalSeconds, 30) +} + func TestGetConfigDir(t *testing.T) { dir, err := GetConfigDir(t.Context()) assert.NoError(t, err)