From a079358ed761c1d6d64d1e134530a412c2eee45b Mon Sep 17 00:00:00 2001 From: chruffins <23645059+chruffins@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:10:17 +0000 Subject: [PATCH] Replay PTY output to late attaches on top of main Drain PTY output for the process lifetime into a replay buffer so a client that attaches after output was written still receives it. - Attach replays buffered output, then forwards live output; the exit event is sent only after the attached sink has flushed. - Detach waits for the forwarder to leave any in-flight sink call so a stale forwarder cannot move the next attachment's cursor. - The waiter stops the PTY reader through its stop channel after a bounded drain before closing the master. The reader polls the raw fd and spins on a closed one. --- server/cmd/api/api/process.go | 37 ++-- server/cmd/api/api/process_attach_test.go | 73 ++++++++ server/cmd/api/api/process_runner.go | 24 +++ server/cmd/api/api/process_test.go | 27 +++ server/cmd/api/api/pty_output.go | 152 +++++++++++++++++ server/cmd/api/api/pty_output_test.go | 195 ++++++++++++++++++++++ 6 files changed, 490 insertions(+), 18 deletions(-) create mode 100644 server/cmd/api/api/process_attach_test.go create mode 100644 server/cmd/api/api/pty_output.go create mode 100644 server/cmd/api/api/pty_output_test.go diff --git a/server/cmd/api/api/process.go b/server/cmd/api/api/process.go index 4c36bf36..b60f96fc 100644 --- a/server/cmd/api/api/process.go +++ b/server/cmd/api/api/process.go @@ -38,6 +38,9 @@ type processHandle struct { stdout io.ReadCloser stderr io.ReadCloser ptyFile *os.File + ptyOut ptyOutput + ptyDrained chan struct{} + ptyStop chan struct{} isTTY bool output *processOutput reason string @@ -710,32 +713,28 @@ func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Reques } }() - // Goroutine: read from PTY (stdout), write to WebSocket - wg.Add(1) - go func() { - defer wg.Done() - err := ptyio.ReadPTYToWriter(h.ptyFile, func(data []byte) error { - select { - case writeCh <- wsWriteOp{msgType: websocket.MessageBinary, data: data}: - return nil - case <-done: - return io.EOF // Signal to stop reading - } - }, done) - if err != nil { - log.Error("pty read error", "err", err) + // Replay output already drained from the PTY, then forward live output. The + // ptyOutput lock keeps replay ahead of any writes that arrive concurrently. + h.ptyOut.attach(func(data []byte) error { + select { + case writeCh <- wsWriteOp{msgType: websocket.MessageBinary, data: data}: + return nil + case <-done: + return io.EOF } - shutdown() - }() + }) - // Goroutine: watch for process exit and send exit code. + // Goroutine: watch for process exit and send exit code after PTY output drains. // This must send the exit code BEFORE triggering shutdown so the writer can deliver it. wg.Add(1) go func() { defer wg.Done() select { case <-h.doneCh: - // Process exited - send exit code before shutdown + // The process waiter closes doneCh only after PTY output has drained. + // Flush the attached sink before enqueueing the exit event so replayed + // and live output stay ahead of it even when the socket was backpressured. + h.ptyOut.flush() h.mu.RLock() exitCode := h.exitCode h.mu.RUnlock() @@ -760,6 +759,8 @@ func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Reques // Wait for shutdown signal <-done + h.ptyOut.detach() + // Wait for all goroutines to complete before clearing attachActive. // This prevents a race where a new client attaches while goroutines are still running. wg.Wait() diff --git a/server/cmd/api/api/process_attach_test.go b/server/cmd/api/api/process_attach_test.go new file mode 100644 index 00000000..2aee6bb9 --- /dev/null +++ b/server/cmd/api/api/process_attach_test.go @@ -0,0 +1,73 @@ +package api + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" + oapi "github.com/kernel/kernel-images/server/lib/oapi" + "github.com/kernel/kernel-images/server/lib/scaletozero" + "github.com/kernel/kernel-images/server/lib/wsdrain" + "github.com/stretchr/testify/require" +) + +func TestProcessAttachReplaysOutputThenExits(t *testing.T) { + t.Parallel() + ctx := context.Background() + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} + + // The first line is written before any client attaches, so it must come back as replay. + args := []string{"-c", "echo first; sleep 0.3; echo second"} + allocateTTY := true + body := &oapi.ProcessSpawnRequest{Command: "sh", Args: &args, AllocateTty: &allocateTTY} + spawnResp, err := svc.ProcessSpawn(ctx, oapi.ProcessSpawnRequestObject{Body: body}) + require.NoError(t, err) + s200, ok := spawnResp.(oapi.ProcessSpawn200JSONResponse) + require.True(t, ok, "unexpected spawn resp: %T", spawnResp) + id := s200.ProcessId.String() + + time.Sleep(100 * time.Millisecond) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + svc.HandleProcessAttachWS(w, r, id, wsdrain.New()) + })) + defer srv.Close() + + conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(srv.URL, "http"), nil) + require.NoError(t, err) + defer conn.CloseNow() + + var out strings.Builder + exitCode := -1 + readCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + for { + msgType, data, err := conn.Read(readCtx) + require.NoError(t, err, "attach ended before exit message; output so far %q", out.String()) + if msgType == websocket.MessageText { + var ctrl struct { + Type string `json:"type"` + ExitCode *int `json:"exitCode"` + } + require.NoError(t, json.Unmarshal(data, &ctrl)) + if ctrl.Type == "exit" { + require.NotNil(t, ctrl.ExitCode) + exitCode = *ctrl.ExitCode + break + } + continue + } + out.Write(data) + } + + require.Equal(t, 0, exitCode) + got := out.String() + require.Contains(t, got, "first") + require.Contains(t, got, "second") + require.Less(t, strings.Index(got, "first"), strings.Index(got, "second")) +} diff --git a/server/cmd/api/api/process_runner.go b/server/cmd/api/api/process_runner.go index db818537..27ec82e8 100644 --- a/server/cmd/api/api/process_runner.go +++ b/server/cmd/api/api/process_runner.go @@ -13,6 +13,7 @@ import ( "github.com/creack/pty" "github.com/google/uuid" "github.com/kernel/kernel-images/server/lib/logger" + "github.com/kernel/kernel-images/server/lib/ptyio" ) const processOutputDrainTimeout = time.Second @@ -57,6 +58,18 @@ func (s *ApiService) startProcess(ctx context.Context, cmd *exec.Cmd, tty, stdin return nil, err } h.pid = cmd.Process.Pid + if tty { + h.ptyDrained = make(chan struct{}) + h.ptyStop = make(chan struct{}) + // PTY output is drained for the process lifetime so output remains available + // when a client attaches after the process has written it. + go func() { + defer close(h.ptyDrained) + if err := ptyio.ReadPTYToWriter(h.ptyFile, h.ptyOut.write, h.ptyStop); err != nil { + logger.FromContext(ctx).Error("pty read error", "err", err) + } + }() + } s.procMu.Lock() if s.procs == nil { s.procs = make(map[string]*processHandle) @@ -190,6 +203,17 @@ func (h *processHandle) run(ctx, lifecycle context.Context, timeoutSec *int, rel drain = nil } } + if h.isTTY { + // The reader polls the raw fd, so stop it before closing the master. + // Closing first leaves it polling a dead descriptor forever. A background + // child can hold the slave open after the shell exits, so the wait is bounded. + select { + case <-h.ptyDrained: + case <-time.After(ptyDrainWait): + close(h.ptyStop) + <-h.ptyDrained + } + } h.closeOutputPipes() if h.stdin != nil { h.stdin.Close() diff --git a/server/cmd/api/api/process_test.go b/server/cmd/api/api/process_test.go index 4893d920..a1c89723 100644 --- a/server/cmd/api/api/process_test.go +++ b/server/cmd/api/api/process_test.go @@ -286,3 +286,30 @@ func TestBuildCmd_AsRootAndAsUserConflict(t *testing.T) { _, err := buildCmd(body) require.Error(t, err, "expected error when both as_root and as_user are set") } + +func TestProcessTTYExitWithBackgroundChildHoldingPTYDoesNotBlock(t *testing.T) { + t.Parallel() + ctx := context.Background() + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} + + // The background sleep ignores SIGHUP, so it keeps the PTY slave open after the + // shell exits and the reader never sees EOF. + args := []string{"-c", "trap \"\" HUP; sleep 30 & exit 0"} + allocateTTY := true + body := &oapi.ProcessSpawnRequest{Command: "sh", Args: &args, AllocateTty: &allocateTTY} + spawnResp, err := svc.ProcessSpawn(ctx, oapi.ProcessSpawnRequestObject{Body: body}) + require.NoError(t, err) + s200, ok := spawnResp.(oapi.ProcessSpawn200JSONResponse) + require.True(t, ok, "unexpected spawn resp: %T", spawnResp) + + svc.procMu.Lock() + h := svc.procs[s200.ProcessId.String()] + svc.procMu.Unlock() + require.NotNil(t, h) + + select { + case <-h.doneCh: + case <-time.After(10 * time.Second): + t.Fatal("process waiter did not finish after the PTY drain deadline") + } +} diff --git a/server/cmd/api/api/pty_output.go b/server/cmd/api/api/pty_output.go new file mode 100644 index 00000000..db2d574b --- /dev/null +++ b/server/cmd/api/api/pty_output.go @@ -0,0 +1,152 @@ +package api + +import ( + "sync" + "time" +) + +const ( + ptyReplayBytes = 64 * 1024 + ptyDrainWait = 2 * time.Second +) + +type ptyOutput struct { + mu sync.Mutex + buf []byte + start uint64 + end uint64 + sink func([]byte) error + stop chan struct{} + exited chan struct{} + changed chan struct{} + delivered uint64 +} + +func (o *ptyOutput) write(data []byte) error { + o.mu.Lock() + defer o.mu.Unlock() + if len(o.buf) == 0 { + o.buf = make([]byte, ptyReplayBytes) + } + + if len(data) > ptyReplayBytes { + o.start = o.end + uint64(len(data)-ptyReplayBytes) + data = data[len(data)-ptyReplayBytes:] + } + for len(data) > 0 { + index := int(o.end % ptyReplayBytes) + n := min(len(data), ptyReplayBytes-index) + copy(o.buf[index:index+n], data[:n]) + o.end += uint64(n) + data = data[n:] + if o.end-o.start > ptyReplayBytes { + o.start = o.end - ptyReplayBytes + } + } + o.signalLocked() + return nil +} + +func (o *ptyOutput) attach(sink func([]byte) error) { + o.mu.Lock() + defer o.mu.Unlock() + if o.changed == nil { + o.changed = make(chan struct{}) + } + o.sink = sink + o.stop = make(chan struct{}) + o.exited = make(chan struct{}) + o.delivered = o.start + go o.forward(o.stop, o.exited) +} + +func (o *ptyOutput) forward(stop <-chan struct{}, exited chan<- struct{}) { + defer close(exited) + for { + o.mu.Lock() + if o.delivered < o.start { + o.delivered = o.start + } + if o.sink == nil { + o.mu.Unlock() + return + } + if o.delivered < o.end { + count := min(o.end-o.delivered, uint64(ptyReplayBytes)) + data := o.copyLocked(o.delivered, count) + sink := o.sink + o.mu.Unlock() + if err := sink(data); err != nil { + o.mu.Lock() + o.sink = nil + o.signalLocked() + o.mu.Unlock() + return + } + o.mu.Lock() + if o.sink == nil { + o.signalLocked() + o.mu.Unlock() + return + } + o.delivered += uint64(len(data)) + o.signalLocked() + o.mu.Unlock() + continue + } + changed := o.changed + o.mu.Unlock() + select { + case <-changed: + case <-stop: + return + } + } +} + +func (o *ptyOutput) copyLocked(cursor, count uint64) []byte { + data := make([]byte, count) + index := int(cursor % ptyReplayBytes) + first := min(len(data), ptyReplayBytes-index) + copy(data, o.buf[index:index+first]) + copy(data[first:], o.buf[:len(data)-first]) + return data +} + +func (o *ptyOutput) flush() { + o.mu.Lock() + target := o.end + for o.sink != nil && o.delivered < target { + changed := o.changed + o.mu.Unlock() + <-changed + o.mu.Lock() + target = o.end + } + o.mu.Unlock() +} + +func (o *ptyOutput) detach() { + o.mu.Lock() + var exited chan struct{} + if o.stop != nil { + close(o.stop) + o.stop = nil + exited = o.exited + } + o.sink = nil + o.signalLocked() + o.mu.Unlock() + // Wait for the forwarder to leave any in-flight sink call so it cannot move + // the cursor of the next attachment. + if exited != nil { + <-exited + } +} + +func (o *ptyOutput) signalLocked() { + if o.changed != nil { + close(o.changed) + } + o.changed = make(chan struct{}) +} diff --git a/server/cmd/api/api/pty_output_test.go b/server/cmd/api/api/pty_output_test.go new file mode 100644 index 00000000..34cd6c69 --- /dev/null +++ b/server/cmd/api/api/pty_output_test.go @@ -0,0 +1,195 @@ +package api + +import ( + "bytes" + "sync" + "testing" + "time" +) + +func TestPTYOutputRetainsLatestBytes(t *testing.T) { + var output ptyOutput + output.write(bytes.Repeat([]byte("a"), ptyReplayBytes-2)) + output.write([]byte("tail")) + + var replay []byte + output.attach(func(data []byte) error { + replay = append(replay, data...) + return nil + }) + output.flush() + + if len(replay) != ptyReplayBytes { + t.Fatalf("replay length = %d, want %d", len(replay), ptyReplayBytes) + } + if !bytes.Equal(replay[len(replay)-4:], []byte("tail")) { + t.Fatalf("replay tail = %q, want %q", replay[len(replay)-4:], "tail") + } +} + +func TestPTYOutputReplayAndLiveOutputPrecedeExit(t *testing.T) { + var output ptyOutput + if err := output.write([]byte("replay")); err != nil { + t.Fatal(err) + } + + var received bytes.Buffer + output.attach(func(data []byte) error { + _, _ = received.Write(data) + return nil + }) + if err := output.write([]byte("live")); err != nil { + t.Fatal(err) + } + output.flush() + _, _ = received.WriteString("exit") + + if got, want := received.String(), "replayliveexit"; got != want { + t.Fatalf("received %q, want %q", got, want) + } +} + +func TestPTYOutputContinuesDrainingDuringSinkBackpressure(t *testing.T) { + var output ptyOutput + entered := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + var received bytes.Buffer + output.attach(func(data []byte) error { + once.Do(func() { + close(entered) + <-release + }) + _, _ = received.Write(data) + return nil + }) + if err := output.write([]byte("blocked")); err != nil { + t.Fatal(err) + } + <-entered + + writesDone := make(chan struct{}) + go func() { + for i := 0; i < 4; i++ { + _ = output.write(bytes.Repeat([]byte("x"), ptyReplayBytes/2)) + } + _ = output.write([]byte("tail")) + close(writesDone) + }() + select { + case <-writesDone: + case <-time.After(time.Second): + t.Fatal("PTY writes blocked behind the WebSocket sink") + } + + close(release) + output.flush() + if !bytes.Contains(received.Bytes(), []byte("tail")) { + t.Fatal("buffered PTY tail was not forwarded after sink recovered") + } +} + +func TestPTYOutputFullBufferDoesNotReallocateOnEachWrite(t *testing.T) { + var output ptyOutput + if err := output.write(bytes.Repeat([]byte("a"), ptyReplayBytes)); err != nil { + t.Fatal(err) + } + if len(output.buf) != ptyReplayBytes { + t.Fatalf("buffer length = %d, want %d", len(output.buf), ptyReplayBytes) + } + firstByte := &output.buf[0] + + for i := 0; i < 10; i++ { + if err := output.write([]byte("x")); err != nil { + t.Fatal(err) + } + if &output.buf[0] != firstByte { + t.Fatalf("buffer reallocated on write %d", i) + } + } + if got := output.buf[(output.end-1)%ptyReplayBytes]; got != 'x' { + t.Fatalf("last buffered byte = %q, want %q", got, 'x') + } +} + +func TestPTYOutputDeliversWritesAfterAttachBeforeFirstOutput(t *testing.T) { + var output ptyOutput + var mu sync.Mutex + var received bytes.Buffer + output.attach(func(data []byte) error { + mu.Lock() + defer mu.Unlock() + _, _ = received.Write(data) + return nil + }) + // Let the forwarder reach its wait before the first write. + time.Sleep(10 * time.Millisecond) + if err := output.write([]byte("first")); err != nil { + t.Fatal(err) + } + + flushed := make(chan struct{}) + go func() { + output.flush() + close(flushed) + }() + select { + case <-flushed: + case <-time.After(time.Second): + t.Fatal("output written after an early attach was never delivered") + } + mu.Lock() + defer mu.Unlock() + if got := received.String(); got != "first" { + t.Fatalf("received %q, want %q", got, "first") + } +} + +func TestPTYOutputDetachWaitsForInFlightSink(t *testing.T) { + var output ptyOutput + entered := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + output.attach(func(data []byte) error { + once.Do(func() { + close(entered) + <-release + }) + return nil + }) + if err := output.write([]byte("a")); err != nil { + t.Fatal(err) + } + <-entered + + detached := make(chan struct{}) + go func() { + output.detach() + close(detached) + }() + select { + case <-detached: + t.Fatal("detach returned while the previous attachment's sink was still running") + case <-time.After(50 * time.Millisecond): + } + + close(release) + select { + case <-detached: + case <-time.After(time.Second): + t.Fatal("detach did not return after the sink finished") + } + + var received bytes.Buffer + output.attach(func(data []byte) error { + _, _ = received.Write(data) + return nil + }) + if err := output.write([]byte("b")); err != nil { + t.Fatal(err) + } + output.flush() + if got := received.String(); got != "ab" { + t.Fatalf("reattached sink received %q, want %q", got, "ab") + } +}