public_git
/
orchard
mirror of https://github.com/cirruslabs/orchard.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor exec sessions into a shared SSH engine

This commit is contained in:
Fedor Korotkov 2026-04-28 16:30:09 -04:00
parent 8f5e55bcca
commit 82f827cffb
3 changed files with 257 additions and 178 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

248
internal/controller/api_vms_exec.go
View File

@ -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,6 +138,52 @@ func (controller *Controller) execVMReconnectable(
var err error
session, _, err = controller.execSessions.getOrCreate(waitContext, key, func() (*execSession, error) {
return controller.newSSHExecSession(
ctx,
waitContext,
vm,
key,
command,
stdin,
controller.execSessions,
reconnectableExecSessionPolicy,
)
})
if err != nil {
return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
}
if !session.commandMatches(command) {
return responder.JSON(http.StatusConflict,
NewErrorResponse("exec session %q is already running a different command", sessionID))
}
}
wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
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),
@ -221,7 +201,7 @@ func (controller *Controller) execVMReconnectable(
if err != nil {
_ = portForwardConn.Close()
return nil, err
return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err)
}
return newExecSession(
@ -229,30 +209,17 @@ func (controller *Controller) execVMReconnectable(
command,
exec,
portForwardConn,
controller.execSessions,
registry,
controller.execSessionExitTTL,
policy,
), nil
})
if err != nil {
return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
}
if !session.commandMatches(command) {
return responder.JSON(http.StatusConflict,
NewErrorResponse("exec session %q is already running a different command", sessionID))
}
}
wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
OriginPatterns: []string{"*"},
})
if err != nil {
return responder.Error(err)
}
defer func() {
_ = wsConn.CloseNow()
}()
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)
}
}
}

75
internal/controller/exec_sessions.go
View File

@ -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,10 +183,12 @@ type execSession struct {
bufferBytes int
nextWatermark uint64
ackedWatermark uint64
started bool
finished bool
closed bool
expiryTimer *time.Timer
startOnce sync.Once
done chan struct{}
doneOnce sync.Once
}
@ -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,8 +391,10 @@ func (session *execSession) close() {
if session.transport != nil {
_ = session.transport.Close()
}
if session.registry != nil {
session.registry.remove(session.key, session)
}
}
func (session *execSession) run() {
outgoingFrames := make(chan *execstream.Frame)
@ -368,8 +428,9 @@ func (session *execSession) recordFrame(frame *execstream.Frame) {
return
}
session.nextWatermark++
frame = cloneExecFrame(frame)
if session.policy.replayEnabled {
session.nextWatermark++
frame.Watermark = session.nextWatermark
session.frames = append(session.frames, execReplayFrame{
@ -379,6 +440,7 @@ func (session *execSession) recordFrame(frame *execstream.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() {

82
internal/controller/exec_sessions_test.go
View File

@ -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"}