diff --git a/internal/session/client_registry.go b/internal/session/client_registry.go index 7441de7..f95eba7 100644 --- a/internal/session/client_registry.go +++ b/internal/session/client_registry.go @@ -24,9 +24,14 @@ type serverMessage struct { } type attachClient struct { - conn net.Conn - out chan serverMessage - done chan struct{} + conn net.Conn + out chan serverMessage + done chan struct{} + // flush is closed by shutdown to ask this client's writer to drain what is + // queued and then close the connection itself. done means "stop now"; + // flush means "finish, then stop". + flush chan struct{} + flushOnce sync.Once version protocolVersion id string requestedRole clientRole @@ -49,6 +54,17 @@ func (client *attachClient) drop() { }) } +// requestFlush asks the client's writer to drain its queue and close. It is a +// no-op for a client that has no writer running; shutdown falls back to drop +// once its flush window expires. +func (client *attachClient) requestFlush() { + client.flushOnce.Do(func() { + if client.flush != nil { + close(client.flush) + } + }) +} + type clientRegistration struct { requestedRole clientRole hello clientHello diff --git a/internal/session/host.go b/internal/session/host.go index 03c49b0..f1024a5 100644 --- a/internal/session/host.go +++ b/internal/session/host.go @@ -55,6 +55,12 @@ const killGrace = 1500 * time.Millisecond // disconnected rather than allowed to stall the session. const attachBufFrames = 512 +// shutdownFlushWindow bounds how long teardown waits for attached viewers to +// receive the agent's final output. Long enough for a queued screen on a slow +// link, short enough that a viewer which stopped reading cannot hold the host +// open. +const shutdownFlushWindow = 250 * time.Millisecond + const ( markClosedRetryWindow = 2 * time.Second markClosedRetryBase = 25 * time.Millisecond @@ -464,7 +470,8 @@ func (h *host) handleAttach(conn net.Conn, br *bufio.Reader, req request) { return } client := &attachClient{ - conn: conn, out: make(chan serverMessage, attachBufFrames), done: make(chan struct{}), version: version, + conn: conn, out: make(chan serverMessage, attachBufFrames), done: make(chan struct{}), + flush: make(chan struct{}), version: version, fallback: version == protocolV1 && !req.versionPresent, } attachResponse, err := h.registerAttachClient(client, registration) @@ -491,6 +498,9 @@ func (h *host) attachWriter(client *attachClient) { select { case <-client.done: return + case <-client.flush: + h.flushClient(client) + return case message := <-client.out: if err := h.writeServerMessage(client, message); err != nil { h.dropClientReason(client, "connection_write") @@ -500,6 +510,28 @@ func (h *host) attachWriter(client *attachClient) { } } +// flushClient writes what the shutting-down client has already queued and then +// closes the connection itself. The write deadline bounds a viewer that has +// stopped reading, so a stalled socket cannot hold up host teardown. +func (h *host) flushClient(client *attachClient) { + defer client.drop() + if client.conn != nil { + if err := client.conn.SetWriteDeadline(time.Now().Add(shutdownFlushWindow)); err != nil { + return + } + } + for { + select { + case message := <-client.out: + if err := h.writeServerMessage(client, message); err != nil { + return + } + default: + return + } + } +} + func (h *host) writeServerMessage(client *attachClient, message serverMessage) error { if client.version == protocolV1 { if message.kind != serverFramePTY { @@ -659,6 +691,17 @@ func (h *host) signalChild(sig syscall.Signal) { // closed (the native replacement for the tmux session-closed hook), tell any // attached clients, and remove the runtime files. func (h *host) shutdown(exitCode int) { + h.shutdownClients() + providerID := readProviderIdentityHandoff(h.dir, h.name, h.providerIdentityFile) + if err := removeSessionFiles(h.dir, h.name); err != nil { + log.Warn("remove session files failed", "session", h.name, "error", err) + } + h.recordExit(exitCode, providerID) +} + +// shutdownClients releases every attached client, giving each writer a bounded +// window to put the agent's last output on the wire first. +func (h *host) shutdownClients() { h.mu.Lock() clients := h.registry.drain() h.mu.Unlock() @@ -667,13 +710,23 @@ func (h *host) shutdown(exitCode int) { Event: "attach.lifecycle", Session: h.name, ClientID: client.id, Protocol: int(client.version), Role: string(client.assignedRole), Reason: "host_shutdown", }) - client.drop() - } - providerID := readProviderIdentityHandoff(h.dir, h.name, h.providerIdentityFile) - if err := removeSessionFiles(h.dir, h.name); err != nil { - log.Warn("remove session files failed", "session", h.name, "error", err) + client.requestFlush() + } + // pumpPTY has already returned, so everything the agent ever wrote is + // queued; each writer only needs the chance to put it on the wire. Closing + // the connection first — as this used to — truncated the agent's last + // screen, and cutting a frame in half turned a clean exit into + // "attach output: unexpected EOF" on the viewer. + flushBy := time.Now().Add(shutdownFlushWindow) + for _, client := range clients { + select { + case <-client.done: + case <-time.After(time.Until(flushBy)): + // No writer is running for this client, or it is wedged behind a + // stalled socket. Either way the deadline is the whole budget. + client.drop() + } } - h.recordExit(exitCode, providerID) } func (h *host) recordExit(exitCode int, providerID string) { diff --git a/internal/session/host_shutdown_flush_test.go b/internal/session/host_shutdown_flush_test.go new file mode 100644 index 0000000..f66fedd --- /dev/null +++ b/internal/session/host_shutdown_flush_test.go @@ -0,0 +1,153 @@ +package session + +import ( + "bufio" + "bytes" + "net" + "testing" + "time" +) + +// Output still queued when the host tears down must reach the viewer. The +// distinction the writer has to honour: `done` means stop now, `flush` means +// finish first. Dropping the client outright — as shutdown used to — truncated +// the agent's last screen, and a half-written frame surfaced on the viewer as +// "attach output: unexpected EOF" instead of a clean session end. +func TestWriterFlushesQueuedOutputOnShutdownSignal(t *testing.T) { + // Given three frames queued and a shutdown already requested. + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := newFlushTestClient(t, h, server) + payloads := [][]byte{[]byte("first\r\n"), []byte("second\r\n"), []byte("agent final line\r\n")} + for _, payload := range payloads { + if !h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: payload}) { + t.Fatal("queueing the agent's output failed") + } + } + client.requestFlush() + + // When the writer runs. + go h.attachWriter(client) + + // Then every queued frame arrives before the connection closes. + if err := viewer.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil { + t.Fatal(err) + } + reader := bufio.NewReader(viewer) + for i, want := range payloads { + kind, payload, err := readFrame(reader) + if err != nil { + t.Fatalf("frame %d lost on shutdown: %v", i, err) + } + if kind != serverFramePTY || !bytes.Equal(payload, want) { + t.Fatalf("frame %d = %d %q, want %d %q", i, kind, payload, serverFramePTY, want) + } + } + waitClosed(t, client) +} + +// The contrast that makes the distinction meaningful: a client told to stop +// now discards what is queued. +func TestWriterStopsImmediatelyOnDone(t *testing.T) { + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := newFlushTestClient(t, h, server) + if err := viewer.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: []byte("discarded")}) + client.drop() + + go h.attachWriter(client) + + if _, _, err := readFrame(bufio.NewReader(viewer)); err == nil { + t.Fatal("a dropped client still wrote its queue") + } +} + +func newFlushTestClient(t *testing.T, h *host, conn net.Conn) *attachClient { + t.Helper() + client := &attachClient{ + conn: conn, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + client.ready = true + return client +} + +func waitClosed(t *testing.T, client *attachClient) { + t.Helper() + select { + case <-client.done: + case <-time.After(3 * time.Second): + t.Fatal("client was not closed after the flush") + } +} + +// A viewer that has stopped reading must not hold teardown open. +func TestShutdownDoesNotWaitForeverOnAStalledViewer(t *testing.T) { + // Given a client whose peer never reads. + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := &attachClient{ + conn: server, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + client.ready = true + go h.attachWriter(client) + h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: bytes.Repeat([]byte("x"), 4096)}) + + // When the host tears down. + start := time.Now() + h.shutdownClients() + + // Then it returns within the flush budget and the client is closed. + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Fatalf("shutdown blocked for %s on a stalled viewer", elapsed) + } + select { + case <-client.done: + default: + t.Fatal("stalled client was not force-closed") + } +} + +// A registered client whose writer never started still has to be closed. +func TestShutdownClosesClientsWithoutAWriter(t *testing.T) { + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := &attachClient{ + conn: server, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + + h.shutdownClients() + + select { + case <-client.done: + default: + t.Fatal("client with no writer was left open") + } +}