Merge pull request #151430 from yuzefovich/backport25.3-150165

release-25.3: sql: ensure that index backfill monitor is closed in error case

Author: yuzefovich

pkg/sql/backfill.go

@@ -1043,10 +1043,7 @@ func (sc *SchemaChanger) distIndexBackfill(
 		)
 		indexBatchSize := indexBackfillBatchSize.Get(&sc.execCfg.Settings.SV)
 		chunkSize := sc.getChunkSize(indexBatchSize)
-		spec, err := initIndexBackfillerSpec(*tableDesc.TableDesc(), writeAsOf, writeAtRequestTimestamp, chunkSize, addedIndexes, 0)
-		if err != nil {
-			return err
-		}
+		spec := initIndexBackfillerSpec(*tableDesc.TableDesc(), writeAsOf, writeAtRequestTimestamp, chunkSize, addedIndexes, 0)
 		p, err = sc.distSQLPlanner.createBackfillerPhysicalPlan(ctx, planCtx, spec, todoSpans)
 		return err
 	}); err != nil {
@@ -1351,10 +1348,7 @@ func (sc *SchemaChanger) distColumnBackfill(
 			planCtx := sc.distSQLPlanner.NewPlanningCtx(
 				ctx, &evalCtx, nil /* planner */, txn.KV(), FullDistribution,
 			)
-			spec, err := initColumnBackfillerSpec(tableDesc, duration, chunkSize, backfillUpdateChunkSizeThresholdBytes, readAsOf)
-			if err != nil {
-				return err
-			}
+			spec := initColumnBackfillerSpec(tableDesc, duration, chunkSize, backfillUpdateChunkSizeThresholdBytes, readAsOf)
 			plan, err := sc.distSQLPlanner.createBackfillerPhysicalPlan(ctx, planCtx, spec, todoSpans)
 			if err != nil {
 				return err
@@ -2917,11 +2911,9 @@ func indexBackfillInTxn(
 	tableDesc catalog.TableDescriptor,
 	traceKV bool,
 ) error {
-	var indexBackfillerMon *mon.BytesMonitor
-	if evalCtx.Planner.Mon() != nil {
-		indexBackfillerMon = execinfra.NewMonitor(ctx, evalCtx.Planner.Mon(),
-			mon.MakeName("local-index-backfill-mon"))
-	}
+	indexBackfillerMon := execinfra.NewMonitor(
+		ctx, evalCtx.Planner.Mon(), mon.MakeName("local-index-backfill-mon"),
+	)
 
 	var backfiller backfill.IndexBackfiller
 	if err := backfiller.InitForLocalUse(

pkg/sql/backfill/backfill.go

@@ -595,13 +595,25 @@ func (ib *IndexBackfiller) ContainsInvertedIndex() bool {
 // InitForLocalUse initializes an IndexBackfiller for use during local execution
 // within a transaction. In this case, the entire backfill process is occurring
 // on the gateway as part of the user's transaction.
+//
+// Non-nil memory monitor must be provided. If an error is returned, it'll be
+// stopped automatically; otherwise, the backfiller takes ownership of the
+// monitor.
 func (ib *IndexBackfiller) InitForLocalUse(
 	ctx context.Context,
 	evalCtx *eval.Context,
 	semaCtx *tree.SemaContext,
 	desc catalog.TableDescriptor,
 	mon *mon.BytesMonitor,
-) error {
+) (retErr error) {
+	if mon == nil {
+		return errors.AssertionFailedf("memory monitor must be provided")
+	}
+	defer func() {
+		if retErr != nil {
+			mon.Stop(ctx)
+		}
+	}()
 
 	// Initialize ib.added.
 	if err := ib.initIndexes(ctx, evalCtx, desc, nil /* allowList */, 0 /*sourceIndex*/, nil); err != nil {
@@ -626,7 +638,8 @@ func (ib *IndexBackfiller) InitForLocalUse(
 		ib.valNeededForCol.Add(ib.colIdxMap.GetDefault(col))
 	})
 
-	return ib.init(evalCtx, predicates, colExprs, mon)
+	ib.init(evalCtx, predicates, colExprs, mon)
+	return nil
 }
 
 // constructExprs is a helper to construct the index and column expressions
@@ -737,14 +750,27 @@ func constructExprs(
 // backfill operation manages its own transactions. This separation is necessary
 // due to the different procedure for accessing user defined type metadata as
 // part of a distributed flow.
+//
+// Non-nil memory monitor must be provided. If an error is returned, it'll be
+// stopped automatically; otherwise, the backfiller takes ownership of the
+// monitor.
 func (ib *IndexBackfiller) InitForDistributedUse(
 	ctx context.Context,
 	flowCtx *execinfra.FlowCtx,
 	desc catalog.TableDescriptor,
 	allowList []catid.IndexID,
 	sourceIndexID catid.IndexID,
 	mon *mon.BytesMonitor,
-) error {
+) (retErr error) {
+	if mon == nil {
+		return errors.AssertionFailedf("memory monitor must be provided")
+	}
+	defer func() {
+		if retErr != nil {
+			mon.Stop(ctx)
+		}
+	}()
+
 	// We'll be modifying the eval.Context in BuildIndexEntriesChunk, so we need
 	// to make a copy.
 	evalCtx := flowCtx.NewEvalCtx()
@@ -801,7 +827,8 @@ func (ib *IndexBackfiller) InitForDistributedUse(
 		ib.valNeededForCol.Add(ib.colIdxMap.GetDefault(col))
 	})
 
-	return ib.init(evalCtx, predicates, colExprs, mon)
+	ib.init(evalCtx, predicates, colExprs, mon)
+	return nil
 }
 
 // Close releases the resources used by the IndexBackfiller. It can be called
@@ -933,12 +960,15 @@ func (ib *IndexBackfiller) initIndexes(
 }
 
 // init completes the initialization of an IndexBackfiller.
+//
+// The IndexBackfiller takes ownership of the monitor which must be non-nil.
+// It'll be closed when the backfiller is closed.
 func (ib *IndexBackfiller) init(
 	evalCtx *eval.Context,
 	predicateExprs map[descpb.IndexID]tree.TypedExpr,
 	colExprs map[descpb.ColumnID]tree.TypedExpr,
 	mon *mon.BytesMonitor,
-) error {
+) {
 	ib.evalCtx = evalCtx
 	ib.predicates = predicateExprs
 	ib.colExprs = colExprs
@@ -959,12 +989,8 @@ func (ib *IndexBackfiller) init(
 	}
 
 	// Create a bound account associated with the index backfiller monitor.
-	if mon == nil {
-		return errors.AssertionFailedf("no memory monitor linked to IndexBackfiller during init")
-	}
 	ib.mon = mon
 	ib.muBoundAccount.boundAccount = mon.MakeBoundAccount()
-	return nil
 }
 
 // BuildIndexEntriesChunk reads a chunk of rows from a table using the span sp

pkg/sql/distsql_plan_backfill.go

@@ -28,15 +28,15 @@ func initColumnBackfillerSpec(
 	chunkSize int64,
 	updateChunkSizeThresholdBytes uint64,
 	readAsOf hlc.Timestamp,
-) (execinfrapb.BackfillerSpec, error) {
+) execinfrapb.BackfillerSpec {
 	return execinfrapb.BackfillerSpec{
 		Table:                         *tbl.TableDesc(),
 		Duration:                      duration,
 		ChunkSize:                     chunkSize,
 		UpdateChunkSizeThresholdBytes: updateChunkSizeThresholdBytes,
 		ReadAsOf:                      readAsOf,
 		Type:                          execinfrapb.BackfillerSpec_Column,
-	}, nil
+	}
 }
 
 func initIndexBackfillerSpec(
@@ -46,7 +46,7 @@ func initIndexBackfillerSpec(
 	chunkSize int64,
 	indexesToBackfill []descpb.IndexID,
 	sourceIndexID descpb.IndexID,
-) (execinfrapb.BackfillerSpec, error) {
+) execinfrapb.BackfillerSpec {
 	return execinfrapb.BackfillerSpec{
 		Table:                 desc,
 		WriteAsOf:             writeAsOf,
@@ -55,21 +55,21 @@ func initIndexBackfillerSpec(
 		ChunkSize:             chunkSize,
 		IndexesToBackfill:     indexesToBackfill,
 		SourceIndexID:         sourceIndexID,
-	}, nil
+	}
 }
 
 func initIndexBackfillMergerSpec(
 	desc descpb.TableDescriptor,
 	addedIndexes []descpb.IndexID,
 	temporaryIndexes []descpb.IndexID,
 	mergeTimestamp hlc.Timestamp,
-) (execinfrapb.IndexBackfillMergerSpec, error) {
+) execinfrapb.IndexBackfillMergerSpec {
 	return execinfrapb.IndexBackfillMergerSpec{
 		Table:            desc,
 		AddedIndexes:     addedIndexes,
 		TemporaryIndexes: temporaryIndexes,
 		MergeTimestamp:   mergeTimestamp,
-	}, nil
+	}
 }
 
 var initialSplitsPerProcessor = settings.RegisterIntSetting(

pkg/sql/index_backfiller.go

@@ -207,13 +207,11 @@ func (ib *IndexBackfillPlanner) plan(
 		// batch size. Also plumb in a testing knob.
 		chunkSize := indexBackfillBatchSize.Get(&ib.execCfg.Settings.SV)
 		const writeAtRequestTimestamp = true
-		spec, err := initIndexBackfillerSpec(
+		spec := initIndexBackfillerSpec(
 			*td.TableDesc(), writeAsOf, writeAtRequestTimestamp, chunkSize,
 			indexesToBackfill, sourceIndexID,
 		)
-		if err != nil {
-			return err
-		}
+		var err error
 		p, err = ib.execCfg.DistSQLPlanner.createBackfillerPhysicalPlan(ctx, planCtx, spec, sourceSpans)
 		return err
 	}); err != nil {

pkg/sql/mvcc_backfiller.go

@@ -137,10 +137,8 @@ func (im *IndexBackfillerMergePlanner) plan(
 			ctx, &extEvalCtx, nil /* planner */, txn.KV(), FullDistribution,
 		)
 
-		spec, err := initIndexBackfillMergerSpec(*tableDesc.TableDesc(), addedIndexes, temporaryIndexes, mergeTimestamp)
-		if err != nil {
-			return err
-		}
+		spec := initIndexBackfillMergerSpec(*tableDesc.TableDesc(), addedIndexes, temporaryIndexes, mergeTimestamp)
+		var err error
 		p, err = im.execCfg.DistSQLPlanner.createIndexBackfillerMergePhysicalPlan(ctx, planCtx, spec, todoSpanList)
 		return err
 	}); err != nil {

pkg/sql/rowexec/indexbackfiller.go

@@ -84,8 +84,9 @@ func newIndexBackfiller(
 	processorID int32,
 	spec execinfrapb.BackfillerSpec,
 ) (*indexBackfiller, error) {
-	indexBackfillerMon := execinfra.NewMonitor(ctx, flowCtx.Cfg.BackfillerMonitor,
-		mon.MakeName("index-backfill-mon"))
+	indexBackfillerMon := execinfra.NewMonitor(
+		ctx, flowCtx.Cfg.BackfillerMonitor, mon.MakeName("index-backfill-mon"),
+	)
 	ib := &indexBackfiller{
 		desc:        flowCtx.TableDescriptor(ctx, &spec.Table),
 		spec:        spec,
@@ -94,8 +95,9 @@ func newIndexBackfiller(
 		filter:      backfill.IndexMutationFilter,
 	}
 
-	if err := ib.IndexBackfiller.InitForDistributedUse(ctx, flowCtx, ib.desc,
-		ib.spec.IndexesToBackfill, ib.spec.SourceIndexID, indexBackfillerMon); err != nil {
+	if err := ib.IndexBackfiller.InitForDistributedUse(
+		ctx, flowCtx, ib.desc, ib.spec.IndexesToBackfill, ib.spec.SourceIndexID, indexBackfillerMon,
+	); err != nil {
 		return nil, err
 	}