fix(host): flush queued output before closing clients on shutdown

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

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) {

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")
+	}
+}