From 82f827cffbe06664c1dbfda03780cd314c771f34 Mon Sep 17 00:00:00 2001 From: Fedor Korotkov Date: Tue, 28 Apr 2026 16:30:09 -0400 Subject: [PATCH] Refactor exec sessions into a shared SSH engine --- internal/controller/api_vms_exec.go | 256 ++++++++-------------- internal/controller/exec_sessions.go | 97 ++++++-- internal/controller/exec_sessions_test.go | 82 +++++++ 3 files changed, 257 insertions(+), 178 deletions(-) diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go index 92fa319..5e9a868 100644 --- a/internal/controller/api_vms_exec.go +++ b/internal/controller/api_vms_exec.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "io" "net" "net/http" "strconv" @@ -70,33 +69,27 @@ func (controller *Controller) execVMLegacy( return responderImpl } - // Establish a port-forwarding connection to a VM's SSH port - portForwardConn, err := retry.NewWithData[net.Conn]( - retry.Context(waitContext), - retry.DelayType(retry.FixedDelay), - retry.Delay(time.Second), - retry.Attempts(0), - retry.LastErrorOnly(true), - ).Do(func() (net.Conn, error) { - return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22) - }) + session, err := controller.newSSHExecSession( + ctx, + waitContext, + vm, + execSessionKey{vmName: name}, + command, + stdin, + nil, + legacyExecSessionPolicy, + ) if err != nil { return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err)) } - defer portForwardConn.Close() - - // Establish an SSH connection to a VM - exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin) - if err != nil { - return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("failed to establish SSH connection to a VM: %v", err)) - } - defer exec.Close() // Upgrade HTTP request to a WebSocket connection wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{ OriginPatterns: []string{"*"}, }) if err != nil { + session.closeIfUnused() + return responder.Error(err) } defer func() { @@ -107,66 +100,7 @@ func (controller *Controller) execVMLegacy( _ = wsConn.CloseNow() }() - // Read WebSocket frames - readFramesErrCh := make(chan error, 1) - go func() { - readFramesErrCh <- controller.readFrames(ctx, wsConn, exec.Stdin()) - }() - - // Run the command - sshErrCh := make(chan error, 1) - outgoingFrames := make(chan *execstream.Frame) - go func() { - sshErrCh <- exec.Run(ctx, command, outgoingFrames) - }() - - for { - select { - case readFramesErr := <-readFramesErrCh: - controller.logger.Warnf("failed to read and process frames from WebSocket: %v", readFramesErr) - - return responder.Empty() - case outgoingFrame := <-outgoingFrames: - if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil { - controller.logger.Warnf("failed to write WebSocket frame to the client: %v", err) - - return responder.Empty() - } - case sshErr := <-sshErrCh: - if sshErr != nil { - if err := execstream.WriteFrame(ctx, wsConn, &execstream.Frame{ - Type: execstream.FrameTypeError, - Error: sshErr.Error(), - }); err != nil { - controller.logger.Warnf("exec: failed to write error frame to WebSocket: %v", err) - } - } - - if err := wsConn.Close(websocket.StatusNormalClosure, "Command finished"); err != nil { - controller.logger.Warnf("exec: failed to close WebSocket cleanly: %v", err) - } - - if readFramesErrCh != nil { - // Read() on a WebSocket should unblock shortly after calling Close() - <-readFramesErrCh - } - - return responder.Empty() - case <-time.After(controller.pingInterval): - pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second) - - if err := wsConn.Ping(pingCtx); err != nil { - controller.logger.Warnf("port forwarding: failed to ping the client, "+ - "connection might time out: %v", err) - } - - pingCtxCancel() - case <-ctx.Done(): - controller.logger.Warnf("client disconnected prematurely") - - return responder.Empty() - } - } + return controller.serveExecSession(ctx, wsConn, session) } func (controller *Controller) execVMReconnectable( @@ -204,34 +138,16 @@ func (controller *Controller) execVMReconnectable( var err error session, _, err = controller.execSessions.getOrCreate(waitContext, key, func() (*execSession, error) { - portForwardConn, err := retry.NewWithData[net.Conn]( - retry.Context(waitContext), - retry.DelayType(retry.FixedDelay), - retry.Delay(time.Second), - retry.Attempts(0), - retry.LastErrorOnly(true), - ).Do(func() (net.Conn, error) { - return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22) - }) - if err != nil { - return nil, err - } - - exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin) - if err != nil { - _ = portForwardConn.Close() - - return nil, err - } - - return newExecSession( + return controller.newSSHExecSession( + ctx, + waitContext, + vm, key, command, - exec, - portForwardConn, + stdin, controller.execSessions, - controller.execSessionExitTTL, - ), nil + reconnectableExecSessionPolicy, + ) }) if err != nil { return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err)) @@ -247,12 +163,63 @@ func (controller *Controller) execVMReconnectable( OriginPatterns: []string{"*"}, }) if err != nil { + session.closeIfUnused() + return responder.Error(err) } defer func() { _ = wsConn.CloseNow() }() + return controller.serveExecSession(ctx, wsConn, session) +} + +func (controller *Controller) newSSHExecSession( + ctx *gin.Context, + waitContext context.Context, + vm *v1.VM, + key execSessionKey, + command string, + stdin bool, + registry *execSessionRegistry, + policy execSessionPolicy, +) (*execSession, error) { + portForwardConn, err := retry.NewWithData[net.Conn]( + retry.Context(waitContext), + retry.DelayType(retry.FixedDelay), + retry.Delay(time.Second), + retry.Attempts(0), + retry.LastErrorOnly(true), + ).Do(func() (net.Conn, error) { + return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22) + }) + if err != nil { + return nil, err + } + + exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin) + if err != nil { + _ = portForwardConn.Close() + + return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err) + } + + return newExecSession( + key, + command, + exec, + portForwardConn, + registry, + controller.execSessionExitTTL, + policy, + ), nil +} + +func (controller *Controller) serveExecSession( + ctx *gin.Context, + wsConn *websocket.Conn, + session *execSession, +) responder.Responder { subscriber, err := session.attach() if err != nil { _ = wsConn.Close(websocket.StatusNormalClosure, err.Error()) @@ -260,10 +227,11 @@ func (controller *Controller) execVMReconnectable( return responder.Empty() } defer session.detach(subscriber) + session.start() readFramesErrCh := make(chan error, 1) go func() { - readFramesErrCh <- controller.readReconnectableFrames(ctx, wsConn, session, subscriber) + readFramesErrCh <- controller.readExecSessionFrames(ctx, wsConn, session, subscriber) }() for { @@ -272,7 +240,7 @@ func (controller *Controller) execVMReconnectable( if readFramesErr != nil && !errors.Is(readFramesErr, errExecSessionDetached) && !errors.Is(readFramesErr, errExecSessionClosed) { - controller.logger.Warnf("failed to read and process reconnectable exec frames from WebSocket: %v", + controller.logger.Warnf("failed to read and process exec frames from WebSocket: %v", readFramesErr) } @@ -287,7 +255,7 @@ func (controller *Controller) execVMReconnectable( } if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil { - controller.logger.Warnf("failed to write reconnectable exec frame to the client: %v", err) + controller.logger.Warnf("failed to write exec frame to the client: %v", err) return responder.Empty() } @@ -295,7 +263,7 @@ func (controller *Controller) execVMReconnectable( pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second) if err := wsConn.Ping(pingCtx); err != nil { - controller.logger.Warnf("reconnectable exec: failed to ping the client, "+ + controller.logger.Warnf("exec: failed to ping the client, "+ "connection might time out: %v", err) } @@ -313,7 +281,7 @@ var ( errExecSessionClosed = errors.New("exec session closed") ) -func (controller *Controller) readReconnectableFrames( +func (controller *Controller) readExecSessionFrames( ctx context.Context, wsConn *websocket.Conn, session *execSession, @@ -346,12 +314,28 @@ func (controller *Controller) readReconnectableFrames( return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err) } case execstream.FrameTypeHistory: + if !session.policy.replayEnabled { + return fmt.Errorf("unexpected frame type received: %q", frame.Type) + } + session.sendHistory(subscriber, frame.Watermark) case execstream.FrameTypeAck: + if !session.policy.replayEnabled { + return fmt.Errorf("unexpected frame type received: %q", frame.Type) + } + session.ack(frame.Watermark) case execstream.FrameTypeDetach: + if !session.policy.replayEnabled { + return fmt.Errorf("unexpected frame type received: %q", frame.Type) + } + return errExecSessionDetached case execstream.FrameTypeClose: + if !session.policy.replayEnabled { + return fmt.Errorf("unexpected frame type received: %q", frame.Type) + } + session.close() return errExecSessionClosed @@ -360,57 +344,3 @@ func (controller *Controller) readReconnectableFrames( } } } - -func (controller *Controller) readFrames( - ctx context.Context, - wsConn *websocket.Conn, - stdinHandle io.WriteCloser, -) error { - for { - var frame execstream.Frame - - messageType, payloadBytes, err := wsConn.Read(ctx) - if err != nil { - var closeErr websocket.CloseError - if errors.As(err, &closeErr) && closeErr.Code == websocket.StatusNormalClosure { - return nil - } - - return fmt.Errorf("failed to read next frame from WebSocket: %w", err) - } - - if messageType != websocket.MessageText { - continue - } - - if err := json.Unmarshal(payloadBytes, &frame); err != nil { - return err - } - - switch frame.Type { - case execstream.FrameTypeStdin: - if stdinHandle == nil { - return fmt.Errorf("failed to handle %q frame: this exec session "+ - "has no stdin is enabled or already closed", frame.Type) - } - - if len(frame.Data) == 0 { - if err := stdinHandle.Close(); err != nil { - return fmt.Errorf("failed to handle %q frame: failed to close "+ - "stdin: %w", frame.Type, err) - } - - stdinHandle = nil - - continue - } - - if _, err := stdinHandle.Write(frame.Data); err != nil { - return fmt.Errorf("failed to handle %q frame: failed to write "+ - "to stdin: %w", frame.Type, err) - } - default: - return fmt.Errorf("unexpected frame type received: %q", frame.Type) - } - } -} diff --git a/internal/controller/exec_sessions.go b/internal/controller/exec_sessions.go index cee6a57..3115017 100644 --- a/internal/controller/exec_sessions.go +++ b/internal/controller/exec_sessions.go @@ -13,6 +13,22 @@ import ( const execSessionReplayBufferBytes = 4 * 1024 * 1024 +type execSessionPolicy struct { + closeOnDetach bool + retainAfterExit bool + replayEnabled bool +} + +var ( + legacyExecSessionPolicy = execSessionPolicy{ + closeOnDetach: true, + } + reconnectableExecSessionPolicy = execSessionPolicy{ + retainAfterExit: true, + replayEnabled: true, + } +) + type sshExecRunner interface { Stdin() io.WriteCloser Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error @@ -154,6 +170,7 @@ type execSession struct { transport net.Conn registry *execSessionRegistry exitTTL time.Duration + policy execSessionPolicy ctx context.Context cancel context.CancelFunc @@ -166,12 +183,14 @@ type execSession struct { bufferBytes int nextWatermark uint64 ackedWatermark uint64 + started bool finished bool closed bool expiryTimer *time.Timer - done chan struct{} - doneOnce sync.Once + startOnce sync.Once + done chan struct{} + doneOnce sync.Once } func newExecSession( @@ -181,6 +200,7 @@ func newExecSession( transport net.Conn, registry *execSessionRegistry, exitTTL time.Duration, + policy execSessionPolicy, ) *execSession { ctx, cancel := context.WithCancel(context.Background()) @@ -191,6 +211,7 @@ func newExecSession( transport: transport, registry: registry, exitTTL: exitTTL, + policy: policy, ctx: ctx, cancel: cancel, stdin: exec.Stdin(), @@ -198,8 +219,6 @@ func newExecSession( done: make(chan struct{}), } - go session.run() - return session } @@ -207,6 +226,31 @@ func (session *execSession) commandMatches(command string) bool { return command == "" || session.command == command } +func (session *execSession) start() { + session.startOnce.Do(func() { + session.mu.Lock() + if session.closed { + session.mu.Unlock() + + return + } + session.started = true + session.mu.Unlock() + + go session.run() + }) +} + +func (session *execSession) closeIfUnused() { + session.mu.Lock() + unused := !session.started && len(session.subscribers) == 0 + session.mu.Unlock() + + if unused { + session.close() + } +} + func (session *execSession) attach() (*execSessionSubscriber, error) { session.mu.Lock() defer session.mu.Unlock() @@ -222,6 +266,12 @@ func (session *execSession) attach() (*execSessionSubscriber, error) { } func (session *execSession) detach(subscriber *execSessionSubscriber) { + if session.policy.closeOnDetach { + session.close() + + return + } + session.mu.Lock() defer session.mu.Unlock() @@ -261,6 +311,10 @@ func (session *execSession) writeStdin(data []byte) error { } func (session *execSession) ack(watermark uint64) { + if !session.policy.replayEnabled { + return + } + session.mu.Lock() defer session.mu.Unlock() @@ -276,6 +330,10 @@ func (session *execSession) sendHistory( subscriber *execSessionSubscriber, watermark uint64, ) { + if !session.policy.replayEnabled { + return + } + session.mu.Lock() defer session.mu.Unlock() @@ -333,7 +391,9 @@ func (session *execSession) close() { if session.transport != nil { _ = session.transport.Close() } - session.registry.remove(session.key, session) + if session.registry != nil { + session.registry.remove(session.key, session) + } } func (session *execSession) run() { @@ -368,17 +428,19 @@ func (session *execSession) recordFrame(frame *execstream.Frame) { return } - session.nextWatermark++ frame = cloneExecFrame(frame) - frame.Watermark = session.nextWatermark + if session.policy.replayEnabled { + session.nextWatermark++ + frame.Watermark = session.nextWatermark - session.frames = append(session.frames, execReplayFrame{ - frame: frame, - size: execFrameSize(frame), - }) - session.bufferBytes += execFrameSize(frame) - session.trimAcknowledgedLocked() - session.trimToLimitLocked() + session.frames = append(session.frames, execReplayFrame{ + frame: frame, + size: execFrameSize(frame), + }) + session.bufferBytes += execFrameSize(frame) + session.trimAcknowledgedLocked() + session.trimToLimitLocked() + } for subscriber := range session.subscribers { if subscriber.enqueue(frame) { @@ -398,7 +460,8 @@ func (session *execSession) markFinished() { } session.finished = true - if !session.closed { + shouldClose := !session.policy.retainAfterExit + if !session.closed && session.policy.retainAfterExit { session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire) } @@ -416,6 +479,10 @@ func (session *execSession) markFinished() { session.doneOnce.Do(func() { close(session.done) }) + + if shouldClose { + session.close() + } } func (session *execSession) expire() { diff --git a/internal/controller/exec_sessions_test.go b/internal/controller/exec_sessions_test.go index ce9f5be..e5042d5 100644 --- a/internal/controller/exec_sessions_test.go +++ b/internal/controller/exec_sessions_test.go @@ -51,6 +51,7 @@ func newManualExecSessionForTest( exec: &fakeExec{}, registry: registry, exitTTL: time.Minute, + policy: reconnectableExecSessionPolicy, ctx: ctx, cancel: cancel, subscribers: map[*execSessionSubscriber]struct{}{}, @@ -98,6 +99,36 @@ func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) { require.EqualValues(t, 1, createCalls.Load()) } +func TestExecSessionStartRunsCommandOnlyOnce(t *testing.T) { + var runCalls atomic.Int32 + runStarted := make(chan struct{}) + + session := newExecSession( + execSessionKey{vmName: "vm", sessionID: "session"}, + "echo test", + &fakeExec{ + run: func(ctx context.Context, _ string, _ chan<- *execstream.Frame) error { + runCalls.Add(1) + close(runStarted) + <-ctx.Done() + + return ctx.Err() + }, + }, + nil, + nil, + time.Minute, + reconnectableExecSessionPolicy, + ) + defer session.close() + + session.start() + session.start() + + <-runStarted + require.EqualValues(t, 1, runCalls.Load()) +} + func TestExecSessionHistoryReplayAndAck(t *testing.T) { registry := newExecSessionRegistry() session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) @@ -140,6 +171,57 @@ func TestExecSessionDetachKeepsProcessAlive(t *testing.T) { require.EqualValues(t, 0, exec.closeCalls.Load()) } +func TestLegacyExecSessionDetachStopsProcess(t *testing.T) { + registry := newExecSessionRegistry() + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) + session.policy = legacyExecSessionPolicy + exec := session.exec.(*fakeExec) + + subscriber, err := session.attach() + require.NoError(t, err) + + session.detach(subscriber) + + require.True(t, session.closed) + require.EqualValues(t, 1, exec.closeCalls.Load()) +} + +func TestLegacyExecSessionDoesNotRetainReplayHistory(t *testing.T) { + registry := newExecSessionRegistry() + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) + session.policy = legacyExecSessionPolicy + + session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")}) + + require.Empty(t, session.frames) + require.Zero(t, session.nextWatermark) +} + +func TestExecSessionCloseIfUnusedClosesIdleSession(t *testing.T) { + registry := newExecSessionRegistry() + key := execSessionKey{vmName: "vm", sessionID: "session"} + session := newManualExecSessionForTest(key, registry) + exec := session.exec.(*fakeExec) + registry.sessions[key] = session + + session.closeIfUnused() + + require.True(t, session.closed) + require.EqualValues(t, 1, exec.closeCalls.Load()) +} + +func TestExecSessionCloseIfUnusedKeepsAttachedSession(t *testing.T) { + registry := newExecSessionRegistry() + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) + + _, err := session.attach() + require.NoError(t, err) + + session.closeIfUnused() + + require.False(t, session.closed) +} + func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) { registry := newExecSessionRegistry() key := execSessionKey{vmName: "vm", sessionID: "session"}