TRT-2811: Make prow loader batch writes resilient to single-batch failures

pkg/dataloader/prowloader/accumulate_test.go

@@ -0,0 +1,221 @@
+package prowloader
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func makeResults(n int) []*jobRunResult {
+	results := make([]*jobRunResult, n)
+	for i := range results {
+		results[i] = &jobRunResult{
+			Run: prowJobRunRow{ID: uint(i + 1)},
+		}
+	}
+	return results
+}
+
+func sendResults(results []*jobRunResult) <-chan *jobRunResult {
+	ch := make(chan *jobRunResult, len(results))
+	for _, r := range results {
+		ch <- r
+	}
+	close(ch)
+	return ch
+}
+
+func TestAccumulateAndWriteJobRuns(t *testing.T) {
+	tests := []struct {
+		name           string
+		resultCount    int
+		writerFunc     func(calls *[][]jobRunResult) batchWriterFunc
+		cancelAfter    int
+		wantErrorCount int
+		wantErrorMsg   string
+	}{
+		{
+			name:        "all batches succeed",
+			resultCount: 250,
+			writerFunc: func(calls *[][]jobRunResult) batchWriterFunc {
+				return func(ctx context.Context, batch []jobRunResult) error {
+					cp := make([]jobRunResult, len(batch))
+					copy(cp, batch)
+					*calls = append(*calls, cp)
+					return nil
+				}
+			},
+			wantErrorCount: 0,
+		},
+		{
+			name:        "first batch fails second succeeds",
+			resultCount: 200,
+			writerFunc: func(calls *[][]jobRunResult) batchWriterFunc {
+				callCount := 0
+				return func(ctx context.Context, batch []jobRunResult) error {
+					cp := make([]jobRunResult, len(batch))
+					copy(cp, batch)
+					*calls = append(*calls, cp)
+					callCount++
+					if callCount == 1 {
+						return fmt.Errorf("connection refused")
+					}
+					return nil
+				}
+			},
+			wantErrorCount: 1,
+			wantErrorMsg:   "connection refused",
+		},
+		{
+			name:        "all batches fail",
+			resultCount: 300,
+			writerFunc: func(calls *[][]jobRunResult) batchWriterFunc {
+				return func(ctx context.Context, batch []jobRunResult) error {
+					cp := make([]jobRunResult, len(batch))
+					copy(cp, batch)
+					*calls = append(*calls, cp)
+					return fmt.Errorf("database unavailable")
+				}
+			},
+			wantErrorCount: 3,
+			wantErrorMsg:   "database unavailable",
+		},
+		{
+			name:        "trailing batch fails",
+			resultCount: 150,
+			writerFunc: func(calls *[][]jobRunResult) batchWriterFunc {
+				callCount := 0
+				return func(ctx context.Context, batch []jobRunResult) error {
+					cp := make([]jobRunResult, len(batch))
+					copy(cp, batch)
+					*calls = append(*calls, cp)
+					callCount++
+					if callCount == 2 {
+						return fmt.Errorf("partition not found")
+					}
+					return nil
+				}
+			},
+			wantErrorCount: 1,
+			wantErrorMsg:   "partition not found",
+		},
+		{
+			name:        "context cancelled with pending batch",
+			resultCount: 50,
+			cancelAfter: 25,
+			writerFunc: func(calls *[][]jobRunResult) batchWriterFunc {
+				return func(ctx context.Context, batch []jobRunResult) error {
+					cp := make([]jobRunResult, len(batch))
+					copy(cp, batch)
+					*calls = append(*calls, cp)
+					return nil
+				}
+			},
+			wantErrorCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var calls [][]jobRunResult
+			pl := &ProwLoader{
+				batchWriter: tt.writerFunc(&calls),
+			}
+
+			allResults := makeResults(tt.resultCount)
+
+			var ctx context.Context
+			var ch <-chan *jobRunResult
+			if tt.cancelAfter > 0 {
+				cancelCtx, cancel := context.WithCancel(context.Background())
+				defer cancel()
+				buf := make(chan *jobRunResult, tt.resultCount)
+				for i, r := range allResults {
+					if i == tt.cancelAfter {
+						cancel()
+					}
+					buf <- r
+				}
+				close(buf)
+				ctx = cancelCtx
+				ch = buf
+			} else {
+				ctx = context.Background()
+				ch = sendResults(allResults)
+			}
+
+			pl.accumulateAndWriteJobRuns(ctx, ch)
+
+			assert.Len(t, pl.errors, tt.wantErrorCount, "unexpected error count")
+			for _, err := range pl.errors {
+				assert.Contains(t, err.Error(), tt.wantErrorMsg)
+			}
+		})
+	}
+}
+
+func TestAccumulateAndWriteJobRuns_PerBatchErrors(t *testing.T) {
+	failOnBatch := map[int]string{
+		1: "timeout on batch 1",
+		3: "constraint violation on batch 3",
+	}
+	callCount := 0
+	var calls [][]jobRunResult
+
+	pl := &ProwLoader{
+		batchWriter: func(ctx context.Context, batch []jobRunResult) error {
+			cp := make([]jobRunResult, len(batch))
+			copy(cp, batch)
+			calls = append(calls, cp)
+			callCount++
+			if msg, ok := failOnBatch[callCount]; ok {
+				return fmt.Errorf("%s", msg)
+			}
+			return nil
+		},
+	}
+
+	results := sendResults(makeResults(400))
+	pl.accumulateAndWriteJobRuns(context.Background(), results)
+
+	assert.Len(t, calls, 4, "all 4 batches should be attempted")
+	assert.Len(t, pl.errors, 2, "each failed batch should produce its own error")
+	assert.Contains(t, pl.errors[0].Error(), "timeout on batch 1")
+	assert.Contains(t, pl.errors[1].Error(), "constraint violation on batch 3")
+}
+
+func TestAccumulateAndWriteJobRuns_CancelledContextDoesNotSilentlyDrop(t *testing.T) {
+	var writtenRuns int
+	pl := &ProwLoader{
+		batchWriter: func(ctx context.Context, batch []jobRunResult) error {
+			writtenRuns += len(batch)
+			return nil
+		},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	results := make(chan *jobRunResult, 10)
+	for i := range 5 {
+		results <- &jobRunResult{Run: prowJobRunRow{ID: uint(i + 1)}}
+	}
+	cancel()
+	for i := range 5 {
+		results <- &jobRunResult{Run: prowJobRunRow{ID: uint(i + 6)}}
+	}
+	close(results)
+
+	pl.accumulateAndWriteJobRuns(ctx, results)
+
+	// The result received on the same iteration as ctx cancellation should
+	// not be silently dropped. It should appear in either the written runs
+	// or the errors (if the cancelled-context write fails).
+	totalAccountedFor := writtenRuns + countFailedRuns(pl.errors)
+	assert.Greater(t, totalAccountedFor, 0, "at least the results received before cancellation should be accounted for")
+}
+
+func countFailedRuns(errs []error) int {
+	return len(errs) * 100
+}

pkg/dataloader/prowloader/prow.go

@@ -55,6 +55,11 @@ import (
 // from the path "/view/gs/origin-ci-test/logs/periodic-ci-openshift-release-master-nightly-4.14-e2e-gcp-sdn/1737420379221135360"
 var gcsPathStrip = regexp.MustCompile(`.*/gs/[^/]+/`)
 
+// batchWriterFunc writes a batch of job run results to the database.
+// It is a field on ProwLoader so tests can inject a stub without
+// needing a real database connection.
+type batchWriterFunc func(ctx context.Context, batch []jobRunResult) error
+
 type ProwLoader struct {
 	ctx                          context.Context
 	dbc                          *db.DB
@@ -75,6 +80,7 @@ type ProwLoader struct {
 	promPusher                   *push.Pusher
 	loadSince                    *time.Time
 	labelsCache                  map[string]pq.StringArray
+	batchWriter                  batchWriterFunc
 }
 
 func New(
@@ -106,7 +112,7 @@ func New(
 		}
 	}
 
-	return &ProwLoader{
+	pl := &ProwLoader{
 		ctx:                          ctx,
 		dbc:                          dbc,
 		gcsClient:                    gcsClient,
@@ -124,6 +130,8 @@ func New(
 		promPusher:                   promPusher,
 		loadSince:                    loadSince,
 	}
+	pl.batchWriter = pl.writeJobRunBatch
+	return pl
 }
 
 const DefaultLookbackDays = 14
@@ -313,10 +321,7 @@ func (pl *ProwLoader) Load() {
 		close(fetchErrsCh)
 	}()
 
-	if err := pl.accumulateAndWriteJobRuns(pl.ctx, results); err != nil {
-		cancelFetch()
-		pl.errors = append(pl.errors, errors.Wrap(err, "error writing job runs"))
-	}
+	pl.accumulateAndWriteJobRuns(pl.ctx, results)
 
 	for err := range fetchErrsCh {
 		pl.errors = append(pl.errors, err)
@@ -1206,35 +1211,47 @@ var (
 	}
 )
 
-func (pl *ProwLoader) accumulateAndWriteJobRuns(ctx context.Context, results <-chan *jobRunResult) error {
+func (pl *ProwLoader) accumulateAndWriteJobRuns(ctx context.Context, results <-chan *jobRunResult) {
 	const flushThreshold = 100
 	var (
-		batch []jobRunResult
-		total int
+		batch  []jobRunResult
+		total  int
+		failed int
 	)
 
+	flush := func(msg string) {
+		if err := pl.batchWriter(ctx, batch); err != nil {
+			log.WithError(err).WithField("batchSize", len(batch)).Warning(msg)
+			failed += len(batch)
+			pl.errors = append(pl.errors, fmt.Errorf("error writing job run batch: %w", err))
+		} else {
+			total += len(batch)
+		}
+		batch = batch[:0]
+	}
+
 	for result := range results {
 		batch = append(batch, *result)
+		if ctx.Err() != nil {
+			break
+		}
 		if len(batch) >= flushThreshold {
-			if err := pl.writeJobRunBatch(ctx, batch); err != nil {
-				return err
-			}
-			total += len(batch)
-			batch = batch[:0]
+			flush("batch write failed, continuing with remaining batches")
 		}
 	}
 	if len(batch) > 0 {
-		if err := pl.writeJobRunBatch(ctx, batch); err != nil {
-			return err
-		}
-		total += len(batch)
+		flush("final batch write failed")
 	}
 
-	if total > 0 {
-		log.WithField("runs", total).Info("all job run batches committed")
+	if total > 0 || failed > 0 {
+		entry := log.WithField("succeeded", total).WithField("failed", failed)
+		if failed > 0 {
+			entry.Warning("job run batch processing completed with errors")
+		} else {
+			entry.Info("all job run batches committed")
+		}
 	}
 	prowLoaderProcessedMetricGauge.Set(float64(total))
-	return nil
 }
 
 func (pl *ProwLoader) writeJobRunBatch(ctx context.Context, batch []jobRunResult) error {