Merge #150161 #150165

150161: scripts: add basic pcr setup/admin wrapper r=dt a=dt

like the 'ldr' script, this script automates the process of making two clusters and provides a quick mechanism for issuing commands to both of them.

Release note: none.
Epic: none.

150165: sql: ensure that index backfill monitor is closed in error case r=yuzefovich a=yuzefovich

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.

Fixes: #150118.

Release note: None

Co-authored-by: David Taylor <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: 3 people

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
 	}
 

scripts/pcr

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+export COCKROACH_USER=root
+
+A="$USER-a"
+B="$USER-b"
+
+
+if [ "$#" -lt 1 ]; then
+  cat << EOF
+usage: $0 <command>
+EOF
+  exit 0
+fi
+
+
+case $1 in
+  "setup")
+    shift
+    $0 create "$@"
+    $0 go
+    ;;
+
+  "create")
+    shift
+    roachprod create $A \
+      --clouds gce --gce-machine-type n2-standard-16 --nodes 4 --username "$USER" --local-ssd=false --gce-pd-volume-size 1000 --lifetime 24h "$@" 
+    roachprod create $B \
+      --clouds gce --gce-machine-type n2-standard-16 --nodes 4 --username "$USER" --local-ssd=false --gce-pd-volume-size 1000 --lifetime 24h "$@"
+    $0 stage cockroach
+    $0 stage workload
+    ;;
+  
+  "init")
+    $0 start
+    roachprod sql --cluster=system $B:1 -- -e "CREATE EXTERNAL CONNECTION IF NOT EXISTS a AS $(roachprod pgurl --cluster=system $A:1)"
+    roachprod sql --cluster=system $A:1 -- -e "CREATE VIRTUAL CLUSTER main;"
+    roachprod sql --cluster=system $B:1 -- -e "CREATE VIRTUAL CLUSTER main FROM REPLICATION OF main ON 'external://a' WITH READ VIRTUAL CLUSTER;"
+    roachprod sql --cluster=system $A:1 -- -e "ALTER VIRTUAL CLUSTER main START SERVICE SHARED;"
+    roachprod sql --cluster=system $A:1 -- -e "SET CLUSTER SETTING server.controller.default_target_cluster = 'main';"
+    roachprod sql --cluster=system $B:1 -- -e "SET CLUSTER SETTING server.controller.default_target_cluster = 'main';"
+    ;;
+  *)
+    cmd="${1}"
+    shift
+
+    # We're going to run the same command against A and B, but note that we have
+    # set -e above which normally would cause the first to stop the script if it
+    # exited non-zero. So we capture the result in an `||` so we keep going to
+    # the second one, then if we're still running, exit with the first's result.
+    ret=0
+
+    echo "${A}:"
+    roachprod "${cmd}" $A "$@" || ret=$?
+    echo "${B}:"
+    roachprod "${cmd}" $B "$@"
+    exit $ret
+  ;;
+esac