Bug bot comments

2026-09-11 16:38:01 -07:00
parent 9fd1868617
commit 0202cf33df
3 changed files with 166 additions and 4 deletions
@@ -139,6 +139,18 @@ func (p *Checkpointer) Complete(_ context.Context, message string) error {
 	return nil
 }

+// ResumeFrom pins the resume point at the start of the given bucket, discarding
+// any finer-grained progress. An interrupted pass uses this so resumption cannot
+// skip past a bucket that still needs another attempt. Only meaningful for a
+// bucket-sequence checkpointer; on a unit scan there is no sequence to rewind and
+// this would clear the unit's resume point instead.
+func (p *Checkpointer) ResumeFrom(bucket, role string) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return p.updateCheckpoint(bucket, role, "")
+}
+
 // UpdateObjectCompletion records successfully processed objects within the current page
 // and maintains fine-grained resumption checkpoints. It uses a conservative tracking
 // strategy that ensures no objects are missed by only checkpointing consecutively
@@ -403,8 +403,12 @@ func (s *Source) scanBuckets(
 	}

 	// Collected rather than returned immediately so one bad bucket does not cut
-	// the pass short.
-	var bucketErrs []error
+	// the pass short; the first failure also anchors resumption if the pass is
+	// interrupted before those buckets get a retry.
+	var (
+		bucketErrs        []error
+		firstFailedBucket string
+	)

 	bucketsToScanCount := len(bucketsToScan)
 	for bucketIdx := pos.index; bucketIdx < bucketsToScanCount; bucketIdx++ {
@@ -433,12 +437,21 @@ func (s *Source) scanBuckets(
 		*totalObjectCount += objectCount
 		if err != nil {
 			if isContextCancellation(err) {
-				// Return before the completion call below so resume info survives
-				// and the next run picks up here instead of at the first bucket.
 				ctx.Logger().V(3).Info("bucket scan interrupted", "bucket", bucket, "err", err)
+				// Returns before the completion call below so resume info survives,
+				// rewound to the earliest failure first: resuming past a bucket that
+				// still needs another attempt drops it from the scan entirely.
+				if firstFailedBucket != "" {
+					if rewindErr := checkpointer.ResumeFrom(firstFailedBucket, role); rewindErr != nil {
+						ctx.Logger().Error(rewindErr, "could not rewind resume point to failed bucket", "bucket", firstFailedBucket)
+					}
+				}
 				return errors.Join(bucketErrs...)
 			}
 			bucketErrs = append(bucketErrs, err)
+			if firstFailedBucket == "" {
+				firstFailedBucket = bucket
+			}
 			continue
 		}
 	}
@@ -479,6 +492,12 @@ func (s *Source) scanBucket(

 	regionalClient, err := s.getRegionalClientForBucket(ctx, client, role, bucket)
 	if err != nil {
+		// Checked ahead of the split below: a cancellation says nothing about the
+		// bucket, so classifying it as an expected denial would report an
+		// interrupted bucket as a finished one.
+		if isContextCancellation(err) {
+			return 0, err
+		}
 		// Same expected-vs-fatal split as listing below: enumeration walks every
 		// bucket and will fail region lookup on many of them.
 		if s.listErrorsAreExpected() {
@@ -501,6 +520,11 @@ func (s *Source) scanBucket(
 	for paginator.HasMorePages() {
 		output, err := paginator.NextPage(ctx)
 		if err != nil {
+			// Checked before the expected-vs-fatal split for the same reason as the
+			// region lookup above.
+			if isContextCancellation(err) {
+				return objectCount, err
+			}
 			s.metricsCollector.RecordBucketListError(bucket, role)
 			if s.listErrorsAreExpected() {
 				// Scanning without naming buckets is supported, and the identity is
@@ -119,6 +119,11 @@ type listObjectsServer struct {
 	// before issuing requests; the handler reads it unlocked.
 	inventory []string

+	// onRequest runs while a listing request is in flight, letting a test cancel
+	// the scan mid-call rather than before it starts. Set it before issuing
+	// requests; the handler reads it unlocked.
+	onRequest func(bucket string)
+
 	mu        sync.Mutex
 	requested []string
 }
@@ -156,6 +161,10 @@ func newListObjectsServer(t *testing.T, deniedBuckets ...string) *listObjectsSer
 		srv.requested = append(srv.requested, bucket)
 		srv.mu.Unlock()

+		if srv.onRequest != nil {
+			srv.onRequest(bucket)
+		}
+
 		if _, isDenied := denied[bucket]; isDenied {
 			w.WriteHeader(http.StatusForbidden)
 			_, _ = fmt.Fprint(w, `<Error><Code>AccessDenied</Code><Message>Access Denied</Message></Error>`)
@@ -399,6 +408,123 @@ func TestChunksDiscoveredBucketDenialsAreHealthy(t *testing.T) {
 		"the denied bucket must be skipped, not abort discovery scanning")
 }

+// cancelWhileListing makes the server cancel the scan as a listing request
+// arrives, so the error surfaces from inside the AWS SDK instead of from the
+// check at the top of scanBucket. The pause gives the client time to notice the
+// cancellation rather than reading a complete response first.
+func cancelWhileListing(srv *listObjectsServer, target string, cancel context.CancelFunc) {
+	srv.onRequest = func(bucket string) {
+		if target != "" && bucket != target {
+			return
+		}
+		cancel()
+		time.Sleep(50 * time.Millisecond)
+	}
+}
+
+// TestChunkUnitDiscoveryCancellationKeepsResumeInfo covers cancellation arriving
+// while a discovered bucket is being listed. Denials are expected on that path,
+// so without a cancellation check first the interrupted bucket looks like a
+// completed one and loses the checkpoint it had already written.
+func TestChunkUnitDiscoveryCancellationKeepsResumeInfo(t *testing.T) {
+	const resumeKey = "objects/last-scanned-key"
+
+	srv := newListObjectsServer(t)
+	s := newTestSource(t, srv) // no configured buckets: discovery mode
+
+	unit := S3SourceUnit{Bucket: "bucket-a"}
+	unitID, _ := unit.SourceUnitID()
+	s.SetEncodedResumeInfoFor(unitID, resumeKey)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cancelWhileListing(srv, "", cancel)
+
+	// Cancellation is still a clean stop, but the unit did not finish, so its
+	// resume point has to survive.
+	require.NoError(t, s.ChunkUnit(ctx, unit, sources.ChanReporter{Ch: make(chan *sources.Chunk, 1)}))
+	assert.Equal(t, resumeKey, s.GetEncodedResumeInfoFor(unitID))
+}
+
+// TestScanBucketsInterruptRewindsResumeToFirstFailure guards the seam between
+// collecting bucket failures and stopping early. The first bucket is denied, the
+// second checkpoints an object and is then interrupted; without a rewind the
+// checkpoint points at the second bucket, so the retry starts past the denied one
+// and it silently drops out of the scan.
+func TestScanBucketsInterruptRewindsResumeToFirstFailure(t *testing.T) {
+	const (
+		deniedBucket      = "a-denied"
+		interruptedBucket = "b-interrupted"
+	)
+	buckets := []string{deniedBucket, interruptedBucket, "c-unreached"}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	page1 := `<?xml version="1.0" encoding="UTF-8"?>` +
+		`<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">` +
+		`<Name>` + interruptedBucket + `</Name><KeyCount>1</KeyCount><IsTruncated>true</IsTruncated>` +
+		`<NextContinuationToken>page-2-token</NextContinuationToken>` +
+		`<Contents><Key>obj-b.txt</Key><Size>12</Size>` +
+		`<LastModified>2026-01-01T00:00:00.000Z</LastModified></Contents></ListBucketResult>`
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		path := strings.Trim(r.URL.Path, "/")
+
+		if path == interruptedBucket+"/obj-b.txt" {
+			w.Header().Set("Content-Type", "text/plain")
+			_, _ = fmt.Fprint(w, "hello secret")
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/xml")
+		if path == deniedBucket {
+			w.WriteHeader(http.StatusForbidden)
+			_, _ = fmt.Fprint(w, `<Error><Code>AccessDenied</Code><Message>Access Denied</Message></Error>`)
+			return
+		}
+
+		// Page one checkpoints obj-b.txt; the follow-up page is where the scan is
+		// interrupted, so the checkpoint is already sitting on the second bucket.
+		if r.URL.Query().Get("continuation-token") == "page-2-token" {
+			cancel()
+			time.Sleep(50 * time.Millisecond)
+			return
+		}
+		if path == interruptedBucket {
+			_, _ = fmt.Fprint(w, page1)
+			return
+		}
+		_, _ = fmt.Fprintf(w, `<?xml version="1.0" encoding="UTF-8"?>`+
+			`<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">`+
+			`<Name>%s</Name><KeyCount>0</KeyCount><IsTruncated>false</IsTruncated></ListBucketResult>`, path)
+	}))
+	t.Cleanup(srv.Close)
+
+	conn, err := anypb.New(&sourcespb.S3{
+		Credential: &sourcespb.S3_Unauthenticated{},
+		Endpoint:   srv.URL,
+		Buckets:    buckets,
+	})
+	require.NoError(t, err)
+
+	s := &Source{}
+	require.NoError(t, s.Init(context.Background(), "s3 test source", 0, 0, false, conn, 1))
+
+	client, err := s.newClient(context.Background(), s.defaultRegion(), "")
+	require.NoError(t, err)
+
+	var totalObjectCount uint64
+	err = s.scanBuckets(ctx, client, "", buckets, make(chan *sources.Chunk, 8), &totalObjectCount)
+
+	require.Error(t, err)
+	assert.ErrorContains(t, err, `could not list objects in configured bucket "a-denied"`)
+	assert.Contains(t, s.EncodedResumeInfo, deniedBucket,
+		"resume must rewind to the failed bucket so the retry attempts it again")
+	assert.NotContains(t, s.EncodedResumeInfo, interruptedBucket,
+		"resume must not start past a bucket that still needs another attempt")
+}
+
 // 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.