Avoid async IPC panic on peer close

tsc/internal/ipc/conn_async.go

@@ -31,6 +31,7 @@ type AsyncConn struct {
 	pending   map[jsonrpc.ID]chan *Message
 	pendingMu sync.Mutex
 	terminal  error
+	hasCause  bool
 	writeMu   sync.Mutex
 	handlers  sync.WaitGroup
 }
@@ -66,10 +67,17 @@ func (c *AsyncConn) SetCollectTiming(enabled bool) {
 // It blocks until the context is cancelled or an error occurs.
 func (c *AsyncConn) Run(ctx context.Context) (err error) {
 	handlerCtx, cancelHandlers := context.WithCancel(ctx)
+	requestErrors := make(chan error, 1)
 	defer func() {
 		c.closePendingCalls(err)
 		cancelHandlers()
 		c.handlers.Wait()
+		select {
+		case requestErr := <-requestErrors:
+			err = errors.Join(err, requestErr)
+		default:
+			// No request failed before the read loop exited.
+		}
 	}()
 	for {
 		if ctx.Err() != nil {
@@ -88,7 +96,13 @@ func (c *AsyncConn) Run(ctx context.Context) (err error) {
 			c.handleResponse(msg)
 		} else if msg.IsRequest() {
 			c.handlers.Go(func() {
-				c.handleRequest(handlerCtx, msg)
+				if requestErr := c.handleRequest(handlerCtx, msg); requestErr != nil {
+					if c.recordRequestError(requestErr, requestErrors) {
+						if c.rwc != nil {
+							_ = c.rwc.Close()
+						}
+					}
+				}
 			})
 		} else if msg.IsNotification() {
 			c.handlers.Go(func() {
@@ -102,12 +116,38 @@ func (c *AsyncConn) Run(ctx context.Context) (err error) {
 func (c *AsyncConn) closePendingCalls(runErr error) {
 	c.pendingMu.Lock()
 	defer c.pendingMu.Unlock()
+	c.recordTerminalErrorLocked(runErr)
+	c.closePendingCallsLocked()
+}
+
+func (c *AsyncConn) recordRequestError(requestErr error, requestErrors chan<- error) bool {
+	c.pendingMu.Lock()
+	defer c.pendingMu.Unlock()
+	if !c.recordTerminalErrorLocked(requestErr) {
+		return false
+	}
+	requestErrors <- requestErr
+	c.closePendingCallsLocked()
+	return true
+}
+
+func (c *AsyncConn) recordTerminalErrorLocked(terminalErr error) bool {
 	if c.terminal == nil {
 		c.terminal = ErrConnClosed
-		if runErr != nil {
-			c.terminal = errors.Join(c.terminal, runErr)
+		if terminalErr != nil {
+			c.terminal = errors.Join(c.terminal, terminalErr)
+			c.hasCause = true
+			return true
 		}
+	} else if !c.hasCause && terminalErr != nil {
+		c.terminal = errors.Join(c.terminal, terminalErr)
+		c.hasCause = true
+		return true
 	}
+	return false
+}
+
+func (c *AsyncConn) closePendingCallsLocked() {
 	for id, ch := range c.pending {
 		close(ch)
 		delete(c.pending, id)
@@ -130,7 +170,7 @@ func (c *AsyncConn) handleResponse(msg *Message) {
 }
 
 // handleRequest processes an incoming request.
-func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) {
+func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) (retErr error) {
 	// Intercept the meta-requests for collected server timing before dispatching
 	// to the handler, so they are answered directly and not themselves recorded.
 	switch msg.Method {
@@ -139,9 +179,9 @@ func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) {
 		writeErr := c.protocol.WriteResponse(msg.ID, serverTimingSnapshot(c.timing))
 		c.writeMu.Unlock()
 		if writeErr != nil {
-			panic(fmt.Sprintf("ipc: failed to write server timing response: %v", writeErr))
+			return fmt.Errorf("ipc: failed to write server timing response: %w", writeErr)
 		}
-		return
+		return nil
 	case string(MethodResetServerTiming):
 		if c.timing != nil {
 			c.timing.reset()
@@ -150,9 +190,9 @@ func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) {
 		writeErr := c.protocol.WriteResponse(msg.ID, nil)
 		c.writeMu.Unlock()
 		if writeErr != nil {
-			panic(fmt.Sprintf("ipc: failed to write reset server timing response: %v", writeErr))
+			return fmt.Errorf("ipc: failed to write reset server timing response: %w", writeErr)
 		}
-		return
+		return nil
 	}
 
 	var result any
@@ -177,7 +217,7 @@ func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) {
 			c.writeMu.Unlock()
 
 			if writeErr != nil {
-				panic(fmt.Sprintf("ipc: failed to write panic error response: %v (original panic: %v)", writeErr, r))
+				retErr = fmt.Errorf("ipc: failed to write panic error response: %w (original panic: %v)", writeErr, r)
 			}
 		}
 	}()
@@ -202,8 +242,9 @@ func (c *AsyncConn) handleRequest(ctx context.Context, msg *Message) {
 	}
 
 	if writeErr != nil {
-		panic(fmt.Sprintf("ipc: failed to write response: %v", writeErr))
+		return fmt.Errorf("ipc: failed to write response: %w", writeErr)
 	}
+	return nil
 }
 
 // handleNotification processes an incoming notification.

tsc/internal/ipc/conn_async_test.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"io"
 	"net"
+	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -25,7 +27,8 @@ func (noOpHandler) HandleNotification(context.Context, string, json.Value) error
 }
 
 type queuedProtocol struct {
-	messages []*ipc.Message
+	messages    []*ipc.Message
+	responseErr error
 }
 
 func (p *queuedProtocol) ReadMessage() (*ipc.Message, error) {
@@ -46,11 +49,11 @@ func (p *queuedProtocol) WriteNotification(string, any) error {
 }
 
 func (p *queuedProtocol) WriteResponse(*jsonrpc.ID, any) error {
-	return nil
+	return p.responseErr
 }
 
 func (p *queuedProtocol) WriteError(*jsonrpc.ID, *jsonrpc.ResponseError) error {
-	return nil
+	return p.responseErr
 }
 
 type blockingHandler struct {
@@ -132,6 +135,70 @@ func TestAsyncConnRunCancelsHandlersOnEOF(t *testing.T) {
 	}
 }
 
+func TestAsyncConnResponseWriteFailureWithNilTransport(t *testing.T) {
+	t.Parallel()
+
+	responseErr := errors.New("response write failed")
+	id := jsonrpc.NewIDString("1")
+	protocol := &queuedProtocol{
+		messages:    []*ipc.Message{{ID: id, Method: "request"}},
+		responseErr: responseErr,
+	}
+	conn := ipc.NewAsyncConnWithProtocol(nil, protocol, noOpHandler{})
+
+	err := conn.Run(t.Context())
+	assert.Assert(t, errors.Is(err, responseErr), "expected response write error, got %v", err)
+}
+
+type closeSignal struct {
+	closed chan struct{}
+	once   sync.Once
+}
+
+func (*closeSignal) Read([]byte) (int, error) {
+	return 0, io.EOF
+}
+
+func (*closeSignal) Write(p []byte) (int, error) {
+	return len(p), nil
+}
+
+func (c *closeSignal) Close() error {
+	c.once.Do(func() { close(c.closed) })
+	return nil
+}
+
+type failingResponseProtocol struct {
+	closed      <-chan struct{}
+	requestRead bool
+	responseErr error
+}
+
+func (p *failingResponseProtocol) ReadMessage() (*ipc.Message, error) {
+	if !p.requestRead {
+		p.requestRead = true
+		return &ipc.Message{ID: jsonrpc.NewIDInt(1), Method: "transform"}, nil
+	}
+	<-p.closed
+	return nil, io.ErrClosedPipe
+}
+
+func (*failingResponseProtocol) WriteRequest(*jsonrpc.ID, string, any) error {
+	return nil
+}
+
+func (*failingResponseProtocol) WriteNotification(string, any) error {
+	return nil
+}
+
+func (p *failingResponseProtocol) WriteResponse(*jsonrpc.ID, any) error {
+	return p.responseErr
+}
+
+func (p *failingResponseProtocol) WriteError(*jsonrpc.ID, *jsonrpc.ResponseError) error {
+	return p.responseErr
+}
+
 func TestAsyncConnCallReturnsWhenPeerCloses(t *testing.T) {
 	t.Parallel()
 	client, server := net.Pipe()
@@ -176,3 +243,72 @@ func TestAsyncConnCallAfterReadLoopFailureReturnsImmediately(t *testing.T) {
 	err = conn.Notify(ctx, "changed", nil)
 	assert.Assert(t, errors.Is(err, ipc.ErrConnClosed), "expected ErrConnClosed, got %v", err)
 }
+
+func TestAsyncConnTerminalErrorIncludesResponseWriteFailure(t *testing.T) {
+	t.Parallel()
+	responseErr := errors.New("response write failed")
+	rwc := &closeSignal{closed: make(chan struct{})}
+	protocol := &failingResponseProtocol{
+		closed:      rwc.closed,
+		responseErr: responseErr,
+	}
+	conn := ipc.NewAsyncConnWithProtocol(rwc, protocol, noOpHandler{})
+
+	err := conn.Run(t.Context())
+	assert.Assert(t, errors.Is(err, responseErr), "expected response write error, got %v", err)
+	_, err = conn.Call(t.Context(), "transform", nil)
+	assert.Assert(t, errors.Is(err, responseErr), "expected terminal response write error, got %v", err)
+	assert.Equal(t, strings.Count(err.Error(), responseErr.Error()), 1)
+}
+
+func TestAsyncConnRunWaitsForRequestAfterPeerCloses(t *testing.T) {
+	t.Parallel()
+	client, server := net.Pipe()
+	defer server.Close()
+	handler := &blockingHandler{
+		started: make(chan struct{}, 1),
+		release: make(chan struct{}),
+	}
+	defer func() {
+		select {
+		case <-handler.release:
+			return
+		default:
+			close(handler.release)
+		}
+	}()
+	conn := ipc.NewAsyncConn(server, handler)
+	runDone := make(chan error, 1)
+	go func() { runDone <- conn.Run(t.Context()) }()
+
+	clientProtocol := ipc.NewJSONRPCProtocol(client)
+	assert.NilError(t, clientProtocol.WriteRequest(jsonrpc.NewIDInt(1), "transform", nil))
+	select {
+	case <-handler.started:
+		break
+	case <-time.After(time.Second):
+		t.Fatal("request handler did not start")
+	}
+	assert.NilError(t, client.Close())
+
+	handlerBlocked := false
+	select {
+	case err := <-runDone:
+		t.Fatalf("connection stopped while request handler was blocked: %v", err)
+	case <-time.After(100 * time.Millisecond):
+		handlerBlocked = true
+	}
+	assert.Assert(t, handlerBlocked)
+
+	close(handler.release)
+	select {
+	case err := <-runDone:
+		assert.ErrorContains(t, err, "ipc: failed to write response")
+		_, err = conn.Call(t.Context(), "transform", nil)
+		assert.ErrorContains(t, err, "ipc: failed to write response")
+		err = conn.Notify(t.Context(), "changed", nil)
+		assert.ErrorContains(t, err, "ipc: failed to write response")
+	case <-time.After(time.Second):
+		t.Fatal("connection did not stop after request handler completed")
+	}
+}

tsc/internal/ipc/conn_sync.go

@@ -70,7 +70,9 @@ func (c *SyncConn) Run(ctx context.Context) error {
 		}
 
 		if msg.IsRequest() {
-			c.handleRequest(ctx, msg)
+			if err := c.handleRequest(ctx, msg); err != nil {
+				return err
+			}
 		} else if msg.IsNotification() {
 			c.handleNotification(ctx, msg)
 		} else {
@@ -81,7 +83,7 @@ func (c *SyncConn) Run(ctx context.Context) error {
 }
 
 // handleRequest processes an incoming request.
-func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) {
+func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) (retErr error) {
 	// Intercept the meta-requests for collected server timing before dispatching
 	// to the handler, so they are answered directly and not themselves recorded.
 	switch msg.Method {
@@ -90,9 +92,9 @@ func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) {
 		writeErr := c.protocol.WriteResponse(msg.ID, serverTimingSnapshot(c.timing))
 		c.mu.Unlock()
 		if writeErr != nil {
-			panic(fmt.Sprintf("ipc: failed to write server timing response: %v", writeErr))
+			return fmt.Errorf("ipc: failed to write server timing response: %w", writeErr)
 		}
-		return
+		return nil
 	case string(MethodResetServerTiming):
 		if c.timing != nil {
 			c.timing.reset()
@@ -101,9 +103,9 @@ func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) {
 		writeErr := c.protocol.WriteResponse(msg.ID, nil)
 		c.mu.Unlock()
 		if writeErr != nil {
-			panic(fmt.Sprintf("ipc: failed to write reset server timing response: %v", writeErr))
+			return fmt.Errorf("ipc: failed to write reset server timing response: %w", writeErr)
 		}
-		return
+		return nil
 	}
 
 	var result any
@@ -128,7 +130,7 @@ func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) {
 			c.mu.Unlock()
 
 			if writeErr != nil {
-				panic(fmt.Sprintf("ipc: failed to write panic error response: %v (original panic: %v)", writeErr, r))
+				retErr = fmt.Errorf("ipc: failed to write panic error response: %w (original panic: %v)", writeErr, r)
 			}
 		}
 	}()
@@ -153,8 +155,9 @@ func (c *SyncConn) handleRequest(ctx context.Context, msg *Message) {
 	}
 
 	if writeErr != nil {
-		panic(fmt.Sprintf("ipc: failed to write response: %v", writeErr))
+		return fmt.Errorf("ipc: failed to write response: %w", writeErr)
 	}
+	return nil
 }
 
 // handleNotification processes an incoming notification.