feat(pipeline): cancel dispatch on collector persistence failure

internal/processor/pipeline/collector.go

@@ -31,11 +31,12 @@ func (o *outputLine) isSuccess() bool {
 // ResultCollector writes ResultItem values to JSONL and records progress.
 // Terminal actor — no out channel.
 type ResultCollector struct {
-	output  *bufio.Writer
-	errors  *bufio.Writer
-	pending *PendingRequests
-	tracker *ProgressTracker
-	logger  logr.Logger
+	output               *bufio.Writer
+	errors               *bufio.Writer
+	pending              *PendingRequests
+	tracker              *ProgressTracker
+	logger               logr.Logger
+	onPersistenceFailure func()
 }
 
 func NewResultCollector(outputFile, errorFile *os.File, pending *PendingRequests, tracker *ProgressTracker, logger logr.Logger) *ResultCollector {
@@ -69,6 +70,9 @@ func (c *ResultCollector) Drain(ctx context.Context, resultCh <-chan ResultItem)
 		if err := c.Receive(msg); err != nil {
 			firstErr = err
 			c.logger.Error(err, "Persistence failure, skipping further writes")
+			if c.onPersistenceFailure != nil {
+				c.onPersistenceFailure()
+			}
 		}
 	}
 	if flushErr := c.flushFiles(); flushErr != nil {

internal/processor/pipeline/collector_test.go

@@ -194,6 +194,42 @@ func TestResultCollector_DrainDecrementsMetricsAfterWriteFailure(t *testing.T) {
 	}
 }
 
+func TestResultCollector_DrainCallsOnPersistenceFailure(t *testing.T) {
+	outputFile := tempFile(t)
+	errorFile := tempFile(t)
+	pending := NewPendingRequests(0)
+	tracker := NewProgressTracker(3, nil, "test-job", 0, logr.Discard())
+	collector := NewResultCollector(outputFile, errorFile, pending, tracker, logr.Discard())
+
+	collector.output = bufio.NewWriterSize(&failAfterNWriter{remaining: 1}, 1)
+
+	var callCount int
+	collector.onPersistenceFailure = func() { callCount++ }
+
+	results := []ResultItem{
+		{RequestID: "req-1", CustomID: "c-1", Response: &batch_types.ResponseData{StatusCode: 200, RequestID: "req-1", Body: map[string]any{"ok": true}}},
+		{RequestID: "req-2", CustomID: "c-2", Response: &batch_types.ResponseData{StatusCode: 200, RequestID: "req-2", Body: map[string]any{"ok": true}}},
+		{RequestID: "req-3", CustomID: "c-3", Response: &batch_types.ResponseData{StatusCode: 200, RequestID: "req-3", Body: map[string]any{"ok": true}}},
+	}
+	for _, r := range results {
+		pending.Store(RequestItem{RequestID: r.RequestID, CustomID: r.CustomID})
+	}
+
+	ch := make(chan ResultItem, len(results))
+	for _, r := range results {
+		ch <- r
+	}
+	close(ch)
+
+	err := collector.Drain(context.Background(), ch)
+	if err == nil {
+		t.Fatal("expected write failure error from Drain")
+	}
+	if callCount != 1 {
+		t.Fatalf("onPersistenceFailure called %d times, want 1", callCount)
+	}
+}
+
 func TestResultCollector_DrainReturnsNilWhenCtxCancelled(t *testing.T) {
 	outputFile := tempFile(t)
 	errorFile := tempFile(t)

internal/processor/pipeline/executor.go

@@ -2,6 +2,8 @@ package pipeline
 
 import (
 	"context"
+	"errors"
+	"sync"
 
 	"github.com/go-logr/logr"
 	"golang.org/x/sync/errgroup"
@@ -39,6 +41,11 @@ func NewJobExecutor(cfg JobExecutorConfig) *JobExecutor {
 func (je *JobExecutor) Execute(ctx context.Context) (*openai.BatchRequestCounts, error) {
 	g, ctx := errgroup.WithContext(ctx)
 
+	dispatchCtx, dispatchCancel := context.WithCancel(ctx)
+	defer dispatchCancel()
+
+	je.collector.onPersistenceFailure = sync.OnceFunc(dispatchCancel)
+
 	requestCh := make(chan RequestItem)
 	resultCh := make(chan ResultItem, defaultResultBuffer)
 
@@ -52,14 +59,26 @@ func (je *JobExecutor) Execute(ctx context.Context) (*openai.BatchRequestCounts,
 		close(trackerDone)
 	}()
 
-	g.Go(func() error { return je.dispatcher.Run(ctx, requestCh, resultCh) })
+	g.Go(func() error { return je.dispatcher.Run(dispatchCtx, requestCh, resultCh) })
 
-	g.Go(func() error { return je.collector.Drain(ctx, resultCh) })
+	var collectorErr error
+	g.Go(func() error {
+		collectorErr = je.collector.Drain(ctx, resultCh)
+		return collectorErr
+	})
 
-	g.Go(func() error { return je.source.Produce(ctx, requestCh) })
+	g.Go(func() error { return je.source.Produce(dispatchCtx, requestCh) })
 
 	err := g.Wait()
 
+	// When onPersistenceFailure cancels dispatchCtx, the source returns
+	// context.Canceled before the collector finishes draining resultCh.
+	// errgroup records the first non-nil error (source's cancel), masking
+	// the root cause. Prefer the collector's persistence error.
+	if errors.Is(err, context.Canceled) && collectorErr != nil {
+		err = collectorErr
+	}
+
 	trackerCancel()
 	<-trackerDone
 

internal/processor/pipeline/executor_test.go

@@ -1,9 +1,11 @@
 package pipeline
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -723,6 +725,130 @@ func getCounterValue(t *testing.T, name, model string) float64 {
 	return 0
 }
 
+// TestJobExecutor_PersistenceFailureCancelsDispatch verifies that when the
+// collector encounters a write failure, dispatch is cancelled early: the slow
+// source should NOT deliver all requests because dispatchCtx is cancelled
+// after the first persistence error.
+func TestJobExecutor_PersistenceFailureCancelsDispatch(t *testing.T) {
+	const total = 20
+	var dispatched atomic.Int32
+	client := &mockInferenceClientForE2E{
+		generateFn: func(_ context.Context, req *inference.GenerateRequest) (*inference.GenerateResponse, *inference.ClientError) {
+			dispatched.Add(1)
+			body, _ := json.Marshal(map[string]any{"ok": true})
+			return &inference.GenerateResponse{RequestID: req.RequestID, Response: body}, nil
+		},
+	}
+	resolver := inference.NewSingleClientResolver(client)
+	defer func() { _ = resolver.Close() }()
+
+	items := makeItems(total, "m1")
+
+	outputFile := tempFile(t)
+	errorFile := tempFile(t)
+	pending := NewPendingRequests(0)
+	tracker := NewProgressTracker(int64(total), nil, "test-job", 0, logr.Discard())
+	collector := NewResultCollector(outputFile, errorFile, pending, tracker, logr.Discard())
+
+	collector.output = bufio.NewWriterSize(&failAfterNWriter{remaining: 1}, 1)
+
+	source := &throttledSource{items: items, delay: 10 * time.Millisecond}
+
+	executor := NewJobExecutor(JobExecutorConfig{
+		Source:     source,
+		Dispatcher: NewPreDispatcher(NewDirectDispatcher(resolver, logr.Discard())),
+		Collector:  collector,
+		Tracker:    tracker,
+		Logger:     logr.Discard(),
+	})
+
+	_, err := executor.Execute(context.Background())
+	if err == nil {
+		t.Fatal("expected persistence error from Execute")
+	}
+
+	d := dispatched.Load()
+	if d >= int32(total) {
+		t.Fatalf("dispatched %d requests (all %d), want fewer: persistence failure should cancel dispatch early", d, total)
+	}
+	t.Logf("dispatched %d/%d requests before cancellation", d, total)
+}
+
+// throttledSource emits items with a configurable delay between sends,
+// respecting context cancellation.
+type throttledSource struct {
+	items []RequestItem
+	delay time.Duration
+}
+
+func (s *throttledSource) Produce(ctx context.Context, out chan<- RequestItem) error {
+	defer close(out)
+	for _, item := range s.items {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case out <- item:
+		}
+		if s.delay > 0 {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(s.delay):
+			}
+		}
+	}
+	return nil
+}
+
+// TestJobExecutor_PersistenceFailureReturnedNotCanceled verifies that Execute()
+// returns the collector's write error, not context.Canceled from the source.
+// When onPersistenceFailure cancels dispatchCtx, the source exits early with
+// context.Canceled. Without the fix, errgroup's sync.Once records that as the
+// first error, masking the persistence failure.
+func TestJobExecutor_PersistenceFailureReturnedNotCanceled(t *testing.T) {
+	const total = 10
+	client := &mockInferenceClientForE2E{
+		generateFn: func(_ context.Context, req *inference.GenerateRequest) (*inference.GenerateResponse, *inference.ClientError) {
+			body, _ := json.Marshal(map[string]any{"ok": true})
+			return &inference.GenerateResponse{RequestID: req.RequestID, Response: body}, nil
+		},
+	}
+	resolver := inference.NewSingleClientResolver(client)
+	defer func() { _ = resolver.Close() }()
+
+	items := makeItems(total, "m1")
+
+	outputFile := tempFile(t)
+	errorFile := tempFile(t)
+	pending := NewPendingRequests(0)
+	tracker := NewProgressTracker(int64(total), nil, "test-job", 0, logr.Discard())
+	collector := NewResultCollector(outputFile, errorFile, pending, tracker, logr.Discard())
+
+	collector.output = bufio.NewWriterSize(&failAfterNWriter{remaining: 1}, 1)
+
+	source := &throttledSource{items: items, delay: 10 * time.Millisecond}
+
+	executor := NewJobExecutor(JobExecutorConfig{
+		Source:     source,
+		Dispatcher: NewPreDispatcher(NewDirectDispatcher(resolver, logr.Discard())),
+		Collector:  collector,
+		Tracker:    tracker,
+		Logger:     logr.Discard(),
+	})
+
+	_, err := executor.Execute(context.Background())
+	if err == nil {
+		t.Fatal("expected error from Execute")
+	}
+	if err == context.Canceled {
+		t.Fatalf("Execute() returned context.Canceled; want the persistence write error")
+	}
+	if !strings.Contains(err.Error(), "write") {
+		t.Errorf("Execute() error = %q; want an error containing 'write' (the persistence failure)", err)
+	}
+	t.Logf("Execute() correctly returned persistence error: %v", err)
+}
+
 var _ inference.InferenceClient = (*mockInferenceClient)(nil)
 var _ inference.InferenceClient = (*mockInferenceClientForE2E)(nil)
 var _ RequestSource = (*sliceSource)(nil)