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

We just saw a test failure where the `index-backfill-mon` memory monitor
wasn't stopped on the server shutdown. I've audited the code, and
I think it likely happened in an error case during the initialization
where we'd previously forget to stop it. This is now fixed.

Additionally, this commit lifts the assumption of the monitor being
non-nil a bit higher in the call stack as well as removes some unused
error return arguments.

Release note: None

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
 	}