diff --git a/pkg/sources/s3/s3.go b/pkg/sources/s3/s3.go index ce91ff0cc..4c085ccd0 100644 --- a/pkg/sources/s3/s3.go +++ b/pkg/sources/s3/s3.go @@ -647,8 +647,10 @@ func (s *Source) pageChunker( ctx = context.WithValues(ctx, "bucket", metadata.bucket, "page_number", metadata.pageNumber) for objIdx, obj := range metadata.page.Contents { octx := context.WithValues(ctx, "key", *obj.Key, "size", *obj.Size) + // break, not return: the pool must drain before returning so an in-flight + // worker cannot move the checkpoint after the caller rewinds resume state. if common.IsDone(octx) { - return + break } // Skip objects excluded by the configured prefixes or extensions. diff --git a/pkg/sources/s3/s3_test.go b/pkg/sources/s3/s3_test.go index b926c5cc9..462026243 100644 --- a/pkg/sources/s3/s3_test.go +++ b/pkg/sources/s3/s3_test.go @@ -10,6 +10,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "testing" "time" @@ -525,6 +526,84 @@ func TestScanBucketsInterruptRewindsResumeToFirstFailure(t *testing.T) { "resume must not start past a bucket that still needs another attempt") } +// pausingCollector parks directory-skip recordings until the test releases them. +// The collector is the one seam inside an object worker that a cancelled context +// cannot unblock, so it stands in for a worker mid-HandleFile. +type pausingCollector struct { + metricsCollector + entered chan struct{} + release chan struct{} + drained atomic.Int32 +} + +func (p *pausingCollector) RecordObjectSkipped(string, string, float64) { + p.entered <- struct{}{} + <-p.release + p.drained.Add(1) +} + +// TestPageChunkerDrainsWorkersOnCancellation pins the invariant the resume +// rewind depends on: pageChunker must not return while object workers are still +// running, or a straggler's checkpoint write could land after scanBuckets has +// rewound resume state to an earlier failed bucket. +func TestPageChunkerDrainsWorkersOnCancellation(t *testing.T) { + conn, err := anypb.New(&sourcespb.S3{Credential: &sourcespb.S3_Unauthenticated{}}) + require.NoError(t, err) + + s := &Source{} + // Two worker slots so both blocking objects are in flight when the scan is + // cancelled between spawns. + require.NoError(t, s.Init(context.Background(), "s3 test source", 0, 0, false, conn, 2)) + + collector := &pausingCollector{ + metricsCollector: s.metricsCollector, + entered: make(chan struct{}, 4), + release: make(chan struct{}, 4), + } + s.metricsCollector = collector + + // Directory keys route each worker straight to the collector without touching + // the network, keeping the test free of S3 stubs. + var contents []s3types.Object + for _, key := range []string{"w0/", "w1/", "w2/", "w3/"} { + contents = append(contents, s3types.Object{Key: aws.String(key), Size: aws.Int64(1)}) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var objectCount, filteredCount uint64 + done := make(chan struct{}) + go func() { + defer close(done) + s.pageChunker( + ctx, + pageMetadata{bucket: "bucket", pageNumber: 1, page: &awss3.ListObjectsV2Output{Contents: contents}}, + processingState{errorCount: &sync.Map{}, objectCount: &objectCount, filteredCount: &filteredCount}, + sources.ChanReporter{Ch: make(chan *sources.Chunk, 4)}, + NewCheckpointer(ctx, &s.Progress, false), + ) + }() + + // Both slots are parked in the collector; cancel mid-page, then free one so + // the loop reaches its context check with the other worker still in flight. + <-collector.entered + <-collector.entered + cancel() + collector.release <- struct{}{} + + select { + case <-done: + t.Fatal("pageChunker returned while a worker was still in flight") + case <-time.After(200 * time.Millisecond): + } + + close(collector.release) + <-done + assert.EqualValues(t, 2, collector.drained.Load(), + "every in-flight worker must finish before pageChunker returns") +} + // TestChunkUnitMidBucketFailureKeepsCheckpointedResume exercises the real // checkpointer rather than a pre-seeded resume value: page one is scanned, page // two is denied, and the checkpoint must survive so a retry does not start over.