diff --git a/internal/pluginhost/stream_bridge.go b/internal/pluginhost/stream_bridge.go index 632cc2bc2..9002e2829 100644 --- a/internal/pluginhost/stream_bridge.go +++ b/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) } diff --git a/internal/pluginhost/stream_bridge_test.go b/internal/pluginhost/stream_bridge_test.go new file mode 100644 index 000000000..8cdb1a6aa --- /dev/null +++ b/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) + } +}