Merge pull request #4481 from KorenKrita/agent/fix-plugin-stream-bridge-close-race

fix(pluginhost): prevent stream close/send panic

internal/pluginhost/stream_bridge.go

@@ -2,6 +2,7 @@ package pluginhost
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 	"sync"
@@ -13,7 +14,33 @@ import (
 type streamBridge struct {
 	next    atomic.Uint64
 	mu      sync.Mutex
-	streams map[string]chan pluginapi.ExecutorStreamChunk
+	streams map[string]*streamBridgeStream
+}
+
+const streamBridgeBufferSize = 16
+
+var errStreamBridgeClosed = errors.New("stream is not open")
+
+type streamBridgeStream struct {
+	chunks    chan pluginapi.ExecutorStreamChunk
+	emits     chan streamBridgeEmit
+	closes    chan streamBridgeClose
+	closed    chan struct{}
+	finished  chan struct{}
+	abort     chan struct{}
+	closeOnce sync.Once
+	abortOnce sync.Once
+}
+
+type streamBridgeEmit struct {
+	ctx   context.Context
+	chunk pluginapi.ExecutorStreamChunk
+	done  chan error
+}
+
+type streamBridgeClose struct {
+	errorMessage string
+	accepted     chan struct{}
 }
 
 type rpcStreamEmitRequest struct {
@@ -28,7 +55,129 @@ type rpcStreamCloseRequest struct {
 }
 
 func newStreamBridge() *streamBridge {
-	return &streamBridge{streams: make(map[string]chan pluginapi.ExecutorStreamChunk)}
+	return &streamBridge{streams: make(map[string]*streamBridgeStream)}
+}
+
+func newStreamBridgeStream() *streamBridgeStream {
+	stream := &streamBridgeStream{
+		chunks:   make(chan pluginapi.ExecutorStreamChunk),
+		emits:    make(chan streamBridgeEmit),
+		closes:   make(chan streamBridgeClose),
+		closed:   make(chan struct{}),
+		finished: make(chan struct{}),
+		abort:    make(chan struct{}),
+	}
+	go stream.run()
+	return stream
+}
+
+func (s *streamBridgeStream) run() {
+	defer func() {
+		s.markClosed()
+		close(s.chunks)
+		close(s.finished)
+	}()
+
+	queue := make([]pluginapi.ExecutorStreamChunk, 0, streamBridgeBufferSize)
+	for {
+		var emitC <-chan streamBridgeEmit
+		if len(queue) < streamBridgeBufferSize {
+			emitC = s.emits
+		}
+		var outputC chan pluginapi.ExecutorStreamChunk
+		var next pluginapi.ExecutorStreamChunk
+		if len(queue) > 0 {
+			outputC = s.chunks
+			next = queue[0]
+		}
+
+		select {
+		case <-s.abort:
+			return
+		case request := <-s.closes:
+			s.markClosed()
+			close(request.accepted)
+			if request.errorMessage != "" {
+				queue = append(queue, pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", request.errorMessage)})
+			}
+			for len(queue) > 0 {
+				select {
+				case <-s.abort:
+					return
+				case s.chunks <- queue[0]:
+					queue = queue[1:]
+				}
+			}
+			return
+		case request := <-emitC:
+			if err := request.ctx.Err(); err != nil {
+				request.done <- err
+				continue
+			}
+			queue = append(queue, request.chunk)
+			request.done <- nil
+		case outputC <- next:
+			queue = queue[1:]
+		}
+	}
+}
+
+func (s *streamBridgeStream) markClosed() {
+	if s == nil {
+		return
+	}
+	s.closeOnce.Do(func() { close(s.closed) })
+}
+
+func (s *streamBridgeStream) abortStream() {
+	if s == nil {
+		return
+	}
+	s.abortOnce.Do(func() {
+		s.markClosed()
+		close(s.abort)
+	})
+}
+
+func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorStreamChunk) error {
+	if s == nil {
+		return errStreamBridgeClosed
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	request := streamBridgeEmit{
+		ctx:   ctx,
+		chunk: chunk,
+		done:  make(chan error, 1),
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-s.closed:
+		return errStreamBridgeClosed
+	case s.emits <- request:
+	}
+	return <-request.done
+}
+
+func (s *streamBridgeStream) close(errorMessage string) {
+	if s == nil {
+		return
+	}
+	request := streamBridgeClose{
+		errorMessage: errorMessage,
+		accepted:     make(chan struct{}),
+	}
+	select {
+	case <-s.finished:
+		return
+	case s.closes <- request:
+	}
+	select {
+	case <-request.accepted:
+	case <-s.finished:
+	}
 }
 
 func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.ExecutorStreamChunk, func()) {
@@ -38,20 +187,26 @@ func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.Execu
 		return "", chunks, func() {}
 	}
 	id := strconv.FormatUint(b.next.Add(1), 10)
-	chunks := make(chan pluginapi.ExecutorStreamChunk, 16)
+	stream := newStreamBridgeStream()
 	b.mu.Lock()
-	b.streams[id] = chunks
+	b.streams[id] = stream
 	b.mu.Unlock()
 	cleanup := func() {
-		b.close(id, "")
+		b.mu.Lock()
+		if b.streams[id] == stream {
+			delete(b.streams, id)
+		}
+		b.mu.Unlock()
+		stream.abortStream()
 	}
 	if ctx != nil && ctx.Done() != nil {
+		// Abort streams canceled before ExecuteStream can install cleanupWhenStreamDone.
 		go func() {
 			<-ctx.Done()
-			b.close(id, ctx.Err().Error())
+			cleanup()
 		}()
 	}
-	return id, chunks, cleanup
+	return id, stream.chunks, cleanup
 }
 
 func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.ExecutorStreamChunk) error {
@@ -59,20 +214,18 @@ func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.Exec
 		return fmt.Errorf("stream id is required")
 	}
 	b.mu.Lock()
-	chunks := b.streams[id]
+	stream := b.streams[id]
 	b.mu.Unlock()
-	if chunks == nil {
+	if stream == nil {
 		return fmt.Errorf("stream %s is not open", id)
 	}
-	if ctx == nil {
-		ctx = context.Background()
-	}
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case chunks <- chunk:
-		return nil
+	if err := stream.emit(ctx, chunk); err != nil {
+		if errors.Is(err, errStreamBridgeClosed) {
+			return fmt.Errorf("stream %s is not open", id)
+		}
+		return err
 	}
+	return nil
 }
 
 func (b *streamBridge) close(id string, errorMessage string) {
@@ -80,14 +233,11 @@ func (b *streamBridge) close(id string, errorMessage string) {
 		return
 	}
 	b.mu.Lock()
-	chunks := b.streams[id]
+	stream := b.streams[id]
 	delete(b.streams, id)
 	b.mu.Unlock()
-	if chunks == nil {
+	if stream == nil {
 		return
 	}
-	if errorMessage != "" {
-		chunks <- pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", errorMessage)}
-	}
-	close(chunks)
+	stream.close(errorMessage)
 }

internal/pluginhost/stream_bridge_test.go

@@ -0,0 +1,197 @@
+package pluginhost
+
+import (
+	"context"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
+)
+
+type streamBridgeNotifyContext struct {
+	context.Context
+	ready chan struct{}
+	once  sync.Once
+}
+
+func (c *streamBridgeNotifyContext) Done() <-chan struct{} {
+	c.once.Do(func() { close(c.ready) })
+	return c.Context.Done()
+}
+
+func TestStreamBridgeCloseUnblocksPendingEmit(t *testing.T) {
+	bridge := newStreamBridge()
+	streamID, chunks, _ := bridge.open(context.Background())
+
+	for range streamBridgeBufferSize {
+		if err := bridge.emit(context.Background(), streamID, pluginapi.ExecutorStreamChunk{Payload: []byte("buffered")}); err != nil {
+			t.Fatalf("fill stream buffer: %v", err)
+		}
+	}
+
+	emitCtx := &streamBridgeNotifyContext{
+		Context: context.Background(),
+		ready:   make(chan struct{}),
+	}
+	emitDone := make(chan error, 1)
+	go func() {
+		emitDone <- bridge.emit(emitCtx, streamID, pluginapi.ExecutorStreamChunk{Payload: []byte("blocked")})
+	}()
+
+	select {
+	case <-emitCtx.ready:
+	case <-time.After(time.Second):
+		t.Fatal("emit did not reach the blocked send")
+	}
+	select {
+	case err := <-emitDone:
+		t.Fatalf("emit returned while the stream buffer was full: %v", err)
+	default:
+	}
+
+	bridge.close(streamID, "")
+
+	select {
+	case err := <-emitDone:
+		if err == nil || !strings.Contains(err.Error(), "is not open") {
+			t.Fatalf("emit error = %v, want stream-not-open error", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("close did not unblock the pending emit")
+	}
+
+	chunkCount := 0
+	for range chunks {
+		chunkCount++
+	}
+	if chunkCount != streamBridgeBufferSize {
+		t.Fatalf("delivered chunks = %d, want %d buffered chunks without the rejected emit", chunkCount, streamBridgeBufferSize)
+	}
+}
+
+func TestStreamBridgeEmitUsesAcceptedPumpResultAfterContextCancellation(t *testing.T) {
+	for range 1000 {
+		ctx, cancel := context.WithCancel(context.Background())
+		stream := &streamBridgeStream{
+			emits:  make(chan streamBridgeEmit),
+			closed: make(chan struct{}),
+		}
+		go func() {
+			request := <-stream.emits
+			cancel()
+			request.done <- nil
+		}()
+
+		if err := stream.emit(ctx, pluginapi.ExecutorStreamChunk{Payload: []byte("accepted")}); err != nil {
+			t.Fatalf("accepted emit returned error: %v", err)
+		}
+	}
+}
+
+func TestStreamBridgeAbortClosesSaturatedStreamWithoutConsumer(t *testing.T) {
+	bridge := newStreamBridge()
+	streamID, chunks, cleanup := bridge.open(context.Background())
+	bridge.mu.Lock()
+	stream := bridge.streams[streamID]
+	bridge.mu.Unlock()
+
+	for range streamBridgeBufferSize {
+		if err := bridge.emit(context.Background(), streamID, pluginapi.ExecutorStreamChunk{Payload: []byte("buffered")}); err != nil {
+			t.Fatalf("fill stream buffer: %v", err)
+		}
+	}
+
+	cleanup()
+
+	select {
+	case <-stream.finished:
+	case <-time.After(time.Second):
+		t.Fatal("abort left the saturated stream pump running")
+	}
+	if _, ok := <-chunks; ok {
+		t.Fatal("aborted stream retained buffered chunks")
+	}
+}
+
+func TestStreamBridgeCleanupAbortsPendingGracefulClose(t *testing.T) {
+	bridge := newStreamBridge()
+	streamID, chunks, cleanup := bridge.open(context.Background())
+	bridge.mu.Lock()
+	stream := bridge.streams[streamID]
+	bridge.mu.Unlock()
+
+	for range streamBridgeBufferSize {
+		if err := bridge.emit(context.Background(), streamID, pluginapi.ExecutorStreamChunk{Payload: []byte("buffered")}); err != nil {
+			t.Fatalf("fill stream buffer: %v", err)
+		}
+	}
+	bridge.close(streamID, "plugin stream failed")
+
+	cleanup()
+
+	select {
+	case <-stream.finished:
+	case <-time.After(time.Second):
+		t.Fatal("cleanup did not abort the graceful close after the stream was removed")
+	}
+	if _, ok := <-chunks; ok {
+		t.Fatal("cleanup retained queued chunks after aborting the graceful close")
+	}
+}
+
+func TestStreamBridgeCloseDeliversTerminalError(t *testing.T) {
+	bridge := newStreamBridge()
+	streamID, chunks, _ := bridge.open(context.Background())
+
+	bridge.close(streamID, "plugin stream failed")
+
+	chunk, ok := <-chunks
+	if !ok {
+		t.Fatal("stream closed before terminal error")
+	}
+	if chunk.Err == nil || chunk.Err.Error() != "plugin stream failed" {
+		t.Fatalf("terminal error = %v, want plugin stream failed", chunk.Err)
+	}
+	if _, ok = <-chunks; ok {
+		t.Fatal("stream remains open after terminal error")
+	}
+}
+
+func TestStreamBridgeClosePreservesTerminalErrorWhenBufferIsFull(t *testing.T) {
+	bridge := newStreamBridge()
+	streamID, chunks, _ := bridge.open(context.Background())
+
+	for range streamBridgeBufferSize {
+		if err := bridge.emit(context.Background(), streamID, pluginapi.ExecutorStreamChunk{Payload: []byte("buffered")}); err != nil {
+			t.Fatalf("fill stream buffer: %v", err)
+		}
+	}
+
+	closeDone := make(chan struct{})
+	go func() {
+		bridge.close(streamID, "plugin stream failed")
+		close(closeDone)
+	}()
+	select {
+	case <-closeDone:
+	case <-time.After(time.Second):
+		t.Fatal("close blocked on the saturated stream")
+	}
+
+	chunkCount := 0
+	var terminalErr error
+	for chunk := range chunks {
+		chunkCount++
+		if chunk.Err != nil {
+			terminalErr = chunk.Err
+		}
+	}
+	if chunkCount != streamBridgeBufferSize+1 {
+		t.Fatalf("delivered chunks = %d, want %d buffered chunks plus terminal error", chunkCount, streamBridgeBufferSize+1)
+	}
+	if terminalErr == nil || terminalErr.Error() != "plugin stream failed" {
+		t.Fatalf("terminal error = %v, want plugin stream failed", terminalErr)
+	}
+}