fix(pipeline): decrement inflight metrics before write-failure guard (llm-d#591)

After a persistence failure in Drain(), the firstErr guard skipped
  Receive() for remaining results. Since Receive() was the only site
  that decremented inflight gauges for async-path results, the counters
  leaked on write failures.

  Move the metrics decrement block from Receive() into Drain(), before
  the firstErr check, so counters are always cleaned up regardless of
  write success.

Signed-off-by: Madhu Goutham Reddy Ambati <mambati@redhat.com>
Co-authored-by: Lior Aronovich <243445518+lioraron@users.noreply.github.com>

Author: madhugoutham

internal/processor/pipeline/collector.go

@@ -56,6 +56,11 @@ func (c *ResultCollector) Drain(ctx context.Context, resultCh <-chan ResultItem)
 		if !c.pending.Resolve(&msg) {
 			continue
 		}
+		if !msg.SubmittedAt.IsZero() {
+			metrics.DecProcessorInflightRequests()
+			metrics.DecModelInflightRequests(msg.ModelID)
+			metrics.RecordModelRequestExecutionDuration(time.Since(msg.SubmittedAt), msg.ModelID)
+		}
 		if firstErr != nil {
 			continue
 		}
@@ -95,12 +100,6 @@ func (c *ResultCollector) Receive(msg ResultItem) error {
 		return fmt.Errorf("write output for %s: %w", msg.RequestID, err)
 	}
 
-	if !msg.SubmittedAt.IsZero() {
-		metrics.DecProcessorInflightRequests()
-		metrics.DecModelInflightRequests(msg.ModelID)
-		metrics.RecordModelRequestExecutionDuration(time.Since(msg.SubmittedAt), msg.ModelID)
-	}
-
 	if line.isSuccess() {
 		c.tracker.RecordSuccess(msg)
 		if msg.Response != nil {

internal/processor/pipeline/collector_test.go

@@ -1,9 +1,12 @@
 package pipeline
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
+	"fmt"
 	"testing"
+	"time"
 
 	"github.com/go-logr/logr"
 
@@ -138,6 +141,57 @@ func TestResultCollector_DrainProcessesAllResultsAfterCancel(t *testing.T) {
 	}
 }
 
+type failAfterNWriter struct {
+	remaining int
+}
+
+func (w *failAfterNWriter) Write(p []byte) (int, error) {
+	if w.remaining <= 0 {
+		return 0, fmt.Errorf("simulated write failure")
+	}
+	w.remaining--
+	return len(p), nil
+}
+
+func TestResultCollector_DrainDecrementsMetricsAfterWriteFailure(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)
+
+	now := time.Now()
+	results := []ResultItem{
+		{RequestID: "req-1", CustomID: "c-1", SubmittedAt: now, Response: &batch_types.ResponseData{StatusCode: 200, RequestID: "req-1", Body: map[string]any{"ok": true}}},
+		{RequestID: "req-2", CustomID: "c-2", SubmittedAt: now, Response: &batch_types.ResponseData{StatusCode: 200, RequestID: "req-2", Body: map[string]any{"ok": true}}},
+		{RequestID: "req-3", CustomID: "c-3", SubmittedAt: now, 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, SubmittedAt: r.SubmittedAt})
+	}
+
+	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")
+	}
+
+	var remaining int
+	pending.DrainUnresolved(func(_ RequestItem) {
+		remaining++
+	})
+	if remaining != 0 {
+		t.Fatalf("pending requests remaining = %d, want 0 (all should be resolved despite write failure)", remaining)
+	}
+}
+
 func TestProgressTracker_AddFailed(t *testing.T) {
 	tracker := NewProgressTracker(10, nil, "test-job", 0, logr.Discard())
 	tracker.AddFailed(5)