reuse exec SSH connections per VM

internal/command/controller/run.go

@@ -35,7 +35,8 @@ var noExperimentalRPCV2 bool
 var experimentalPingInterval time.Duration
 var experimentalDisableDBCompression bool
 var workerOfflineTimeout time.Duration
-var execSessionExitTTL time.Duration
+var execSessionRetentionTTL time.Duration
+var execSSHConnectionKeepaliveInterval time.Duration
 var synthetic bool
 
 func newRunCommand() *cobra.Command {
@@ -88,8 +89,10 @@ func newRunCommand() *cobra.Command {
 		"duration (e.g. 60s or 5m30s) after which a worker is considered offline for the purposes "+
 			"of scheduling (no new VMs will be scheduled on such worker and already assigned VMs will be "+
 			"marked as failed)")
-	cmd.Flags().DurationVar(&execSessionExitTTL, "exec-session-exit-ttl", 10*time.Minute,
+	cmd.Flags().DurationVar(&execSessionRetentionTTL, "exec-session-retention-ttl", 10*time.Minute,
 		"duration to retain reconnectable exec session history after the command exits")
+	cmd.Flags().DurationVar(&execSSHConnectionKeepaliveInterval, "exec-ssh-connection-keepalive-interval", 30*time.Second,
+		"interval between SSH keepalive requests sent by the controller for shared exec connections")
 
 	// Hidden flags
 	cmd.Flags().BoolVar(&synthetic, "synthetic", false, "")
@@ -150,7 +153,8 @@ func runController(cmd *cobra.Command, args []string) (err error) {
 		controller.WithListenAddr(address),
 		controller.WithDataDir(dataDir),
 		controller.WithWorkerOfflineTimeout(workerOfflineTimeout),
-		controller.WithExecSessionExitTTL(execSessionExitTTL),
+		controller.WithExecSessionRetentionTTL(execSessionRetentionTTL),
+		controller.WithExecSSHConnectionKeepaliveInterval(execSSHConnectionKeepaliveInterval),
 		controller.WithLogger(logger),
 	}
 
@@ -215,6 +219,10 @@ func runController(cmd *cobra.Command, args []string) (err error) {
 		controllerOpts = append(controllerOpts, controller.WithDisableDBCompression())
 	}
 
+	if execSSHConnectionKeepaliveInterval < 5*time.Second {
+		return fmt.Errorf("--exec-ssh-connection-keepalive-interval's value cannot be less than 5 seconds")
+	}
+
 	controllerInstance, err := controller.New(controllerOpts...)
 	if err != nil {
 		return err

internal/controller/api_vms_exec.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"net"
 	"net/http"
 	"strconv"
 	"time"
@@ -190,8 +189,7 @@ func (controller *Controller) newSSHExecSession(
 	sessionContext, sessionContextCancel := context.WithCancel(context.Background())
 
 	type sshExecAttempt struct {
-		portForwardConn net.Conn
+		exec *sshexec.Exec
-		exec            *sshexec.Exec
 	}
 
 	attempt, err := retry.NewWithData[sshExecAttempt](
@@ -201,32 +199,38 @@ func (controller *Controller) newSSHExecSession(
 		retry.Attempts(0),
 		retry.LastErrorOnly(true),
 	).Do(func() (sshExecAttempt, error) {
-		portForwardConn, err := controller.portForwardConnection(
+		exec, err := controller.execSSHClients.newExec(vm.UID, sshexec.Options{
-			sessionContext,
-			waitContext,
-			vm.Worker,
-			vm.UID,
-			22,
-		)
-		if err != nil {
-			return sshExecAttempt{}, err
-		}
-
-		exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), sshexec.Options{
 			Interactive: spec.interactive,
 			TTY:         spec.tty,
 			Rows:        spec.rows,
 			Cols:        spec.cols,
+		}, func() (sshExecClient, error) {
+			portForwardConn, err := controller.portForwardConnection(
+				context.Background(),
+				waitContext,
+				vm.Worker,
+				vm.UID,
+				22,
+			)
+			if err != nil {
+				return nil, err
+			}
+
+			client, err := sshexec.NewClient(portForwardConn, vm.SSHUsername(), vm.SSHPassword())
+			if err != nil {
+				_ = portForwardConn.Close()
+
+				return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err)
+			}
+
+			return client, nil
 		})
 		if err != nil {
-			_ = portForwardConn.Close()
-
 			return sshExecAttempt{}, fmt.Errorf("failed to establish SSH connection to a VM: %w", err)
 		}
 
 		return sshExecAttempt{
-			portForwardConn: portForwardConn,
+			exec: exec,
-			exec:            exec,
 		}, nil
 	})
 	if err != nil {
@@ -242,9 +246,9 @@ func (controller *Controller) newSSHExecSession(
 		spec,
 		runCommand,
 		attempt.exec,
-		attempt.portForwardConn,
+		nil,
 		registry,
-		controller.execSessionExitTTL,
+		controller.execSessionRetentionTTL,
 		policy,
 	), nil
 }

internal/controller/controller.go

@@ -39,33 +39,35 @@ var (
 )
 
 type Controller struct {
-	dataDir              *DataDir
+	dataDir                            *DataDir
-	listenAddr           string
+	listenAddr                         string
-	apiPrefix            string
+	apiPrefix                          string
-	tlsConfig            *tls.Config
+	tlsConfig                          *tls.Config
-	listener             net.Listener
+	listener                           net.Listener
-	httpServer           *http.Server
+	httpServer                         *http.Server
-	insecureAuthDisabled bool
+	insecureAuthDisabled               bool
-	scheduler            *scheduler.Scheduler
+	scheduler                          *scheduler.Scheduler
-	store                storepkg.Store
+	store                              storepkg.Store
-	logger               *zap.SugaredLogger
+	logger                             *zap.SugaredLogger
-	grpcServer           *grpc.Server
+	grpcServer                         *grpc.Server
-	workerNotifier       *notifier.Notifier
+	workerNotifier                     *notifier.Notifier
-	connRendezvous       *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]]
+	connRendezvous                     *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]]
-	ipRendezvous         *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]]
+	ipRendezvous                       *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]]
-	enableSwaggerDocs    bool
+	enableSwaggerDocs                  bool
-	workerOfflineTimeout time.Duration
+	workerOfflineTimeout               time.Duration
-	execSessionExitTTL   time.Duration
+	execSessionRetentionTTL            time.Duration
-	experimentalRPCV2    bool
+	execSSHConnectionKeepaliveInterval time.Duration
-	disableDBCompression bool
+	experimentalRPCV2                  bool
-	pingInterval         time.Duration
+	disableDBCompression               bool
-	synthetic            bool
+	pingInterval                       time.Duration
+	synthetic                          bool
 
 	sshListenAddr   string
 	sshSigner       ssh.Signer
 	sshNoClientAuth bool
 	sshServer       *sshserver.SSHServer
 	execSessions    *execSessionRegistry
+	execSSHClients  *execSSHClientPool
 
 	single singleflight.Group
 
@@ -74,13 +76,14 @@ type Controller struct {
 
 func New(opts ...Option) (*Controller, error) {
 	controller := &Controller{
-		connRendezvous:       rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](),
+		connRendezvous:                     rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](),
-		ipRendezvous:         rendezvous.New[rendezvous.ResultWithErrorMessage[string]](),
+		ipRendezvous:                       rendezvous.New[rendezvous.ResultWithErrorMessage[string]](),
-		workerOfflineTimeout: 3 * time.Minute,
+		workerOfflineTimeout:               3 * time.Minute,
-		execSessionExitTTL:   10 * time.Minute,
+		execSessionRetentionTTL:            10 * time.Minute,
-		pingInterval:         30 * time.Second,
+		execSSHConnectionKeepaliveInterval: 30 * time.Second,
-		execSessions:         newExecSessionRegistry(),
+		pingInterval:                       30 * time.Second,
-		single:               singleflight.Group{},
+		execSessions:                       newExecSessionRegistry(),
+		single:                             singleflight.Group{},
 	}
 
 	// Apply options
@@ -99,6 +102,10 @@ func New(opts ...Option) (*Controller, error) {
 	if controller.logger == nil {
 		controller.logger = zap.NewNop().Sugar()
 	}
+	controller.execSSHClients = newExecSSHClientPool(
+		controller.execSSHConnectionKeepaliveInterval,
+		controller.logger.With("component", "exec-ssh"),
+	)
 
 	// Instantiate the database
 	store, err := badger.NewBadgerStore(controller.dataDir.DBPath(), controller.disableDBCompression,
@@ -313,6 +320,7 @@ func (controller *Controller) Run(ctx context.Context) error {
 		<-ctx.Done()
 
 		controller.execSessions.closeAll()
+		controller.execSSHClients.closeAll()
 
 		if err := controller.httpServer.Shutdown(ctx); err != nil {
 			controller.logger.Errorf("failed to cleanly shutdown the HTTP server: %v", err)

internal/controller/exec_sessions.go

@@ -325,14 +325,14 @@ func (subscriber *execSessionSubscriber) close() {
 }
 
 type execSession struct {
-	key       execSessionKey
+	key          execSessionKey
-	spec      execSessionSpec
+	spec         execSessionSpec
-	command   string
+	command      string
-	exec      sshExecRunner
+	exec         sshExecRunner
-	transport net.Conn
+	transport    net.Conn
-	registry  *execSessionRegistry
+	registry     *execSessionRegistry
-	exitTTL   time.Duration
+	retentionTTL time.Duration
-	policy    execSessionPolicy
+	policy       execSessionPolicy
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -358,7 +358,7 @@ func newExecSession(
 	exec sshExecRunner,
 	transport net.Conn,
 	registry *execSessionRegistry,
-	exitTTL time.Duration,
+	retentionTTL time.Duration,
 	policy execSessionPolicy,
 ) *execSession {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -372,7 +372,7 @@ func newExecSession(
 		exec,
 		transport,
 		registry,
-		exitTTL,
+		retentionTTL,
 		policy,
 	)
 }
@@ -386,7 +386,7 @@ func newExecSessionWithContextAndSpec(
 	exec sshExecRunner,
 	transport net.Conn,
 	registry *execSessionRegistry,
-	exitTTL time.Duration,
+	retentionTTL time.Duration,
 	policy execSessionPolicy,
 ) *execSession {
 	if ctx == nil || cancel == nil {
@@ -394,19 +394,19 @@ func newExecSessionWithContextAndSpec(
 	}
 
 	session := &execSession{
-		key:         key,
+		key:          key,
-		spec:        spec.clone(),
+		spec:         spec.clone(),
-		command:     command,
+		command:      command,
-		exec:        exec,
+		exec:         exec,
-		transport:   transport,
+		transport:    transport,
-		registry:    registry,
+		registry:     registry,
-		exitTTL:     exitTTL,
+		retentionTTL: retentionTTL,
-		policy:      policy,
+		policy:       policy,
-		ctx:         ctx,
+		ctx:          ctx,
-		cancel:      cancel,
+		cancel:       cancel,
-		stdin:       exec.Stdin(),
+		stdin:        exec.Stdin(),
-		subscribers: map[*execSessionSubscriber]struct{}{},
+		subscribers:  map[*execSessionSubscriber]struct{}{},
-		done:        make(chan struct{}),
+		done:         make(chan struct{}),
 	}
 
 	return session
@@ -646,7 +646,7 @@ func (session *execSession) markFinished() {
 	session.finished = true
 	shouldClose := !session.policy.retainAfterExit
 	if !session.closed && session.policy.retainAfterExit {
-		session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire)
+		session.expiryTimer = time.AfterFunc(session.retentionTTL, session.expire)
 	}
 
 	var subscribers []*execSessionSubscriber

internal/controller/exec_sessions_test.go

@@ -55,17 +55,17 @@ func newManualExecSessionForTest(
 	ctx, cancel := context.WithCancel(context.Background())
 
 	return &execSession{
-		key:         key,
+		key:          key,
-		spec:        execSessionSpec{command: "echo test"},
+		spec:         execSessionSpec{command: "echo test"},
-		command:     "echo test",
+		command:      "echo test",
-		exec:        &fakeExec{},
+		exec:         &fakeExec{},
-		registry:    registry,
+		registry:     registry,
-		exitTTL:     time.Minute,
+		retentionTTL: time.Minute,
-		policy:      reconnectableExecSessionPolicy,
+		policy:       reconnectableExecSessionPolicy,
-		ctx:         ctx,
+		ctx:          ctx,
-		cancel:      cancel,
+		cancel:       cancel,
-		subscribers: map[*execSessionSubscriber]struct{}{},
+		subscribers:  map[*execSessionSubscriber]struct{}{},
-		done:        make(chan struct{}),
+		done:         make(chan struct{}),
 	}
 }
 
@@ -350,7 +350,7 @@ func TestExecSessionFinishedEntryExpiresAfterTTL(t *testing.T) {
 	registry := newExecSessionRegistry()
 	key := execSessionKey{vmName: "vm", sessionID: "session"}
 	session := newManualExecSessionForTest(key, registry)
-	session.exitTTL = 10 * time.Millisecond
+	session.retentionTTL = 10 * time.Millisecond
 	registry.sessions[key] = session
 
 	session.markFinished()

internal/controller/exec_ssh_clients.go

@@ -0,0 +1,202 @@
+package controller
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/cirruslabs/orchard/internal/controller/sshexec"
+	"go.uber.org/zap"
+)
+
+var errExecSSHClientPoolClosed = errors.New("SSH exec client pool is closed")
+
+type execSSHClientPool struct {
+	// clients maps VM UID strings to *pooledExecSSHClient values so each VM has
+	// at most one shared SSH connection.
+	clients sync.Map
+
+	closed atomic.Bool
+
+	keepaliveInterval time.Duration
+	logger            *zap.SugaredLogger
+}
+
+func newExecSSHClientPool(keepaliveInterval time.Duration, logger *zap.SugaredLogger) *execSSHClientPool {
+	return &execSSHClientPool{
+		keepaliveInterval: keepaliveInterval,
+		logger:            logger,
+	}
+}
+
+func (pool *execSSHClientPool) newExec(
+	vmUID string,
+	options sshexec.Options,
+	create func() (sshExecClient, error),
+) (*sshexec.Exec, error) {
+	if pool.closed.Load() {
+		return nil, errExecSSHClientPoolClosed
+	}
+
+	return pool.client(vmUID).newExec(func() (sshExecClient, error) {
+		initializedClient, err := create()
+		if err != nil {
+			return nil, err
+		}
+		if pool.closed.Load() {
+			_ = initializedClient.Close()
+
+			return nil, errExecSSHClientPoolClosed
+		}
+
+		return initializedClient, nil
+	}, options)
+}
+
+func (pool *execSSHClientPool) closeAll() {
+	pool.closed.Store(true)
+
+	pool.clients.Range(func(key any, value any) bool {
+		pool.clients.Delete(key)
+		value.(*pooledExecSSHClient).close()
+
+		return true
+	})
+}
+
+func (pool *execSSHClientPool) client(vmUID string) *pooledExecSSHClient {
+	client, _ := pool.clients.LoadOrStore(vmUID, &pooledExecSSHClient{
+		vmUID:             vmUID,
+		keepaliveInterval: pool.keepaliveInterval,
+		logger:            pool.logger,
+	})
+
+	return client.(*pooledExecSSHClient)
+}
+
+type pooledExecSSHClient struct {
+	vmUID string
+
+	mu      sync.Mutex
+	current sshExecClient
+
+	keepaliveInterval time.Duration
+	logger            *zap.SugaredLogger
+}
+
+func (client *pooledExecSSHClient) newExec(
+	create func() (sshExecClient, error),
+	options sshexec.Options,
+) (*sshexec.Exec, error) {
+	current, err := client.getOrInit(create)
+	if err != nil {
+		return nil, err
+	}
+
+	exec, err := current.NewExec(options)
+	if err != nil && current.ShouldRecreateAfter(err) {
+		client.invalidate(current)
+	}
+
+	return exec, err
+}
+
+func (client *pooledExecSSHClient) getOrInit(
+	create func() (sshExecClient, error),
+) (sshExecClient, error) {
+	client.mu.Lock()
+	defer client.mu.Unlock()
+
+	if client.current != nil {
+		return client.current, nil
+	}
+
+	initializedClient, err := create()
+	if err != nil {
+		return nil, err
+	}
+
+	client.current = initializedClient
+	client.monitor(initializedClient)
+
+	return initializedClient, nil
+}
+
+func (client *pooledExecSSHClient) invalidate(expected sshExecClient) bool {
+	if !client.clear(expected) {
+		return false
+	}
+
+	_ = expected.Close()
+
+	return true
+}
+
+func (client *pooledExecSSHClient) clear(expected sshExecClient) bool {
+	client.mu.Lock()
+	defer client.mu.Unlock()
+
+	if client.current != expected {
+		return false
+	}
+
+	client.current = nil
+
+	return true
+}
+
+func (client *pooledExecSSHClient) close() {
+	client.mu.Lock()
+	current := client.current
+	client.current = nil
+	client.mu.Unlock()
+
+	if current == nil {
+		return
+	}
+
+	_ = current.Close()
+}
+
+func (client *pooledExecSSHClient) monitor(current sshExecClient) {
+	go func() {
+		var keepalive <-chan time.Time
+		if client.keepaliveInterval > 0 {
+			ticker := time.NewTicker(client.keepaliveInterval)
+			defer ticker.Stop()
+
+			keepalive = ticker.C
+		}
+
+		for {
+			select {
+			case <-current.Done():
+				if client.clear(current) {
+					client.logger.Debugf("evicted disconnected SSH exec client for VM UID %s: %v",
+						client.vmUID, current.Err())
+				}
+
+				return
+			case <-keepalive:
+				if err := current.Keepalive(); err != nil {
+					if client.invalidate(current) {
+						client.logger.Debugf("evicted SSH exec client for VM UID %s after keepalive failure: %v",
+							client.vmUID, err)
+					}
+
+					return
+				}
+			}
+		}
+	}()
+}
+
+type sshExecClient interface {
+	NewExec(options sshexec.Options) (*sshexec.Exec, error)
+	Keepalive() error
+	Done() <-chan struct{}
+	Err() error
+	Close() error
+	ShouldRecreateAfter(err error) bool
+}

internal/controller/exec_ssh_clients_test.go

@@ -0,0 +1,170 @@
+package controller
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/cirruslabs/orchard/internal/controller/sshexec"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+type fakeSSHExecClient struct {
+	keepaliveErr error
+
+	closeCalls     atomic.Int32
+	keepaliveCalls atomic.Int32
+
+	done      chan struct{}
+	closeOnce sync.Once
+}
+
+func newFakeSSHExecClient() *fakeSSHExecClient {
+	return &fakeSSHExecClient{
+		done: make(chan struct{}),
+	}
+}
+
+func (client *fakeSSHExecClient) NewExec(options sshexec.Options) (*sshexec.Exec, error) {
+	return nil, nil
+}
+
+func (client *fakeSSHExecClient) Keepalive() error {
+	client.keepaliveCalls.Add(1)
+
+	return client.keepaliveErr
+}
+
+func (client *fakeSSHExecClient) Done() <-chan struct{} {
+	return client.done
+}
+
+func (client *fakeSSHExecClient) Err() error {
+	<-client.done
+
+	return errors.New("disconnected")
+}
+
+func (client *fakeSSHExecClient) Close() error {
+	client.closeCalls.Add(1)
+	client.closeOnce.Do(func() {
+		close(client.done)
+	})
+
+	return nil
+}
+
+func (client *fakeSSHExecClient) ShouldRecreateAfter(err error) bool {
+	return true
+}
+
+func TestExecSSHClientPoolDeduplicatesConcurrentInitialization(t *testing.T) {
+	pool := newExecSSHClientPool(0, zap.NewNop().Sugar())
+	client := newFakeSSHExecClient()
+
+	createStarted := make(chan struct{})
+	releaseCreate := make(chan struct{})
+
+	var createCalls atomic.Int32
+	create := func() (sshExecClient, error) {
+		if createCalls.Add(1) == 1 {
+			close(createStarted)
+		}
+
+		<-releaseCreate
+
+		return client, nil
+	}
+
+	const callers = 16
+
+	start := make(chan struct{})
+	errCh := make(chan error, callers)
+
+	var wg sync.WaitGroup
+	wg.Add(callers)
+
+	for range callers {
+		go func() {
+			defer wg.Done()
+
+			<-start
+
+			_, err := pool.newExec("vm-1", sshexec.Options{}, create)
+			if err != nil {
+				errCh <- err
+			}
+		}()
+	}
+
+	close(start)
+	<-createStarted
+	close(releaseCreate)
+
+	wg.Wait()
+	close(errCh)
+
+	for err := range errCh {
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 1, createCalls.Load())
+
+	pool.closeAll()
+}
+
+func TestExecSSHClientPoolKeepaliveInvalidatesClient(t *testing.T) {
+	pool := newExecSSHClientPool(time.Millisecond, zap.NewNop().Sugar())
+	client := newFakeSSHExecClient()
+	client.keepaliveErr = errors.New("boom")
+
+	_, err := pool.newExec("vm-1", sshexec.Options{}, func() (sshExecClient, error) {
+		return client, nil
+	})
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		return client.keepaliveCalls.Load() > 0 && client.closeCalls.Load() == 1
+	}, time.Second, time.Millisecond)
+
+	replacement := newFakeSSHExecClient()
+	var createCalls atomic.Int32
+	_, err = pool.newExec("vm-1", sshexec.Options{}, func() (sshExecClient, error) {
+		createCalls.Add(1)
+
+		return replacement, nil
+	})
+	require.NoError(t, err)
+	require.EqualValues(t, 1, createCalls.Load())
+
+	pool.closeAll()
+}
+
+func TestExecSSHClientPoolWaitClearsDisconnectedClient(t *testing.T) {
+	pool := newExecSSHClientPool(0, zap.NewNop().Sugar())
+	client := newFakeSSHExecClient()
+
+	_, err := pool.newExec("vm-1", sshexec.Options{}, func() (sshExecClient, error) {
+		return client, nil
+	})
+	require.NoError(t, err)
+
+	require.NoError(t, client.Close())
+
+	replacement := newFakeSSHExecClient()
+	var createCalls atomic.Int32
+	require.Eventually(t, func() bool {
+		_, err := pool.newExec("vm-1", sshexec.Options{}, func() (sshExecClient, error) {
+			createCalls.Add(1)
+
+			return replacement, nil
+		})
+		require.NoError(t, err)
+
+		return createCalls.Load() == 1
+	}, time.Second, time.Millisecond)
+
+	pool.closeAll()
+}

internal/controller/option.go

@@ -60,9 +60,15 @@ func WithWorkerOfflineTimeout(workerOfflineTimeout time.Duration) Option {
 	}
 }
 
-func WithExecSessionExitTTL(execSessionExitTTL time.Duration) Option {
+func WithExecSessionRetentionTTL(execSessionRetentionTTL time.Duration) Option {
 	return func(controller *Controller) {
-		controller.execSessionExitTTL = execSessionExitTTL
+		controller.execSessionRetentionTTL = execSessionRetentionTTL
+	}
+}
+
+func WithExecSSHConnectionKeepaliveInterval(execSSHConnectionKeepaliveInterval time.Duration) Option {
+	return func(controller *Controller) {
+		controller.execSSHConnectionKeepaliveInterval = execSSHConnectionKeepaliveInterval
 	}
 }
 

internal/controller/sshexec/sshexec.go

@@ -10,6 +10,7 @@ import (
 	"slices"
 	"sort"
 	"strings"
+	"sync"
 
 	"github.com/cirruslabs/orchard/internal/execstream"
 	"golang.org/x/crypto/ssh"
@@ -27,8 +28,15 @@ type Options struct {
 	Workdir     string
 }
 
+type Client struct {
+	sshClient *ssh.Client
+	done      chan struct{}
+	waitErr   error
+	waitMu    sync.Mutex
+}
+
 type Exec struct {
-	sshClient   *ssh.Client
+	ownedClient *Client
 	sshSession  *ssh.Session
 	stdout      io.Reader
 	stderr      io.Reader
@@ -37,7 +45,7 @@ type Exec struct {
 	tty         bool
 }
 
-func New(netConn net.Conn, user string, password string, options Options) (*Exec, error) {
+func NewClient(netConn net.Conn, user string, password string) (*Client, error) {
 	// Establish an SSH connection
 	sshConn, sshChans, sshReqs, err := ssh.NewClientConn(netConn, "", &ssh.ClientConfig{
 		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
@@ -52,18 +60,50 @@ func New(netConn net.Conn, user string, password string, options Options) (*Exec
 		return nil, fmt.Errorf("failed to create an SSH connection: %w", err)
 	}
 
-	sshClient := ssh.NewClient(sshConn, sshChans, sshReqs)
+	client := &Client{
+		sshClient: ssh.NewClient(sshConn, sshChans, sshReqs),
+		done:      make(chan struct{}),
+	}
 
-	// Create a new SSH session
+	go func() {
-	sshSession, err := sshClient.NewSession()
+		err := client.sshClient.Wait()
+
+		client.waitMu.Lock()
+		client.waitErr = err
+		client.waitMu.Unlock()
+
+		close(client.done)
+	}()
+
+	return client, nil
+}
+
+func New(netConn net.Conn, user string, password string, options Options) (*Exec, error) {
+	client, err := NewClient(netConn, user, password)
 	if err != nil {
-		_ = sshClient.Close()
+		return nil, err
+	}
 
+	exec, err := client.NewExec(options)
+	if err != nil {
+		_ = client.Close()
+
+		return nil, err
+	}
+
+	exec.ownedClient = client
+
+	return exec, nil
+}
+
+func (client *Client) NewExec(options Options) (*Exec, error) {
+	// Create a new SSH session
+	sshSession, err := client.sshClient.NewSession()
+	if err != nil {
 		return nil, fmt.Errorf("failed to create an SSH session: %w", err)
 	}
 
 	exec := &Exec{
-		sshClient:  sshClient,
 		sshSession: sshSession,
 		tty:        options.TTY,
 	}
@@ -83,7 +123,6 @@ func New(netConn net.Conn, user string, password string, options Options) (*Exec
 			ssh.TerminalModes{},
 		); err != nil {
 			_ = sshSession.Close()
-			_ = sshClient.Close()
 
 			return nil, fmt.Errorf("failed to request PTY for the SSH session: %w", err)
 		}
@@ -92,7 +131,6 @@ func New(netConn net.Conn, user string, password string, options Options) (*Exec
 	exec.stdout, err = sshSession.StdoutPipe()
 	if err != nil {
 		_ = sshSession.Close()
-		_ = sshClient.Close()
 
 		return nil, fmt.Errorf("failed to create standard output pipe "+
 			"for the SSH session: %w", err)
@@ -101,7 +139,6 @@ func New(netConn net.Conn, user string, password string, options Options) (*Exec
 	exec.stderr, err = sshSession.StderrPipe()
 	if err != nil {
 		_ = sshSession.Close()
-		_ = sshClient.Close()
 
 		return nil, fmt.Errorf("failed to create standard error pipe "+
 			"for the SSH session: %w", err)
@@ -110,6 +147,41 @@ func New(netConn net.Conn, user string, password string, options Options) (*Exec
 	return exec, nil
 }
 
+func (client *Client) Keepalive() error {
+	_, _, err := client.sshClient.SendRequest("keepalive@openssh.com", true, nil)
+
+	return err
+}
+
+func (client *Client) Done() <-chan struct{} {
+	return client.done
+}
+
+func (client *Client) Err() error {
+	<-client.done
+
+	client.waitMu.Lock()
+	defer client.waitMu.Unlock()
+
+	return client.waitErr
+}
+
+func (client *Client) Close() error {
+	return client.sshClient.Close()
+}
+
+func (client *Client) ShouldRecreateAfter(err error) bool {
+	select {
+	case <-client.done:
+		return true
+	default:
+	}
+
+	return errors.Is(err, io.EOF) ||
+		errors.Is(err, net.ErrClosed) ||
+		strings.Contains(err.Error(), "use of closed network connection")
+}
+
 func (exec *Exec) Stdin() io.WriteCloser {
 	return exec.stdin
 }
@@ -286,10 +358,16 @@ func (exec *Exec) Close() error {
 	}
 
 	if err := exec.sshSession.Close(); err != nil {
-		_ = exec.sshClient.Close()
+		if exec.ownedClient != nil {
+			_ = exec.ownedClient.Close()
+		}
 
 		return err
 	}
 
-	return exec.sshClient.Close()
+	if exec.ownedClient != nil {
+		return exec.ownedClient.Close()
+	}
+
+	return nil
 }

internal/tests/exec_ssh_server_test.go

@@ -17,6 +17,12 @@ type execSSHServer struct {
 	config   *ssh.ServerConfig
 
 	rejectFirstConnections atomic.Int32
+	rejectFirstSessions    atomic.Int32
+	successfulConnections  atomic.Int32
+	keepaliveRequests      atomic.Int32
+
+	mu    sync.Mutex
+	conns map[net.Conn]struct{}
 
 	wg sync.WaitGroup
 }
@@ -35,6 +41,7 @@ func startExecSSHServer(t *testing.T, rejectFirstConnections int32) *execSSHServ
 
 	server := &execSSHServer{
 		listener: listener,
+		conns:    map[net.Conn]struct{}{},
 		config: &ssh.ServerConfig{
 			PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
 				if conn.User() != "admin" || string(password) != "admin" {
@@ -63,6 +70,31 @@ func (server *execSSHServer) Addr() string {
 	return server.listener.Addr().String()
 }
 
+func (server *execSSHServer) SuccessfulConnections() int32 {
+	return server.successfulConnections.Load()
+}
+
+func (server *execSSHServer) KeepaliveRequests() int32 {
+	return server.keepaliveRequests.Load()
+}
+
+func (server *execSSHServer) RejectNextSessions(count int32) {
+	server.rejectFirstSessions.Store(count)
+}
+
+func (server *execSSHServer) CloseClientConnections() {
+	server.mu.Lock()
+	conns := make([]net.Conn, 0, len(server.conns))
+	for conn := range server.conns {
+		conns = append(conns, conn)
+	}
+	server.mu.Unlock()
+
+	for _, conn := range conns {
+		_ = conn.Close()
+	}
+}
+
 func (server *execSSHServer) run() {
 	defer server.wg.Done()
 
@@ -88,6 +120,15 @@ func (server *execSSHServer) run() {
 }
 
 func (server *execSSHServer) serve(conn net.Conn) {
+	server.mu.Lock()
+	server.conns[conn] = struct{}{}
+	server.mu.Unlock()
+
+	defer func() {
+		server.mu.Lock()
+		delete(server.conns, conn)
+		server.mu.Unlock()
+	}()
 	defer conn.Close()
 
 	serverConn, newChannels, requests, err := ssh.NewServerConn(conn, server.config)
@@ -96,7 +137,9 @@ func (server *execSSHServer) serve(conn net.Conn) {
 	}
 	defer serverConn.Close()
 
-	go ssh.DiscardRequests(requests)
+	server.successfulConnections.Add(1)
+
+	go server.serveGlobalRequests(requests)
 
 	for newChannel := range newChannels {
 		if newChannel.ChannelType() != "session" {
@@ -104,6 +147,11 @@ func (server *execSSHServer) serve(conn net.Conn) {
 
 			continue
 		}
+		if server.rejectFirstSessions.Add(-1) >= 0 {
+			_ = newChannel.Reject(ssh.Prohibited, "session rejected for test")
+
+			continue
+		}
 
 		channel, requests, err := newChannel.Accept()
 		if err != nil {
@@ -119,6 +167,17 @@ func (server *execSSHServer) serve(conn net.Conn) {
 	}
 }
 
+func (server *execSSHServer) serveGlobalRequests(requests <-chan *ssh.Request) {
+	for request := range requests {
+		if request.Type == "keepalive@openssh.com" {
+			server.keepaliveRequests.Add(1)
+		}
+		if request.WantReply {
+			_ = request.Reply(false, nil)
+		}
+	}
+}
+
 func serveExecSSHSession(channel ssh.Channel, requests <-chan *ssh.Request) {
 	defer channel.Close()
 

internal/tests/exec_test.go

@@ -228,7 +228,7 @@ func TestVMExecScript(t *testing.T) {
 }
 
 func TestVMExecManyConcurrentSessions(t *testing.T) {
-	sshServer := startExecSSHServer(t, 24)
+	sshServer := startExecSSHServer(t, 2)
 
 	devClient, vmName := prepareForSyntheticExec(t, dialer.DialFunc(
 		func(ctx context.Context, network string, addr string) (net.Conn, error) {
@@ -307,6 +307,62 @@ func TestVMExecManyConcurrentSessions(t *testing.T) {
 	for err := range errCh {
 		require.NoError(t, err)
 	}
+
+	require.EqualValues(t, 1, sshServer.SuccessfulConnections())
+}
+
+func TestVMExecRecreatesSharedSSHClientAfterDisconnect(t *testing.T) {
+	sshServer := startExecSSHServer(t, 0)
+
+	devClient, vmName := prepareForSyntheticExec(t, dialer.DialFunc(
+		func(ctx context.Context, network string, addr string) (net.Conn, error) {
+			var netDialer net.Dialer
+
+			return netDialer.DialContext(ctx, network, sshServer.Addr())
+		},
+	))
+
+	runSyntheticExec(t, devClient, vmName)
+	require.EqualValues(t, 1, sshServer.SuccessfulConnections())
+
+	sshServer.CloseClientConnections()
+
+	runSyntheticExec(t, devClient, vmName)
+	require.EqualValues(t, 2, sshServer.SuccessfulConnections())
+}
+
+func TestVMExecSharedSSHClientSendsKeepalives(t *testing.T) {
+	sshServer := startExecSSHServer(t, 0)
+
+	devClient, vmName := prepareForSyntheticExec(t, dialer.DialFunc(
+		func(ctx context.Context, network string, addr string) (net.Conn, error) {
+			var netDialer net.Dialer
+
+			return netDialer.DialContext(ctx, network, sshServer.Addr())
+		},
+	), controller.WithExecSSHConnectionKeepaliveInterval(10*time.Millisecond))
+
+	runSyntheticExec(t, devClient, vmName)
+
+	require.Eventually(t, func() bool {
+		return sshServer.KeepaliveRequests() > 0
+	}, time.Second, 10*time.Millisecond)
+}
+
+func TestVMExecKeepsSharedSSHClientAfterSessionRejection(t *testing.T) {
+	sshServer := startExecSSHServer(t, 0)
+	sshServer.RejectNextSessions(1)
+
+	devClient, vmName := prepareForSyntheticExec(t, dialer.DialFunc(
+		func(ctx context.Context, network string, addr string) (net.Conn, error) {
+			var netDialer net.Dialer
+
+			return netDialer.DialContext(ctx, network, sshServer.Addr())
+		},
+	))
+
+	runSyntheticExec(t, devClient, vmName)
+	require.EqualValues(t, 1, sshServer.SuccessfulConnections())
 }
 
 func TestVMExecSessionReconnectHistory(t *testing.T) {
@@ -499,9 +555,15 @@ func prepareForExec(t *testing.T) (*client.Client, string) {
 	return devClient, vmName
 }
 
-func prepareForSyntheticExec(t *testing.T, vmDialer dialer.Dialer) (*client.Client, string) {
+func prepareForSyntheticExec(
+	t *testing.T,
+	vmDialer dialer.Dialer,
+	additionalControllerOpts ...controller.Option,
+) (*client.Client, string) {
+	controllerOpts := append([]controller.Option{controller.WithSynthetic()}, additionalControllerOpts...)
+
 	devClient, _, _ := devcontroller.StartIntegrationTestEnvironmentWithAdditionalOpts(t,
-		false, []controller.Option{controller.WithSynthetic()},
+		false, controllerOpts,
 		false, []worker.Option{
 			worker.WithSynthetic(),
 			worker.WithDialer(vmDialer),
@@ -525,6 +587,22 @@ func prepareForSyntheticExec(t *testing.T, vmDialer dialer.Dialer) (*client.Clie
 	return devClient, vmName
 }
 
+func runSyntheticExec(t *testing.T, devClient *client.Client, vmName string) {
+	t.Helper()
+
+	wsConn, err := devClient.VMs().Exec(t.Context(), vmName, "echo ignored", false, 30)
+	require.NoError(t, err)
+	defer wsConn.CloseNow()
+
+	frame := readFrame(t, wsConn)
+	require.Equal(t, execstream.FrameTypeStdout, frame.Type)
+	require.Equal(t, "ok", string(frame.Data))
+
+	frame = readFrame(t, wsConn)
+	require.Equal(t, execstream.FrameTypeExit, frame.Type)
+	require.EqualValues(t, 0, frame.Exit.Code)
+}
+
 func readFrame(t *testing.T, wsConn *websocket.Conn) *execstream.Frame {
 	t.Helper()