From b30bb5bedf23791bf7237138c79a1f458fb94a9b Mon Sep 17 00:00:00 2001 From: Fedor Kororkov Date: Mon, 11 May 2026 15:56:46 -0400 Subject: [PATCH] Reuse exec SSH connections per VM (#440) --- .cirrus.yml | 2 +- internal/command/controller/run.go | 14 +- internal/controller/api_vms_exec.go | 46 +++-- internal/controller/controller.go | 64 +++--- internal/controller/exec_sessions.go | 50 ++--- internal/controller/exec_sessions_test.go | 24 +-- internal/controller/exec_ssh_clients.go | 202 +++++++++++++++++++ internal/controller/exec_ssh_clients_test.go | 170 ++++++++++++++++ internal/controller/option.go | 10 +- internal/controller/sshexec/sshexec.go | 102 ++++++++-- internal/tests/exec_ssh_server_test.go | 61 +++++- internal/tests/exec_test.go | 84 +++++++- 12 files changed, 721 insertions(+), 108 deletions(-) create mode 100644 internal/controller/exec_ssh_clients.go create mode 100644 internal/controller/exec_ssh_clients_test.go diff --git a/.cirrus.yml b/.cirrus.yml index 4a9f1d3..ab467cf 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -8,7 +8,7 @@ task: pre_pull_script: - tart pull ghcr.io/cirruslabs/macos-tahoe-base:latest test_script: - - go test -ldflags="-B gobuildid" -v -count=1 ./... + - go test -timeout=20m -ldflags="-B gobuildid" -v -count=1 ./... always: cleanup_script: - tart list diff --git a/internal/command/controller/run.go b/internal/command/controller/run.go index 8833f93..b8cfff0 100644 --- a/internal/command/controller/run.go +++ b/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 diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go index cf684d7..a3ebf0d 100644 --- a/internal/controller/api_vms_exec.go +++ b/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( - sessionContext, - waitContext, - vm.Worker, - vm.UID, - 22, - ) - if err != nil { - return sshExecAttempt{}, err - } - - exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), sshexec.Options{ + exec, err := controller.execSSHClients.newExec(vm.UID, 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 } diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 1521e0d..70090d8 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -39,33 +39,35 @@ var ( ) type Controller struct { - dataDir *DataDir - listenAddr string - apiPrefix string - tlsConfig *tls.Config - listener net.Listener - httpServer *http.Server - insecureAuthDisabled bool - scheduler *scheduler.Scheduler - store storepkg.Store - logger *zap.SugaredLogger - grpcServer *grpc.Server - workerNotifier *notifier.Notifier - connRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]] - ipRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]] - enableSwaggerDocs bool - workerOfflineTimeout time.Duration - execSessionExitTTL time.Duration - experimentalRPCV2 bool - disableDBCompression bool - pingInterval time.Duration - synthetic bool + dataDir *DataDir + listenAddr string + apiPrefix string + tlsConfig *tls.Config + listener net.Listener + httpServer *http.Server + insecureAuthDisabled bool + scheduler *scheduler.Scheduler + store storepkg.Store + logger *zap.SugaredLogger + grpcServer *grpc.Server + workerNotifier *notifier.Notifier + connRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]] + ipRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]] + enableSwaggerDocs bool + workerOfflineTimeout time.Duration + execSessionRetentionTTL time.Duration + execSSHConnectionKeepaliveInterval time.Duration + experimentalRPCV2 bool + disableDBCompression 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]](), - ipRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[string]](), - workerOfflineTimeout: 3 * time.Minute, - execSessionExitTTL: 10 * time.Minute, - pingInterval: 30 * time.Second, - execSessions: newExecSessionRegistry(), - single: singleflight.Group{}, + connRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](), + ipRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[string]](), + workerOfflineTimeout: 3 * time.Minute, + execSessionRetentionTTL: 10 * time.Minute, + execSSHConnectionKeepaliveInterval: 30 * time.Second, + pingInterval: 30 * time.Second, + 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) diff --git a/internal/controller/exec_sessions.go b/internal/controller/exec_sessions.go index 13a6212..8ef0f31 100644 --- a/internal/controller/exec_sessions.go +++ b/internal/controller/exec_sessions.go @@ -325,14 +325,14 @@ func (subscriber *execSessionSubscriber) close() { } type execSession struct { - key execSessionKey - spec execSessionSpec - command string - exec sshExecRunner - transport net.Conn - registry *execSessionRegistry - exitTTL time.Duration - policy execSessionPolicy + key execSessionKey + spec execSessionSpec + command string + exec sshExecRunner + transport net.Conn + registry *execSessionRegistry + retentionTTL time.Duration + 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, - spec: spec.clone(), - command: command, - exec: exec, - transport: transport, - registry: registry, - exitTTL: exitTTL, - policy: policy, - ctx: ctx, - cancel: cancel, - stdin: exec.Stdin(), - subscribers: map[*execSessionSubscriber]struct{}{}, - done: make(chan struct{}), + key: key, + spec: spec.clone(), + command: command, + exec: exec, + transport: transport, + registry: registry, + retentionTTL: retentionTTL, + policy: policy, + ctx: ctx, + cancel: cancel, + stdin: exec.Stdin(), + subscribers: map[*execSessionSubscriber]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 diff --git a/internal/controller/exec_sessions_test.go b/internal/controller/exec_sessions_test.go index 7164635..e569def 100644 --- a/internal/controller/exec_sessions_test.go +++ b/internal/controller/exec_sessions_test.go @@ -55,17 +55,17 @@ func newManualExecSessionForTest( ctx, cancel := context.WithCancel(context.Background()) return &execSession{ - key: key, - spec: execSessionSpec{command: "echo test"}, - command: "echo test", - exec: &fakeExec{}, - registry: registry, - exitTTL: time.Minute, - policy: reconnectableExecSessionPolicy, - ctx: ctx, - cancel: cancel, - subscribers: map[*execSessionSubscriber]struct{}{}, - done: make(chan struct{}), + key: key, + spec: execSessionSpec{command: "echo test"}, + command: "echo test", + exec: &fakeExec{}, + registry: registry, + retentionTTL: time.Minute, + policy: reconnectableExecSessionPolicy, + ctx: ctx, + cancel: cancel, + subscribers: map[*execSessionSubscriber]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() diff --git a/internal/controller/exec_ssh_clients.go b/internal/controller/exec_ssh_clients.go new file mode 100644 index 0000000..b89b6cb --- /dev/null +++ b/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 +} diff --git a/internal/controller/exec_ssh_clients_test.go b/internal/controller/exec_ssh_clients_test.go new file mode 100644 index 0000000..8cd9987 --- /dev/null +++ b/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() +} diff --git a/internal/controller/option.go b/internal/controller/option.go index 325c912..adda403 100644 --- a/internal/controller/option.go +++ b/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 } } diff --git a/internal/controller/sshexec/sshexec.go b/internal/controller/sshexec/sshexec.go index 91cc8ab..2f8c56e 100644 --- a/internal/controller/sshexec/sshexec.go +++ b/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 - sshSession, err := sshClient.NewSession() + go func() { + 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 } diff --git a/internal/tests/exec_ssh_server_test.go b/internal/tests/exec_ssh_server_test.go index 3ea9782..971d111 100644 --- a/internal/tests/exec_ssh_server_test.go +++ b/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() diff --git a/internal/tests/exec_test.go b/internal/tests/exec_test.go index 31c44a7..2c3b314 100644 --- a/internal/tests/exec_test.go +++ b/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()