Keep idle SSH tunnel sessions alive with a websocket ping

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

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,

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

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")

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.

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,38 @@ 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:
+					// 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 {
+						// 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.
+						log.Debugf(gCtx, "Sent websocket keepalive ping")
+					}
+				}
+			}
+		})
 		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)
 		})

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,14 +156,25 @@ func TestMaxClients(t *testing.T) {
 }
 
 func TestHandover(t *testing.T) {
+	t.Run("without keepalive", func(t *testing.T) {
+		runHandoverExchange(t, time.Hour)
+	})
+	// 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)
+	})
+}
+
+func runHandoverExchange(t *testing.T, keepaliveInterval time.Duration) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
 
 	handoverChan := make(chan time.Time)
 	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 {