internal/daemon: batch and backpressure PTY output

Author: seruman

internal/daemon/server_attach.go

@@ -235,6 +235,7 @@ func (s *Server) watchSession(sess *Session) {
 		empty := len(s.sessions) == 0
 		s.mu.Unlock()
 		if s.autoExit && empty {
+			sess.waitClients()
 			slog.Info("auto-exit: last session ended, shutting down")
 			s.Shutdown()
 		}

internal/daemon/session.go

@@ -15,7 +15,7 @@ import (
 
 var feedPool = sync.Pool{
 	New: func() any {
-		b := make([]byte, 32*1024)
+		b := make([]byte, ptyBatchSize)
 		return &b
 	},
 }
@@ -45,6 +45,9 @@ type sessionClient struct {
 	version   string
 	readOnly  bool
 	outCh     chan protocol.Message
+	ready     chan<- struct{}
+	writeDone chan struct{}
+	final     protocol.Message
 }
 
 type sessionAction interface {
@@ -107,22 +110,26 @@ type Session struct {
 	PID       uint32
 	CreatedAt time.Time
 
-	ptmx    *os.File
-	cmd     *exec.Cmd
-	term    *terminalState
-	feedCh  chan feedItem
-	tempDir string
+	ptmx     *os.File
+	cmd      *exec.Cmd
+	term     *terminalState
+	feedCh   chan feedItem
+	feedDone chan struct{}
+	ptyDone  chan struct{}
+	tempDir  string
 
-	actions  chan sessionAction
-	ptyOut   chan []byte
-	done     chan struct{}
-	exitCode int32
+	actions     chan sessionAction
+	ptyOut      chan []byte
+	clientReady chan struct{}
+	done        chan struct{}
+	exitCode    int32
 
 	// sizeVal packs cols|rows as (cols<<16)|rows for lock-free reads.
 	sizeVal atomic.Uint32
 
-	resizePolicy config.ResizePolicy
-	ctx          context.Context
+	resizePolicy  config.ResizePolicy
+	clientWriters sync.WaitGroup
+	ctx           context.Context
 }
 
 func (s *Session) size() (uint16, uint16) {

internal/daemon/session_clients.go

@@ -6,22 +6,36 @@ import (
 	"log/slog"
 	"math"
 	"syscall"
-	"time"
 
 	"code.selman.me/hauntty/internal/config"
 	"code.selman.me/hauntty/internal/protocol"
 	"github.com/creack/pty"
 )
 
-const (
-	sessionClientOutBufferSize = 256
-	slowClientGracePeriod      = 100 * time.Millisecond
-)
+const sessionClientOutBufferSize = 256
 
 func (c *sessionClient) writeLoop() {
+	defer func() {
+		close(c.writeDone)
+		select {
+		case c.ready <- struct{}{}:
+		default:
+		}
+	}()
+
 	for msg := range c.outCh {
+		select {
+		case c.ready <- struct{}{}:
+		default:
+		}
 		if err := c.conn.WriteMessage(msg); err != nil {
-			break
+			_ = c.closeConn()
+			return
+		}
+	}
+	if c.final != nil {
+		if err := c.conn.WriteMessage(c.final); err != nil {
+			_ = c.closeConn()
 		}
 	}
 }
@@ -169,32 +183,37 @@ func removeClient(clients []*sessionClient, target *sessionClient) []*sessionCli
 	return clients
 }
 
-func broadcastOutput(clients []*sessionClient, name string, msg *protocol.Output) []*sessionClient {
-	i := 0
+func queueOutput(clients []*sessionClient, msg *protocol.Output) []*sessionClient {
+	var pending []*sessionClient
 	for _, c := range clients {
 		select {
-		case c.outCh <- msg:
-			clients[i] = c
-			i++
+		case <-c.writeDone:
 			continue
 		default:
 		}
-
-		timer := time.NewTimer(slowClientGracePeriod)
 		select {
 		case c.outCh <- msg:
-			if !timer.Stop() {
-				<-timer.C
-			}
-			clients[i] = c
-			i++
-		case <-timer.C:
-			slog.Debug("evicting slow client", "session", name, "grace", slowClientGracePeriod)
+		default:
+			pending = append(pending, c)
+		}
+	}
+	return pending
+}
+
+func pruneFinishedClients(clients, pending []*sessionClient) ([]*sessionClient, []*sessionClient, bool) {
+	kept := clients[:0]
+	changed := false
+	for _, c := range clients {
+		select {
+		case <-c.writeDone:
+			pending = removeClient(pending, c)
 			close(c.outCh)
-			_ = c.closeConn()
+			changed = true
+		default:
+			kept = append(kept, c)
 		}
 	}
-	return clients[:i]
+	return kept, pending, changed
 }
 
 func notifyClientsChanged(clients []*sessionClient, sizeFn func() (uint16, uint16)) {