Merge #147511

147511: sql: fix memory leak in index backfill progress tracking r=rafiss a=rafiss

Previously, when resuming an index backfill job, the completed spans
from the job's checkpoint would be repeatedly added to the progress
tracker on each update, causing unbounded memory growth. For large
tables with many ranges, this could exhaust available memory and
cause the backfill to fail.

This change fixes the memory leak by initializing the completed spans
group once with the checkpointed progress at job startup, then only
adding new spans from subsequent batches. A test-only assertion is
added to detect overlapping spans in progress updates.

The test is also updated to verify that checkpointed spans properly
enclose all completed work and to remove flaky timing dependencies.
These test changes also address the test issues that were present since
the test was first added in 058a7f5.
The testing changes include:
- Use a context in the testing hook so that when the job is paused, the
  testing hook does not keep blocking.
- Fixed a race condition by getting the progress from the persisted
  progress, rather than a different hook.
- Avoid using `require` / `t.Fatal` from inside a goroutine, which
  violates the expectations of testing.T.

fixes #140358
fixes #147209

Release note (bug fix): Fixed a memory leak in index backfill jobs
where completed spans were duplicated in memory on each progress
update after resuming from a checkpoint. This could cause out-of-memory
errors when backfilling indexes on large tables with many ranges.


Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/sql/execinfra/server_config.go

@@ -244,7 +244,7 @@ type TestingKnobs struct {
 
 	// RunBeforeIndexBackfillProgressUpdate is called before updating the
 	// progress for a single index backfill.
-	RunBeforeIndexBackfillProgressUpdate func(completed []roachpb.Span)
+	RunBeforeIndexBackfillProgressUpdate func(ctx context.Context, completed []roachpb.Span)
 
 	// SerializeIndexBackfillCreationAndIngestion ensures that every index batch
 	// created during an index backfill is also ingested before moving on to the

pkg/sql/index_backfiller.go

@@ -21,8 +21,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/errors"
 )
 
 // IndexBackfillPlanner holds dependencies for an index backfiller
@@ -77,17 +79,34 @@ func (ib *IndexBackfillPlanner) BackfillIndexes(
 		completed.g.Add(c...)
 		return completed.g.Slice()
 	}
+	// Add spans that were already completed before the job resumed.
+	addCompleted(progress.CompletedSpans...)
 	updateFunc := func(
 		ctx context.Context, meta *execinfrapb.ProducerMetadata,
 	) error {
 		if meta.BulkProcessorProgress == nil {
 			return nil
 		}
-		progress.CompletedSpans = append(progress.CompletedSpans, addCompleted(
-			meta.BulkProcessorProgress.CompletedSpans...)...)
+		progress.CompletedSpans = addCompleted(meta.BulkProcessorProgress.CompletedSpans...)
+		// Make sure the progress update does not contain overlapping spans.
+		// This is a sanity check that only runs in test configurations, since it
+		// is an expensive n^2 check.
+		if buildutil.CrdbTestBuild {
+			for i, span1 := range progress.CompletedSpans {
+				for j, span2 := range progress.CompletedSpans {
+					if i <= j {
+						continue
+					}
+					if span1.Overlaps(span2) {
+						return errors.Newf("progress update contains overlapping spans: %s and %s", span1, span2)
+					}
+				}
+			}
+		}
+
 		knobs := &ib.execCfg.DistSQLSrv.TestingKnobs
 		if knobs.RunBeforeIndexBackfillProgressUpdate != nil {
-			knobs.RunBeforeIndexBackfillProgressUpdate(progress.CompletedSpans)
+			knobs.RunBeforeIndexBackfillProgressUpdate(ctx, progress.CompletedSpans)
 		}
 		return tracker.SetBackfillProgress(ctx, progress)
 	}

pkg/sql/indexbackfiller_test.go

@@ -43,13 +43,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
-	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
@@ -576,61 +575,55 @@ INSERT INTO foo VALUES (1), (10), (100);
 func TestIndexBackfillerResumePreservesProgress(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
-	skip.UnderStress(t, "timing-sensitive test")
 
 	ctx := context.Background()
-	backfillProgressUpdateCh := make(chan struct{})
 	backfillProgressCompletedCh := make(chan []roachpb.Span)
-	checkpointedSpansCh := make(chan []roachpb.Span)
 	const numSpans = 100
 	var isBlockingBackfillProgress atomic.Bool
 	isBlockingBackfillProgress.Store(true)
-	var isBlockingCheckpoint atomic.Bool
 
-	clusterArgs := base.TestClusterArgs{
-		ServerArgs: base.TestServerArgs{
-			Knobs: base.TestingKnobs{
-				DistSQL: &execinfra.TestingKnobs{
-					// We want to push progress every batch_size rows to control
-					// the backfill incrementally.
-					BulkAdderFlushesEveryBatch: true,
-					RunBeforeIndexBackfillProgressUpdate: func(completed []roachpb.Span) {
-						if isBlockingBackfillProgress.Load() {
-							<-backfillProgressUpdateCh
-							backfillProgressCompletedCh <- completed
+	// Start the server with testing knob.
+	tc, db, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
+			DistSQL: &execinfra.TestingKnobs{
+				// We want to push progress every batch_size rows to control
+				// the backfill incrementally.
+				BulkAdderFlushesEveryBatch: true,
+				RunBeforeIndexBackfillProgressUpdate: func(ctx context.Context, completed []roachpb.Span) {
+					if isBlockingBackfillProgress.Load() {
+						select {
+						case <-ctx.Done():
+						case backfillProgressCompletedCh <- completed:
+							t.Logf("before index backfill progress update, completed spans: %v", completed)
 						}
-					},
+					}
 				},
-				SQLDeclarativeSchemaChanger: &scexec.TestingKnobs{
-					RunBeforeBackfill: func(progresses []scexec.BackfillProgress) error {
-						if isBlockingCheckpoint.Load() && progresses != nil && progresses[0].CompletedSpans != nil {
-							checkpointedSpansCh <- progresses[0].CompletedSpans
-						}
-						return nil
-					},
+			},
+			SQLDeclarativeSchemaChanger: &scexec.TestingKnobs{
+				RunBeforeBackfill: func(progresses []scexec.BackfillProgress) error {
+					if progresses != nil {
+						t.Logf("before resuming backfill, checkpointed spans: %v", progresses[0].CompletedSpans)
+					}
+					return nil
 				},
 			},
 		},
-	}
-	// Start the server with a testing knob
-	tc := testcluster.NewTestCluster(t, 1, clusterArgs)
-	tc.Start(t)
+	})
 	defer tc.Stopper().Stop(ctx)
-	db := tc.Conns[0]
 
 	_, err := db.Exec(`SET CLUSTER SETTING bulkio.index_backfill.batch_size = 10`)
 	require.NoError(t, err)
 	// Ensure that we checkpoint our progress to the backfill job so that
 	// RESUMEs can get an up-to-date backfill progress.
-	_, err = db.Exec(`SET CLUSTER SETTING bulkio.index_backfill.checkpoint_interval = '0s'`)
+	_, err = db.Exec(`SET CLUSTER SETTING bulkio.index_backfill.checkpoint_interval = '10ms'`)
 	require.NoError(t, err)
 	_, err = db.Exec(`CREATE TABLE t(i INT PRIMARY KEY)`)
 	require.NoError(t, err)
 	_, err = db.Exec(`INSERT INTO t SELECT generate_series(1, $1)`, numSpans)
 	require.NoError(t, err)
 	_, err = db.Exec(`ALTER TABLE t SPLIT AT TABLE generate_series(1, $1)`, numSpans)
 	require.NoError(t, err)
-	require.NoError(t, tc.WaitForFullReplication())
 	var descID catid.DescID
 	descIDRow := db.QueryRow(`SELECT 't'::regclass::oid`)
 	err = descIDRow.Scan(&descID)
@@ -646,112 +639,118 @@ func TestIndexBackfillerResumePreservesProgress(t *testing.T) {
 		return nil
 	})
 
-	g.GoCtx(func(ctx context.Context) error {
-		testutils.SucceedsSoon(t, func() error {
-			jobIDRow := db.QueryRow(`
+	testutils.SucceedsWithin(t, func() error {
+		jobIDRow := db.QueryRow(`
 				SELECT job_id FROM [SHOW JOBS]
-				WHERE job_type = 'NEW SCHEMA CHANGE' AND description ILIKE '%ADD COLUMN%'`,
-			)
-			if err := jobIDRow.Scan(&jobID); err != nil {
+				WHERE job_type = 'NEW SCHEMA CHANGE' AND description ILIKE '%ADD COLUMN j%'`,
+		)
+		if err := jobIDRow.Scan(&jobID); err != nil {
+			return err
+		}
+		return nil
+	}, 5*time.Second)
+
+	ensureJobState := func(targetState string) {
+		testutils.SucceedsWithin(t, func() error {
+			var jobState string
+			statusRow := db.QueryRow(`SELECT status FROM [SHOW JOB $1]`, jobID)
+			if err := statusRow.Scan(&jobState); err != nil {
 				return err
 			}
+			if jobState != targetState {
+				return errors.Errorf("expected job to be %s, but found status: %s",
+					targetState, jobState)
+			}
 			return nil
-		})
+		}, 5*time.Second)
+	}
 
-		ensureJobState := func(targetState string) {
-			testutils.SucceedsSoon(t, func() error {
-				var jobState string
-				statusRow := db.QueryRow(`SELECT status FROM [SHOW JOB $1]`, jobID)
-				if err := statusRow.Scan(&jobState); err != nil {
-					return err
+	var completedSpans roachpb.SpanGroup
+	receiveProgressUpdate := func() {
+		progressUpdate := <-backfillProgressCompletedCh
+
+		// Make sure the progress update does not contain overlapping spans.
+		for i, span1 := range progressUpdate {
+			for j, span2 := range progressUpdate {
+				if i <= j {
+					continue
 				}
-				if jobState != targetState {
-					return errors.Errorf("expected job to be %s, but found status: %s",
-						targetState, jobState)
+				if span1.Overlaps(span2) {
+					t.Fatalf("progress update contains overlapping spans: %s and %s", span1, span2)
 				}
-				return nil
-			})
+			}
 		}
+		completedSpans.Add(progressUpdate...)
+	}
 
-		// Let the backfill step forward a bit before we do our PAUSE/RESUME
-		// dance.
-		var spansCompletedBeforePause []roachpb.Span
-		for i := 0; i < 2; i++ {
-			backfillProgressUpdateCh <- struct{}{}
-			spansCompletedBeforePause = <-backfillProgressCompletedCh
-		}
+	ensureCompletedSpansAreCheckpointed := func() {
+		testutils.SucceedsWithin(t, func() error {
+			stmt := `SELECT payload FROM crdb_internal.system_jobs WHERE id = $1`
+			var payloadBytes []byte
+			if err := db.QueryRowContext(ctx, stmt, jobID).Scan(&payloadBytes); err != nil {
+				return err
+			}
 
-		_, err := db.Exec(`PAUSE JOB $1`, jobID)
-		if err != nil {
-			return err
-		}
-		ensureJobState("paused")
+			payload := &jobspb.Payload{}
+			if err := protoutil.Unmarshal(payloadBytes, payload); err != nil {
+				return err
+			}
 
-		_, err = db.Exec(`RESUME JOB $1`, jobID)
-		if err != nil {
-			return err
-		}
-		ensureJobState("running")
+			schemaChangeProgress := *(payload.Details.(*jobspb.Payload_NewSchemaChange).NewSchemaChange)
+			var checkpointedSpans []roachpb.Span
+			if len(schemaChangeProgress.BackfillProgress) > 0 {
+				checkpointedSpans = schemaChangeProgress.BackfillProgress[0].CompletedSpans
+			}
+			var sg roachpb.SpanGroup
+			sg.Add(checkpointedSpans...)
+			// Ensure that the spans we already completed are fully contained in our
+			// checkpointed completed spans group.
+			if !sg.Encloses(completedSpans.Slice()...) {
+				return errors.Errorf("checkpointed spans %v do not enclose completed spans %v",
+					checkpointedSpans, completedSpans.Slice())
+			}
 
-		for i := 0; i < 2; i++ {
-			backfillProgressUpdateCh <- struct{}{}
-			<-backfillProgressCompletedCh
-		}
-		isBlockingCheckpoint.Store(true)
+			return nil
+		}, 5*time.Second)
+	}
 
-		_, err = db.Exec(`PAUSE JOB $1`, jobID)
-		if err != nil {
-			return err
-		}
-		ensureJobState("paused")
+	// Let the backfill step forward a bit before we do our PAUSE/RESUME
+	// dance.
+	for i := 0; i < 2; i++ {
+		receiveProgressUpdate()
+	}
 
-		_, err = db.Exec(`RESUME JOB $1`, jobID)
-		if err != nil {
-			return err
-		}
-		ensureJobState("running")
-
-		var wg sync.WaitGroup
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			timer := time.NewTimer(1 * time.Minute)
-			defer timer.Stop()
-			select {
-			case checkpointed := <-checkpointedSpansCh:
-				isBlockingCheckpoint.Store(false)
-				var sg roachpb.SpanGroup
-				sg.Add(checkpointed...)
-				// Ensure that the spans we completed before any PAUSE is
-				// fully contained in our checkpointed completed spans group.
-				require.True(t, sg.Encloses(spansCompletedBeforePause...))
-			case <-timer.C:
-				require.Fail(t, "timed out waiting for checkpoint")
-			}
-		}()
-
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-
-			isBlockingBackfillProgress.Store(false)
-			timer := time.NewTimer(1 * time.Minute)
-			defer timer.Stop()
-			select {
-			case backfillProgressUpdateCh <- struct{}{}:
-				<-backfillProgressCompletedCh
-			case <-timer.C:
-				require.Fail(t, "timed out waiting for backfill progress")
-			}
-		}()
+	ensureCompletedSpansAreCheckpointed()
+	t.Logf("pausing backfill")
+	_, err = db.Exec(`PAUSE JOB $1`, jobID)
+	require.NoError(t, err)
+	ensureJobState("paused")
 
-		// Wait for both operations to complete
-		wg.Wait()
+	t.Logf("resuming backfill")
+	_, err = db.Exec(`RESUME JOB $1`, jobID)
+	require.NoError(t, err)
+	ensureJobState("running")
 
-		// Now we can wait for the job to succeed
-		ensureJobState("succeeded")
-		return nil
-	})
+	// Step forward again before re-pausing.
+	for i := 0; i < 2; i++ {
+		receiveProgressUpdate()
+	}
+
+	ensureCompletedSpansAreCheckpointed()
+	isBlockingBackfillProgress.Store(false)
+
+	t.Logf("pausing backfill")
+	_, err = db.Exec(`PAUSE JOB $1`, jobID)
+	require.NoError(t, err)
+	ensureJobState("paused")
+
+	t.Logf("resuming backfill")
+	_, err = db.Exec(`RESUME JOB $1`, jobID)
+	require.NoError(t, err)
+	ensureJobState("running")
+
+	// Now we can wait for the job to succeed
+	ensureJobState("succeeded")
 
 	if err = g.Wait(); err != nil {
 		require.NoError(t, err)