diff --git a/api/openapi.yaml b/api/openapi.yaml index 29823a2..10f4d89 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -429,11 +429,33 @@ paths: parameters: - in: query name: command - description: Command to execute + description: | + Command to execute. + + Required when starting a new exec session. May be omitted when reconnecting to an + existing session identified by the `session` parameter. schema: type: string minLength: 1 - required: true + required: false + - in: query + name: session + description: | + Optional stable exec session identifier. When present, websocket disconnects detach + from the command instead of terminating it, and later requests with the same VM name + and session id may reconnect and request buffered history. + schema: + type: string + minLength: 1 + required: false + - in: query + name: cmux_session_id + description: | + Compatibility alias for `session`. Prefer `session` for new Orchard clients. + schema: + type: string + minLength: 1 + required: false - in: query name: stdin description: | @@ -483,7 +505,9 @@ paths: '400': description: Invalid parameters were supplied '404': - description: VM resource with the given name doesn't exist + description: VM resource with the given name or reconnectable exec session doesn't exist + '409': + description: Reconnectable exec session already exists with a different command '503': description: Controller failed to establish a connection with the VM /vms/{name}/ip: @@ -800,10 +824,18 @@ components: description: WebSocket frame from Orchard Client to the Orchard Controller oneOf: - $ref: '#/components/schemas/ExecClientFrameStdin' + - $ref: '#/components/schemas/ExecClientFrameHistory' + - $ref: '#/components/schemas/ExecClientFrameAck' + - $ref: '#/components/schemas/ExecClientFrameDetach' + - $ref: '#/components/schemas/ExecClientFrameClose' discriminator: propertyName: type mapping: stdin: '#/components/schemas/ExecClientFrameStdin' + history: '#/components/schemas/ExecClientFrameHistory' + ack: '#/components/schemas/ExecClientFrameAck' + detach: '#/components/schemas/ExecClientFrameDetach' + close: '#/components/schemas/ExecClientFrameClose' ExecClientFrameStdin: description: Send bytes to the process standard input type: object @@ -822,6 +854,56 @@ components: example: type: stdin data: aGVsbG8K + ExecClientFrameHistory: + description: Request buffered output strictly newer than the supplied watermark + type: object + required: [ type, watermark ] + properties: + type: + type: string + enum: [ history ] + watermark: + type: integer + format: int64 + minimum: 0 + example: + type: history + watermark: 42 + ExecClientFrameAck: + description: Acknowledge that output has been durably consumed through the supplied watermark + type: object + required: [ type, watermark ] + properties: + type: + type: string + enum: [ ack ] + watermark: + type: integer + format: int64 + minimum: 0 + example: + type: ack + watermark: 42 + ExecClientFrameDetach: + description: Detach this websocket while leaving the remote command running + type: object + required: [ type ] + properties: + type: + type: string + enum: [ detach ] + example: + type: detach + ExecClientFrameClose: + description: Close the reconnectable exec session and terminate the remote command + type: object + required: [ type ] + properties: + type: + type: string + enum: [ close ] + example: + type: close ExecControllerFrame: description: WebSocket frame from Orchard Controller to the Orchard Client oneOf: @@ -829,6 +911,7 @@ components: - $ref: '#/components/schemas/ExecControllerFrameStderr' - $ref: '#/components/schemas/ExecControllerFrameExit' - $ref: '#/components/schemas/ExecControllerFrameError' + - $ref: '#/components/schemas/ExecControllerFrameNoMoreHistory' discriminator: propertyName: type mapping: @@ -836,6 +919,7 @@ components: stderr: '#/components/schemas/ExecControllerFrameStderr' exit: '#/components/schemas/ExecControllerFrameExit' error: '#/components/schemas/ExecControllerFrameError' + no_more_history: '#/components/schemas/ExecControllerFrameNoMoreHistory' ExecControllerFrameStdout: description: Standard output from the process type: object @@ -848,6 +932,10 @@ components: type: string format: byte description: Base64-encoded standard output bytes from the process + watermark: + type: integer + format: int64 + description: Monotonic output watermark present on reconnectable sessions example: type: stdout data: aGVsbG8K @@ -863,6 +951,10 @@ components: type: string format: byte description: Base64-encoded standard error bytes from the process + watermark: + type: integer + format: int64 + description: Monotonic output watermark present on reconnectable sessions example: type: stderr data: aGVsbG8K @@ -882,6 +974,10 @@ components: type: integer format: int32 description: Process exit code + watermark: + type: integer + format: int64 + description: Monotonic output watermark present on reconnectable sessions example: type: exit exit: @@ -897,9 +993,28 @@ components: error: type: string description: Error message text + watermark: + type: integer + format: int64 + description: Monotonic output watermark present on reconnectable sessions example: type: error error: Failed to establish SSH connection to a VM + ExecControllerFrameNoMoreHistory: + description: Marker indicating that the requested replay range has been fully sent + type: object + required: [ type, watermark ] + properties: + type: + type: string + enum: [ no_more_history ] + watermark: + type: integer + format: int64 + description: Highest watermark known to the exec session at the time of replay + example: + type: no_more_history + watermark: 42 Event: title: Generic Resource Event type: object diff --git a/internal/command/controller/run.go b/internal/command/controller/run.go index 5406654..8833f93 100644 --- a/internal/command/controller/run.go +++ b/internal/command/controller/run.go @@ -35,6 +35,7 @@ var noExperimentalRPCV2 bool var experimentalPingInterval time.Duration var experimentalDisableDBCompression bool var workerOfflineTimeout time.Duration +var execSessionExitTTL time.Duration var synthetic bool func newRunCommand() *cobra.Command { @@ -87,6 +88,8 @@ 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, + "duration to retain reconnectable exec session history after the command exits") // Hidden flags cmd.Flags().BoolVar(&synthetic, "synthetic", false, "") @@ -147,6 +150,7 @@ func runController(cmd *cobra.Command, args []string) (err error) { controller.WithListenAddr(address), controller.WithDataDir(dataDir), controller.WithWorkerOfflineTimeout(workerOfflineTimeout), + controller.WithExecSessionExitTTL(execSessionExitTTL), controller.WithLogger(logger), } diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go index de384ca..92fa319 100644 --- a/internal/controller/api_vms_exec.go +++ b/internal/controller/api_vms_exec.go @@ -28,9 +28,13 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { // Retrieve and parse path and query parameters name := ctx.Param("name") + sessionID := ctx.Query("session") + if sessionID == "" { + sessionID = ctx.Query("cmux_session_id") + } command := ctx.Query("command") - if command == "" { + if sessionID == "" && command == "" { return responder.JSON(http.StatusBadRequest, NewErrorResponse("\"command\" parameter cannot be empty")) } @@ -43,6 +47,20 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { return responder.Code(http.StatusBadRequest) } + if sessionID != "" { + return controller.execVMReconnectable(ctx, name, sessionID, command, stdin, wait) + } + + return controller.execVMLegacy(ctx, name, command, stdin, wait) +} + +func (controller *Controller) execVMLegacy( + ctx *gin.Context, + name string, + command string, + stdin bool, + wait uint64, +) responder.Responder { // Look-up the VM waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second) defer waitContextCancel() @@ -151,6 +169,198 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { } } +func (controller *Controller) execVMReconnectable( + ctx *gin.Context, + name string, + sessionID string, + command string, + stdin bool, + wait uint64, +) responder.Responder { + key := execSessionKey{ + vmName: name, + sessionID: sessionID, + } + + session, ok := controller.execSessions.get(key) + if ok { + if !session.commandMatches(command) { + return responder.JSON(http.StatusConflict, + NewErrorResponse("exec session %q is already running a different command", sessionID)) + } + } else { + if command == "" { + return responder.JSON(http.StatusNotFound, + NewErrorResponse("exec session %q does not exist", sessionID)) + } + + waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second) + defer waitContextCancel() + + vm, responderImpl := controller.waitForVM(waitContext, name) + if responderImpl != nil { + return responderImpl + } + + 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( + key, + command, + exec, + portForwardConn, + controller.execSessions, + controller.execSessionExitTTL, + ), 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() + }() + + subscriber, err := session.attach() + if err != nil { + _ = wsConn.Close(websocket.StatusNormalClosure, err.Error()) + + return responder.Empty() + } + defer session.detach(subscriber) + + readFramesErrCh := make(chan error, 1) + go func() { + readFramesErrCh <- controller.readReconnectableFrames(ctx, wsConn, session, subscriber) + }() + + for { + select { + case readFramesErr := <-readFramesErrCh: + 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", + readFramesErr) + } + + return responder.Empty() + case outgoingFrame, ok := <-subscriber.frames: + if !ok { + if err := wsConn.Close(websocket.StatusNormalClosure, "Command finished"); err != nil { + controller.logger.Warnf("exec: failed to close WebSocket cleanly: %v", err) + } + + return responder.Empty() + } + + if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil { + controller.logger.Warnf("failed to write reconnectable exec frame to the client: %v", err) + + 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("reconnectable exec: failed to ping the client, "+ + "connection might time out: %v", err) + } + + pingCtxCancel() + case <-ctx.Done(): + controller.logger.Warnf("client disconnected prematurely") + + return responder.Empty() + } + } +} + +var ( + errExecSessionDetached = errors.New("exec session detached") + errExecSessionClosed = errors.New("exec session closed") +) + +func (controller *Controller) readReconnectableFrames( + ctx context.Context, + wsConn *websocket.Conn, + session *execSession, + subscriber *execSessionSubscriber, +) 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 errExecSessionDetached + } + + 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 err := session.writeStdin(frame.Data); err != nil { + return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err) + } + case execstream.FrameTypeHistory: + session.sendHistory(subscriber, frame.Watermark) + case execstream.FrameTypeAck: + session.ack(frame.Watermark) + case execstream.FrameTypeDetach: + return errExecSessionDetached + case execstream.FrameTypeClose: + session.close() + + return errExecSessionClosed + default: + return fmt.Errorf("unexpected frame type received: %q", frame.Type) + } + } +} + func (controller *Controller) readFrames( ctx context.Context, wsConn *websocket.Conn, diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 906f9d2..1521e0d 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -55,6 +55,7 @@ type Controller struct { ipRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]] enableSwaggerDocs bool workerOfflineTimeout time.Duration + execSessionExitTTL time.Duration experimentalRPCV2 bool disableDBCompression bool pingInterval time.Duration @@ -64,6 +65,7 @@ type Controller struct { sshSigner ssh.Signer sshNoClientAuth bool sshServer *sshserver.SSHServer + execSessions *execSessionRegistry single singleflight.Group @@ -75,7 +77,9 @@ func New(opts ...Option) (*Controller, error) { 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{}, } @@ -308,6 +312,8 @@ func (controller *Controller) Run(ctx context.Context) error { go func() { <-ctx.Done() + controller.execSessions.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 new file mode 100644 index 0000000..cee6a57 --- /dev/null +++ b/internal/controller/exec_sessions.go @@ -0,0 +1,462 @@ +package controller + +import ( + "context" + "errors" + "io" + "net" + "sync" + "time" + + "github.com/cirruslabs/orchard/internal/execstream" +) + +const execSessionReplayBufferBytes = 4 * 1024 * 1024 + +type sshExecRunner interface { + Stdin() io.WriteCloser + Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error + Close() error +} + +type execSessionKey struct { + vmName string + sessionID string +} + +type execSessionCreation struct { + done chan struct{} + session *execSession + err error +} + +type execSessionRegistry struct { + mu sync.Mutex + sessions map[execSessionKey]*execSession + creating map[execSessionKey]*execSessionCreation +} + +func newExecSessionRegistry() *execSessionRegistry { + return &execSessionRegistry{ + sessions: map[execSessionKey]*execSession{}, + creating: map[execSessionKey]*execSessionCreation{}, + } +} + +func (registry *execSessionRegistry) get(key execSessionKey) (*execSession, bool) { + registry.mu.Lock() + defer registry.mu.Unlock() + + session, ok := registry.sessions[key] + + return session, ok +} + +func (registry *execSessionRegistry) getOrCreate( + ctx context.Context, + key execSessionKey, + create func() (*execSession, error), +) (*execSession, bool, error) { + for { + registry.mu.Lock() + + if session, ok := registry.sessions[key]; ok { + registry.mu.Unlock() + + return session, false, nil + } + + if creation, ok := registry.creating[key]; ok { + registry.mu.Unlock() + + select { + case <-ctx.Done(): + return nil, false, ctx.Err() + case <-creation.done: + if creation.err != nil { + return nil, false, creation.err + } + + return creation.session, false, nil + } + } + + creation := &execSessionCreation{done: make(chan struct{})} + registry.creating[key] = creation + registry.mu.Unlock() + + session, err := create() + + registry.mu.Lock() + delete(registry.creating, key) + if err == nil { + registry.sessions[key] = session + } + creation.session = session + creation.err = err + close(creation.done) + registry.mu.Unlock() + + return session, true, err + } +} + +func (registry *execSessionRegistry) remove(key execSessionKey, expected *execSession) { + registry.mu.Lock() + defer registry.mu.Unlock() + + if registry.sessions[key] == expected { + delete(registry.sessions, key) + } +} + +func (registry *execSessionRegistry) closeAll() { + registry.mu.Lock() + sessions := make([]*execSession, 0, len(registry.sessions)) + for _, session := range registry.sessions { + sessions = append(sessions, session) + } + registry.mu.Unlock() + + for _, session := range sessions { + session.close() + } +} + +type execReplayFrame struct { + frame *execstream.Frame + size int +} + +type execSessionSubscriber struct { + frames chan *execstream.Frame +} + +func newExecSessionSubscriber() *execSessionSubscriber { + return &execSessionSubscriber{ + frames: make(chan *execstream.Frame, 128), + } +} + +func (subscriber *execSessionSubscriber) enqueue(frame *execstream.Frame) bool { + select { + case subscriber.frames <- cloneExecFrame(frame): + return true + default: + return false + } +} + +type execSession struct { + key execSessionKey + command string + exec sshExecRunner + transport net.Conn + registry *execSessionRegistry + exitTTL time.Duration + + ctx context.Context + cancel context.CancelFunc + + mu sync.Mutex + stdin io.WriteCloser + stdinClosed bool + subscribers map[*execSessionSubscriber]struct{} + frames []execReplayFrame + bufferBytes int + nextWatermark uint64 + ackedWatermark uint64 + finished bool + closed bool + expiryTimer *time.Timer + + done chan struct{} + doneOnce sync.Once +} + +func newExecSession( + key execSessionKey, + command string, + exec sshExecRunner, + transport net.Conn, + registry *execSessionRegistry, + exitTTL time.Duration, +) *execSession { + ctx, cancel := context.WithCancel(context.Background()) + + session := &execSession{ + key: key, + command: command, + exec: exec, + transport: transport, + registry: registry, + exitTTL: exitTTL, + ctx: ctx, + cancel: cancel, + stdin: exec.Stdin(), + subscribers: map[*execSessionSubscriber]struct{}{}, + done: make(chan struct{}), + } + + go session.run() + + return session +} + +func (session *execSession) commandMatches(command string) bool { + return command == "" || session.command == command +} + +func (session *execSession) attach() (*execSessionSubscriber, error) { + session.mu.Lock() + defer session.mu.Unlock() + + if session.closed { + return nil, errors.New("exec session is closed") + } + + subscriber := newExecSessionSubscriber() + session.subscribers[subscriber] = struct{}{} + + return subscriber, nil +} + +func (session *execSession) detach(subscriber *execSessionSubscriber) { + session.mu.Lock() + defer session.mu.Unlock() + + session.detachLocked(subscriber) +} + +func (session *execSession) detachLocked(subscriber *execSessionSubscriber) { + if _, ok := session.subscribers[subscriber]; !ok { + return + } + + delete(session.subscribers, subscriber) + close(subscriber.frames) +} + +func (session *execSession) writeStdin(data []byte) error { + session.mu.Lock() + defer session.mu.Unlock() + + if session.stdin == nil || session.stdinClosed { + return errors.New("this exec session has no stdin enabled or it is already closed") + } + + if len(data) == 0 { + if err := session.stdin.Close(); err != nil { + return err + } + + session.stdinClosed = true + + return nil + } + + _, err := session.stdin.Write(data) + + return err +} + +func (session *execSession) ack(watermark uint64) { + session.mu.Lock() + defer session.mu.Unlock() + + if watermark <= session.ackedWatermark { + return + } + + session.ackedWatermark = watermark + session.trimAcknowledgedLocked() +} + +func (session *execSession) sendHistory( + subscriber *execSessionSubscriber, + watermark uint64, +) { + session.mu.Lock() + defer session.mu.Unlock() + + if _, ok := session.subscribers[subscriber]; !ok { + return + } + + for _, record := range session.frames { + if record.frame.Watermark <= watermark { + continue + } + + if !subscriber.enqueue(record.frame) { + session.detachLocked(subscriber) + + return + } + } + + if !subscriber.enqueue(&execstream.Frame{ + Type: execstream.FrameTypeNoMoreHistory, + Watermark: session.nextWatermark, + }) { + session.detachLocked(subscriber) + } +} + +func (session *execSession) close() { + session.mu.Lock() + if session.closed { + session.mu.Unlock() + + return + } + + session.closed = true + if session.expiryTimer != nil { + session.expiryTimer.Stop() + session.expiryTimer = nil + } + + subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers)) + for subscriber := range session.subscribers { + subscribers = append(subscribers, subscriber) + } + session.subscribers = map[*execSessionSubscriber]struct{}{} + session.mu.Unlock() + + for _, subscriber := range subscribers { + close(subscriber.frames) + } + + session.cancel() + _ = session.exec.Close() + if session.transport != nil { + _ = session.transport.Close() + } + session.registry.remove(session.key, session) +} + +func (session *execSession) run() { + outgoingFrames := make(chan *execstream.Frame) + runErrCh := make(chan error, 1) + + go func() { + runErrCh <- session.exec.Run(session.ctx, session.command, outgoingFrames) + close(outgoingFrames) + }() + + for frame := range outgoingFrames { + session.recordFrame(frame) + } + + runErr := <-runErrCh + if runErr != nil && !errors.Is(runErr, context.Canceled) { + session.recordFrame(&execstream.Frame{ + Type: execstream.FrameTypeError, + Error: runErr.Error(), + }) + } + + session.markFinished() +} + +func (session *execSession) recordFrame(frame *execstream.Frame) { + session.mu.Lock() + defer session.mu.Unlock() + + if session.closed { + return + } + + session.nextWatermark++ + frame = cloneExecFrame(frame) + frame.Watermark = session.nextWatermark + + 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) { + continue + } + + session.detachLocked(subscriber) + } +} + +func (session *execSession) markFinished() { + session.mu.Lock() + if session.finished { + session.mu.Unlock() + + return + } + + session.finished = true + if !session.closed { + session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire) + } + + subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers)) + for subscriber := range session.subscribers { + subscribers = append(subscribers, subscriber) + } + session.subscribers = map[*execSessionSubscriber]struct{}{} + session.mu.Unlock() + + for _, subscriber := range subscribers { + close(subscriber.frames) + } + + session.doneOnce.Do(func() { + close(session.done) + }) +} + +func (session *execSession) expire() { + session.close() +} + +func (session *execSession) trimAcknowledgedLocked() { + for len(session.frames) > 0 && session.frames[0].frame.Watermark <= session.ackedWatermark { + session.bufferBytes -= session.frames[0].size + session.frames = session.frames[1:] + } +} + +func (session *execSession) trimToLimitLocked() { + for session.bufferBytes > execSessionReplayBufferBytes && len(session.frames) > 0 { + session.bufferBytes -= session.frames[0].size + session.frames = session.frames[1:] + } +} + +func cloneExecFrame(frame *execstream.Frame) *execstream.Frame { + if frame == nil { + return nil + } + + clone := *frame + if frame.Data != nil { + clone.Data = append([]byte(nil), frame.Data...) + } + if frame.Exit != nil { + exit := *frame.Exit + clone.Exit = &exit + } + + return &clone +} + +func execFrameSize(frame *execstream.Frame) int { + if frame == nil { + return 0 + } + + return len(frame.Data) + len(frame.Error) + 16 +} diff --git a/internal/controller/exec_sessions_test.go b/internal/controller/exec_sessions_test.go new file mode 100644 index 0000000..ce9f5be --- /dev/null +++ b/internal/controller/exec_sessions_test.go @@ -0,0 +1,172 @@ +package controller + +import ( + "context" + "io" + "sync/atomic" + "testing" + "time" + + "github.com/cirruslabs/orchard/internal/execstream" + "github.com/stretchr/testify/require" +) + +type fakeExec struct { + stdin io.WriteCloser + run func(context.Context, string, chan<- *execstream.Frame) error + closeCalls atomic.Int32 +} + +func (exec *fakeExec) Stdin() io.WriteCloser { + return exec.stdin +} + +func (exec *fakeExec) Run( + ctx context.Context, + command string, + outgoingFrames chan<- *execstream.Frame, +) error { + if exec.run != nil { + return exec.run(ctx, command, outgoingFrames) + } + + return nil +} + +func (exec *fakeExec) Close() error { + exec.closeCalls.Add(1) + + return nil +} + +func newManualExecSessionForTest( + key execSessionKey, + registry *execSessionRegistry, +) *execSession { + ctx, cancel := context.WithCancel(context.Background()) + + return &execSession{ + key: key, + command: "echo test", + exec: &fakeExec{}, + registry: registry, + exitTTL: time.Minute, + ctx: ctx, + cancel: cancel, + subscribers: map[*execSessionSubscriber]struct{}{}, + done: make(chan struct{}), + } +} + +func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) { + registry := newExecSessionRegistry() + key := execSessionKey{vmName: "vm", sessionID: "session"} + + createStarted := make(chan struct{}) + releaseCreate := make(chan struct{}) + var createCalls atomic.Int32 + + create := func() (*execSession, error) { + createCalls.Add(1) + close(createStarted) + <-releaseCreate + + return newManualExecSessionForTest(key, registry), nil + } + + firstDone := make(chan struct{}) + go func() { + defer close(firstDone) + _, _, err := registry.getOrCreate(context.Background(), key, create) + require.NoError(t, err) + }() + + <-createStarted + + secondDone := make(chan struct{}) + go func() { + defer close(secondDone) + _, created, err := registry.getOrCreate(context.Background(), key, create) + require.NoError(t, err) + require.False(t, created) + }() + + close(releaseCreate) + + <-firstDone + <-secondDone + require.EqualValues(t, 1, createCalls.Load()) +} + +func TestExecSessionHistoryReplayAndAck(t *testing.T) { + registry := newExecSessionRegistry() + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) + + session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")}) + session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStderr, Data: []byte("err")}) + session.recordFrame(&execstream.Frame{ + Type: execstream.FrameTypeExit, + Exit: &execstream.Exit{Code: 7}, + }) + + subscriber, err := session.attach() + require.NoError(t, err) + + session.sendHistory(subscriber, 0) + + require.Equal(t, execstream.FrameTypeStdout, (<-subscriber.frames).Type) + require.Equal(t, execstream.FrameTypeStderr, (<-subscriber.frames).Type) + require.Equal(t, execstream.FrameTypeExit, (<-subscriber.frames).Type) + noMoreHistory := <-subscriber.frames + require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type) + require.EqualValues(t, 3, noMoreHistory.Watermark) + + session.ack(2) + require.Len(t, session.frames, 1) + require.EqualValues(t, 3, session.frames[0].frame.Watermark) +} + +func TestExecSessionDetachKeepsProcessAlive(t *testing.T) { + registry := newExecSessionRegistry() + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) + exec := session.exec.(*fakeExec) + + subscriber, err := session.attach() + require.NoError(t, err) + + session.detach(subscriber) + + require.False(t, session.closed) + require.EqualValues(t, 0, exec.closeCalls.Load()) +} + +func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) { + registry := newExecSessionRegistry() + key := execSessionKey{vmName: "vm", sessionID: "session"} + session := newManualExecSessionForTest(key, registry) + exec := session.exec.(*fakeExec) + registry.sessions[key] = session + + session.close() + + require.True(t, session.closed) + require.EqualValues(t, 1, exec.closeCalls.Load()) + _, ok := registry.get(key) + require.False(t, ok) +} + +func TestExecSessionFinishedEntryExpiresAfterTTL(t *testing.T) { + registry := newExecSessionRegistry() + key := execSessionKey{vmName: "vm", sessionID: "session"} + session := newManualExecSessionForTest(key, registry) + session.exitTTL = 10 * time.Millisecond + registry.sessions[key] = session + + session.markFinished() + + require.Eventually(t, func() bool { + _, ok := registry.get(key) + + return !ok + }, time.Second, 10*time.Millisecond) +} diff --git a/internal/controller/option.go b/internal/controller/option.go index 8370d50..325c912 100644 --- a/internal/controller/option.go +++ b/internal/controller/option.go @@ -60,6 +60,12 @@ func WithWorkerOfflineTimeout(workerOfflineTimeout time.Duration) Option { } } +func WithExecSessionExitTTL(execSessionExitTTL time.Duration) Option { + return func(controller *Controller) { + controller.execSessionExitTTL = execSessionExitTTL + } +} + func WithExperimentalRPCV2() Option { return func(controller *Controller) { controller.experimentalRPCV2 = true diff --git a/internal/execstream/frame.go b/internal/execstream/frame.go index c40f252..ab41f0b 100644 --- a/internal/execstream/frame.go +++ b/internal/execstream/frame.go @@ -10,18 +10,24 @@ import ( type FrameType string const ( - FrameTypeStdin FrameType = "stdin" - FrameTypeStdout FrameType = "stdout" - FrameTypeStderr FrameType = "stderr" - FrameTypeExit FrameType = "exit" - FrameTypeError FrameType = "error" + FrameTypeStdin FrameType = "stdin" + FrameTypeStdout FrameType = "stdout" + FrameTypeStderr FrameType = "stderr" + FrameTypeExit FrameType = "exit" + FrameTypeError FrameType = "error" + FrameTypeHistory FrameType = "history" + FrameTypeNoMoreHistory FrameType = "no_more_history" + FrameTypeAck FrameType = "ack" + FrameTypeDetach FrameType = "detach" + FrameTypeClose FrameType = "close" ) type Frame struct { - Type FrameType `json:"type"` - Data []byte `json:"data,omitempty"` - Exit *Exit `json:"exit,omitempty"` - Error string `json:"error,omitempty"` + Type FrameType `json:"type"` + Data []byte `json:"data,omitempty"` + Exit *Exit `json:"exit,omitempty"` + Error string `json:"error,omitempty"` + Watermark uint64 `json:"watermark,omitempty"` } type Exit struct { diff --git a/internal/execstream/frame_test.go b/internal/execstream/frame_test.go new file mode 100644 index 0000000..673f0c3 --- /dev/null +++ b/internal/execstream/frame_test.go @@ -0,0 +1,23 @@ +package execstream + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestFrameRoundTripsWatermark(t *testing.T) { + frame := Frame{ + Type: FrameTypeHistory, + Watermark: 42, + } + + payload, err := json.Marshal(frame) + require.NoError(t, err) + + var decoded Frame + err = json.Unmarshal(payload, &decoded) + require.NoError(t, err) + require.Equal(t, frame, decoded) +} diff --git a/internal/tests/exec_test.go b/internal/tests/exec_test.go index 227f66d..af87560 100644 --- a/internal/tests/exec_test.go +++ b/internal/tests/exec_test.go @@ -132,6 +132,167 @@ func TestVMExecScript(t *testing.T) { require.Equal(t, websocket.StatusNormalClosure, closeError.Code) } +func TestVMExecSessionReconnectHistory(t *testing.T) { + devClient, vmName := prepareForExec(t) + sessionID := uuid.NewString() + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: "sh -c 'echo first; sleep 1; echo second'", + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + + firstFrame := readFrame(t, wsConn) + require.Equal(t, execstream.FrameTypeStdout, firstFrame.Type) + require.Equal(t, "first\n", string(firstFrame.Data)) + require.EqualValues(t, 1, firstFrame.Watermark) + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach}) + require.NoError(t, err) + _ = wsConn.CloseNow() + + wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeHistory, + Watermark: firstFrame.Watermark, + }) + require.NoError(t, err) + + frames := readFramesUntilExit(t, wsConn) + require.Len(t, framesByType(frames, execstream.FrameTypeStdout), 1) + require.Equal(t, "second\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data)) + require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code) +} + +func TestVMExecSessionReconnectAfterExit(t *testing.T) { + devClient, vmName := prepareForExec(t) + sessionID := uuid.NewString() + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: "sh -c 'echo replay-me'", + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach}) + require.NoError(t, err) + _ = wsConn.CloseNow() + + time.Sleep(time.Second) + + wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory}) + require.NoError(t, err) + + frames := readFramesUntilExit(t, wsConn) + require.Equal(t, "replay-me\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data)) + require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code) +} + +func TestVMExecSessionReplayPreservesStreams(t *testing.T) { + devClient, vmName := prepareForExec(t) + sessionID := uuid.NewString() + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: "sh -c 'echo out1; sleep 1; echo err1 >&2; sleep 1; echo out2; sleep 1; echo err2 >&2'", + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach}) + require.NoError(t, err) + _ = wsConn.CloseNow() + + time.Sleep(4 * time.Second) + + wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory}) + require.NoError(t, err) + + frames := readFramesUntilExit(t, wsConn) + require.Equal(t, []execstream.FrameType{ + execstream.FrameTypeStdout, + execstream.FrameTypeStderr, + execstream.FrameTypeStdout, + execstream.FrameTypeStderr, + execstream.FrameTypeExit, + }, frameTypes(frames)) + require.Equal(t, "out1\n", string(frames[0].Data)) + require.Equal(t, "err1\n", string(frames[1].Data)) + require.Equal(t, "out2\n", string(frames[2].Data)) + require.Equal(t, "err2\n", string(frames[3].Data)) +} + +func TestVMExecSessionStdinSurvivesReconnect(t *testing.T) { + devClient, vmName := prepareForExec(t) + sessionID := uuid.NewString() + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: "/bin/cat", + Stdin: true, + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte("one\n"), + }) + require.NoError(t, err) + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach}) + require.NoError(t, err) + _ = wsConn.CloseNow() + + wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + WaitSeconds: 30, + Session: sessionID, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte("two\n"), + }) + require.NoError(t, err) + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte{}, + }) + require.NoError(t, err) + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory}) + require.NoError(t, err) + + frames := readFramesUntilExit(t, wsConn) + stdoutFrames := framesByType(frames, execstream.FrameTypeStdout) + require.Len(t, stdoutFrames, 2) + require.Equal(t, "one\n", string(stdoutFrames[0].Data)) + require.Equal(t, "two\n", string(stdoutFrames[1].Data)) + require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code) +} + func prepareForExec(t *testing.T) (*client.Client, string) { devClient, _, _ := devcontroller.StartIntegrationTestEnvironment(t) @@ -164,3 +325,43 @@ func readFrame(t *testing.T, wsConn *websocket.Conn) *execstream.Frame { return &frame } + +func readFramesUntilExit(t *testing.T, wsConn *websocket.Conn) []*execstream.Frame { + t.Helper() + + var frames []*execstream.Frame + + for { + frame := readFrame(t, wsConn) + if frame.Type == execstream.FrameTypeNoMoreHistory { + continue + } + + frames = append(frames, frame) + if frame.Type == execstream.FrameTypeExit { + return frames + } + } +} + +func framesByType(frames []*execstream.Frame, frameType execstream.FrameType) []*execstream.Frame { + var result []*execstream.Frame + + for _, frame := range frames { + if frame.Type == frameType { + result = append(result, frame) + } + } + + return result +} + +func frameTypes(frames []*execstream.Frame) []execstream.FrameType { + var result []execstream.FrameType + + for _, frame := range frames { + result = append(result, frame.Type) + } + + return result +} diff --git a/pkg/client/vms.go b/pkg/client/vms.go index 3bcbb6e..7b73544 100644 --- a/pkg/client/vms.go +++ b/pkg/client/vms.go @@ -47,6 +47,13 @@ type EventsPageOptions struct { Cursor string } +type ExecSessionOptions struct { + Command string + Stdin bool + WaitSeconds uint16 + Session string +} + func (service *VMsService) Create(ctx context.Context, vm *v1.VM) error { err := service.client.request(ctx, http.MethodPost, "vms", vm, nil, nil) @@ -164,12 +171,33 @@ func (service *VMsService) Exec( stdin bool, waitSeconds uint16, ) (*websocket.Conn, error) { + return service.ExecSession(ctx, name, ExecSessionOptions{ + Command: command, + Stdin: stdin, + WaitSeconds: waitSeconds, + }) +} + +func (service *VMsService) ExecSession( + ctx context.Context, + name string, + options ExecSessionOptions, +) (*websocket.Conn, error) { + params := map[string]string{ + "wait": strconv.FormatUint(uint64(options.WaitSeconds), 10), + } + if options.Command != "" { + params["command"] = options.Command + } + if options.Stdin { + params["stdin"] = strconv.FormatBool(true) + } + if options.Session != "" { + params["session"] = options.Session + } + return service.client.wsRequestRaw(ctx, fmt.Sprintf("vms/%s/exec", url.PathEscape(name)), - map[string]string{ - "command": command, - "stdin": strconv.FormatBool(stdin), - "wait": strconv.FormatUint(uint64(waitSeconds), 10), - }) + params) } func (service *VMsService) IP(ctx context.Context, name string, waitSeconds uint16) (string, error) { diff --git a/pkg/client/vms_test.go b/pkg/client/vms_test.go new file mode 100644 index 0000000..79f3394 --- /dev/null +++ b/pkg/client/vms_test.go @@ -0,0 +1,40 @@ +package client + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/coder/websocket" + "github.com/stretchr/testify/require" +) + +func TestExecSessionBuildsReconnectableQuery(t *testing.T) { + var query map[string][]string + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + query = request.URL.Query() + + conn, err := websocket.Accept(writer, request, nil) + require.NoError(t, err) + defer conn.CloseNow() + })) + defer server.Close() + + devClient, err := New(WithAddress(server.URL)) + require.NoError(t, err) + + conn, err := devClient.VMs().ExecSession(t.Context(), "vm", ExecSessionOptions{ + Command: "echo hello", + Stdin: true, + WaitSeconds: 7, + Session: "resume-me", + }) + require.NoError(t, err) + defer conn.CloseNow() + + require.Equal(t, []string{"echo hello"}, query["command"]) + require.Equal(t, []string{"true"}, query["stdin"]) + require.Equal(t, []string{"7"}, query["wait"]) + require.Equal(t, []string{"resume-me"}, query["session"]) +}