Merge pull request #147564 from rafiss/blathers/backport-release-25.2.1-rc-147511

release-25.2.1-rc: sql: fix memory leak in index backfill progress tracking

Author: rafiss

pkg/sql/execinfra/server_config.go

@@ -245,7 +245,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)