coder
diff --git a/‎agent/immortalstreams/backedpipe/backed_pipe.go
Lines changed: 76 additions & 73 deletions b/‎agent/immortalstreams/backedpipe/backed_pipe.go
Lines changed: 76 additions & 73 deletions
@@ -5,6 +5,7 @@ import (
 	"io"
 	"sync"
 
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/singleflight"
 	"golang.org/x/xerrors"
 )
@@ -23,8 +24,8 @@ const (
 	DefaultBufferSize = 64 * 1024 * 1024
 )
 
-// ReconnectFunc is called when the BackedPipe needs to establish a new connection.
-// It should:
+// Reconnector is an interface for establishing connections when the BackedPipe needs to reconnect.
+// Implementations should:
 // 1. Establish a new connection to the remote side
 // 2. Exchange sequence numbers with the remote side
 // 3. Return the new connection and the remote's current sequence number
@@ -34,7 +35,9 @@ const (
 //
 // The returned readerSeqNum should be the remote side's current sequence number,
 // which indicates where the local reader should resume from.
-type ReconnectFunc func(ctx context.Context, writerSeqNum uint64) (conn io.ReadWriteCloser, readerSeqNum uint64, err error)
+type Reconnector interface {
+	Reconnect(ctx context.Context, writerSeqNum uint64) (conn io.ReadWriteCloser, readerSeqNum uint64, err error)
+}
 
 // BackedPipe provides a reliable bidirectional byte stream over unreliable network connections.
 // It orchestrates a BackedReader and BackedWriter to provide transparent reconnection
@@ -45,44 +48,39 @@ type BackedPipe struct {
 	mu          sync.RWMutex
 	reader      *BackedReader
 	writer      *BackedWriter
-	reconnectFn ReconnectFunc
+	reconnector Reconnector
 	conn        io.ReadWriteCloser
 	connected   bool
 	closed      bool
 
 	// Reconnection state
 	reconnecting bool
 
-	// Error channel for receiving connection errors from reader/writer
-	errorChan chan error
+	// Error channels for receiving connection errors from reader/writer separately
+	readerErrorChan chan error
+	writerErrorChan chan error
 
 	// singleflight group to dedupe concurrent ForceReconnect calls
 	sf singleflight.Group
 }
 
-// NewBackedPipe creates a new BackedPipe with default options and the specified reconnect function.
+// NewBackedPipe creates a new BackedPipe with default options and the specified reconnector.
 // The pipe starts disconnected and must be connected using Connect().
-func NewBackedPipe(ctx context.Context, reconnectFn ReconnectFunc) *BackedPipe {
+func NewBackedPipe(ctx context.Context, reconnector Reconnector) *BackedPipe {
 	pipeCtx, cancel := context.WithCancel(ctx)
 
-	errorChan := make(chan error, 2) // Buffer for reader and writer errors
+	readerErrorChan := make(chan error, 1) // Buffer for reader errors
+	writerErrorChan := make(chan error, 1) // Buffer for writer errors
 	bp := &BackedPipe{
-		ctx:         pipeCtx,
-		cancel:      cancel,
-		reader:      NewBackedReader(),
-		writer:      NewBackedWriter(DefaultBufferSize, errorChan),
-		reconnectFn: reconnectFn,
-		errorChan:   errorChan,
+		ctx:             pipeCtx,
+		cancel:          cancel,
+		reader:          NewBackedReader(readerErrorChan),
+		writer:          NewBackedWriter(DefaultBufferSize, writerErrorChan),
+		reconnector:     reconnector,
+		readerErrorChan: readerErrorChan,
+		writerErrorChan: writerErrorChan,
 	}
 
-	// Set up error callback for reader only (writer uses error channel directly)
-	bp.reader.SetErrorCallback(func(err error) {
-		select {
-		case bp.errorChan <- err:
-		case <-bp.ctx.Done():
-		}
-	})
-
 	// Start error handler goroutine
 	go bp.handleErrors()
 
@@ -109,16 +107,7 @@ func (bp *BackedPipe) Connect() error {
 
 // Read implements io.Reader by delegating to the BackedReader.
 func (bp *BackedPipe) Read(p []byte) (int, error) {
-	bp.mu.RLock()
-	reader := bp.reader
-	closed := bp.closed
-	bp.mu.RUnlock()
-
-	if closed {
-		return 0, io.EOF
-	}
-
-	return reader.Read(p)
+	return bp.reader.Read(p)
 }
 
 // Write implements io.Writer by delegating to the BackedWriter.
@@ -147,32 +136,38 @@ func (bp *BackedPipe) Close() error {
 	bp.closed = true
 	bp.cancel() // Cancel main context
 
-	// Close underlying components
-	var readerErr, writerErr, connErr error
+	// Close all components in parallel to avoid deadlocks
+	//
+	// IMPORTANT: The connection must be closed first to unblock any
+	// readers or writers that might be holding the mutex on Read/Write
+	var g errgroup.Group
 
-	if bp.reader != nil {
-		readerErr = bp.reader.Close()
+	if bp.conn != nil {
+		conn := bp.conn
+		g.Go(func() error {
+			return conn.Close()
+		})
+		bp.conn = nil
 	}
 
-	if bp.writer != nil {
-		writerErr = bp.writer.Close()
+	if bp.reader != nil {
+		reader := bp.reader
+		g.Go(func() error {
+			return reader.Close()
+		})
 	}
 
-	if bp.conn != nil {
-		connErr = bp.conn.Close()
-		bp.conn = nil
+	if bp.writer != nil {
+		writer := bp.writer
+		g.Go(func() error {
+			return writer.Close()
+		})
 	}
 
 	bp.connected = false
 
-	// Return first error encountered
-	if readerErr != nil {
-		return readerErr
-	}
-	if writerErr != nil {
-		return writerErr
-	}
-	return connErr
+	// Wait for all close operations to complete and return any error
+	return g.Wait()
 }
 
 // Connected returns whether the pipe is currently connected.
@@ -204,7 +199,7 @@ func (bp *BackedPipe) reconnectLocked() error {
 	// Get current writer sequence number to send to remote
 	writerSeqNum := bp.writer.SequenceNum()
 
-	conn, readerSeqNum, err := bp.reconnectFn(bp.ctx, writerSeqNum)
+	conn, readerSeqNum, err := bp.reconnector.Reconnect(bp.ctx, writerSeqNum)
 	if err != nil {
 		return ErrReconnectFailed
 	}
@@ -244,32 +239,40 @@ func (bp *BackedPipe) handleErrors() {
 		select {
 		case <-bp.ctx.Done():
 			return
-		case err := <-bp.errorChan:
-			// Connection error occurred
-			bp.mu.Lock()
-
-			// Skip if already closed or not connected
-			if bp.closed || !bp.connected {
-				bp.mu.Unlock()
-				continue
-			}
-
-			// Mark as disconnected
-			bp.connected = false
-
-			// Try to reconnect using internal context
-			reconnectErr := bp.reconnectLocked()
-			bp.mu.Unlock()
-
-			if reconnectErr != nil {
-				// Reconnection failed - log or handle as needed
-				// For now, we'll just continue and wait for manual reconnection
-				_ = err // Use the original error
-			}
+		case err := <-bp.readerErrorChan:
+			// Reader connection error occurred
+			bp.handleConnectionError(err, "reader")
+		case err := <-bp.writerErrorChan:
+			// Writer connection error occurred
+			bp.handleConnectionError(err, "writer")
 		}
 	}
 }
 
+// handleConnectionError handles errors from either reader or writer components.
+func (bp *BackedPipe) handleConnectionError(err error, component string) {
+	bp.mu.Lock()
+	defer bp.mu.Unlock()
+
+	// Skip if already closed or not connected
+	if bp.closed || !bp.connected {
+		return
+	}
+
+	// Mark as disconnected
+	bp.connected = false
+
+	// Try to reconnect using internal context
+	reconnectErr := bp.reconnectLocked()
+
+	if reconnectErr != nil {
+		// Reconnection failed - log or handle as needed
+		// For now, we'll just continue and wait for manual reconnection
+		_ = err       // Use the original error from the component
+		_ = component // Component info available for potential logging by higher layers
+	}
+}
+
 // ForceReconnect forces a reconnection attempt immediately.
 // This can be used to force a reconnection if a new connection is established.
 func (bp *BackedPipe) ForceReconnect() error {