sql: use MakeSplitterForSideEffect in more places

In #116170 we added span.MakeSplitterForSideEffect which does not
perform a column-family-skipping optimization. We were using this for
locking lookup joins of multi-column-family tables, but forgot to use it
in locking scans, locking index joins, and other places that could be
performing a locked read of a multi-column-family table.

This did not really matter until the previous commit, which made it more
likely that we would use one of these other operations with locking on a
multi-column-family table.

Informs: #116838

Release note: None

Author: michae2

pkg/sql/distsql_physical_planner.go

@@ -3363,7 +3363,15 @@ func (dsp *DistSQLPlanner) planIndexJoin(
 		return err
 	}
 
-	splitter := span.MakeSplitter(planInfo.fetch.desc, index, fetchOrdinals)
+	var splitter span.Splitter
+	// This logic matches opt.Locking.MustLockAllRequestedColumnFamilies.
+	if (joinReaderSpec.LockingStrength != descpb.ScanLockingStrength_FOR_NONE &&
+		joinReaderSpec.LockingDurability == descpb.ScanLockingDurability_GUARANTEED) ||
+		joinReaderSpec.LockingWaitPolicy != descpb.ScanLockingWaitPolicy_BLOCK {
+		splitter = span.MakeSplitterForSideEffect(planInfo.fetch.desc, index, fetchOrdinals)
+	} else {
+		splitter = span.MakeSplitter(planInfo.fetch.desc, index, fetchOrdinals)
+	}
 	joinReaderSpec.SplitFamilyIDs = splitter.FamilyIDs()
 
 	p.PlanToStreamColMap = identityMap(p.PlanToStreamColMap, len(fetchColIDs))
@@ -3458,8 +3466,10 @@ func (dsp *DistSQLPlanner) planLookupJoin(
 	}
 
 	var splitter span.Splitter
-	if joinReaderSpec.LockingStrength != descpb.ScanLockingStrength_FOR_NONE &&
-		joinReaderSpec.LockingDurability == descpb.ScanLockingDurability_GUARANTEED {
+	// This logic matches opt.Locking.MustLockAllRequestedColumnFamilies.
+	if (joinReaderSpec.LockingStrength != descpb.ScanLockingStrength_FOR_NONE &&
+		joinReaderSpec.LockingDurability == descpb.ScanLockingDurability_GUARANTEED) ||
+		joinReaderSpec.LockingWaitPolicy != descpb.ScanLockingWaitPolicy_BLOCK {
 		splitter = span.MakeSplitterForSideEffect(planInfo.fetch.desc, planInfo.fetch.index, fetchOrdinals)
 	} else {
 		splitter = span.MakeSplitter(planInfo.fetch.desc, planInfo.fetch.index, fetchOrdinals)

pkg/sql/distsql_spec_exec_factory.go

@@ -284,7 +284,12 @@ func (e *distSQLSpecExecFactory) ConstructScan(
 	if params.InvertedConstraint != nil {
 		spans, err = sb.SpansFromInvertedSpans(e.ctx, params.InvertedConstraint, params.IndexConstraint, nil /* scratch */)
 	} else {
-		splitter := span.MakeSplitter(tabDesc, idx, params.NeededCols)
+		var splitter span.Splitter
+		if params.Locking.MustLockAllRequestedColumnFamilies() {
+			splitter = span.MakeSplitterForSideEffect(tabDesc, idx, params.NeededCols)
+		} else {
+			splitter = span.MakeSplitter(tabDesc, idx, params.NeededCols)
+		}
 		spans, err = sb.SpansFromConstraint(params.IndexConstraint, splitter)
 	}
 	if err != nil {

pkg/sql/insert_fast_path.go

@@ -112,7 +112,14 @@ func (c *insertFastPathCheck) init(params runParams) error {
 	codec := params.ExecCfg().Codec
 	c.keyPrefix = rowenc.MakeIndexKeyPrefix(codec, c.tabDesc.GetID(), c.idx.GetID())
 	c.spanBuilder.InitAllowingExternalRowData(params.EvalContext(), codec, c.tabDesc, c.idx)
-	c.spanSplitter = span.MakeSplitter(c.tabDesc, c.idx, intsets.Fast{} /* neededColOrdinals */)
+
+	if c.Locking.MustLockAllRequestedColumnFamilies() {
+		c.spanSplitter = span.MakeSplitterForSideEffect(
+			c.tabDesc, c.idx, intsets.Fast{}, /* neededColOrdinals */
+		)
+	} else {
+		c.spanSplitter = span.MakeSplitter(c.tabDesc, c.idx, intsets.Fast{} /* neededColOrdinals */)
+	}
 
 	if len(c.InsertCols) > idx.numLaxKeyCols {
 		return errors.AssertionFailedf(

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -750,3 +750,82 @@ ROLLBACK
 
 statement ok
 RESET enable_durable_locking_for_serializable
+
+# Ensure that we lock all column families of a multi-family table when using
+# durable locking.
+
+statement ok
+CREATE TABLE abc (
+  a INT NOT NULL,
+  b INT NOT NULL,
+  c INT NOT NULL,
+  PRIMARY KEY (a),
+  INDEX (b),
+  FAMILY f0 (a),
+  FAMILY f1 (b),
+  FAMILY f2 (c)
+)
+
+statement ok
+INSERT INTO abc VALUES (6, 7, 8)
+
+statement ok
+SET optimizer_use_lock_op_for_serializable = on
+
+statement ok
+SET enable_durable_locking_for_serializable = on
+
+statement ok
+SET distsql = off
+
+# A scan where we normally skip reading family 0.
+query T kvtrace
+SELECT * FROM abc WHERE a = 6
+----
+Scan /Table/113/1/6/{1/1-2/2}
+
+# But with FOR UPDATE we must read all families.
+query T kvtrace
+SELECT * FROM abc WHERE a = 6 FOR UPDATE
+----
+Scan /Table/113/1/{6-7} lock Exclusive (Block, Replicated)
+
+# An index join where we normally skip reading family 0.
+query T kvtrace
+SELECT * FROM abc WHERE b = 7
+----
+Scan /Table/113/2/{7-8}
+Scan /Table/113/1/6/{1/1-2/2}
+
+# But with FOR UPDATE we must read all families.
+query T kvtrace
+SELECT * FROM abc WHERE b = 7 FOR UPDATE
+----
+Scan /Table/113/2/{7-8}
+Scan /Table/113/1/{6-7} lock Exclusive (Block, Replicated)
+
+statement ok
+INSERT INTO t129145 VALUES (6, 7)
+
+# A lookup join where we normally skip reading family 0.
+query T kvtrace
+SELECT * FROM t129145 INNER LOOKUP JOIN abc ON abc.a = t129145.a
+----
+Scan /Table/112/{1-2}
+Scan /Table/113/1/1/{1/1-2/2}, /Table/113/1/6/{1/1-2/2}
+
+# But with FOR UPDATE we must read all families.
+query T kvtrace
+SELECT * FROM t129145 INNER LOOKUP JOIN abc ON abc.a = t129145.a FOR UPDATE OF abc
+----
+Scan /Table/112/{1-2}
+Scan /Table/113/1/{1-2}, /Table/113/1/{6-7} lock Exclusive (Block, Replicated)
+
+statement ok
+RESET distsql
+
+statement ok
+RESET enable_durable_locking_for_serializable
+
+statement ok
+RESET optimizer_use_lock_op_for_serializable

pkg/sql/opt/locking.go

@@ -93,3 +93,19 @@ func (l Locking) IsLocking() bool {
 func (l Locking) IsNoOp() bool {
 	return l == Locking{}
 }
+
+// MustLockAllRequestedColumnFamilies returns true if the locking semantics
+// require actually acquiring locks on all requested column families. If this
+// returns false, then locking is best-effort, and it's ok to skip over some
+// requested column families.
+func (l Locking) MustLockAllRequestedColumnFamilies() bool {
+	// If durability = guaranteed, then we've promised that locks will be held
+	// until transaction commit. It's important in this case to actually acquire
+	// all the requested locks.
+	return (l.Strength != tree.ForNone && l.Durability == tree.LockDurabilityGuaranteed) ||
+		// If this read is using SKIP LOCKED or NOWAIT then we want to be sure to
+		// check all requested column families for locks. Otherwise we might assume
+		// there are no locks and subsequently hit an unexpected write-write
+		// conflict later in the same transaction.
+		l.WaitPolicy != tree.LockWaitBlock
+}

pkg/sql/opt_exec_factory.go

@@ -183,7 +183,12 @@ func generateScanSpans(
 	if params.InvertedConstraint != nil {
 		return sb.SpansFromInvertedSpans(ctx, params.InvertedConstraint, params.IndexConstraint, nil /* scratch */)
 	}
-	splitter := span.MakeSplitter(tabDesc, index, params.NeededCols)
+	var splitter span.Splitter
+	if params.Locking.MustLockAllRequestedColumnFamilies() {
+		splitter = span.MakeSplitterForSideEffect(tabDesc, index, params.NeededCols)
+	} else {
+		splitter = span.MakeSplitter(tabDesc, index, params.NeededCols)
+	}
 	return sb.SpansFromConstraint(params.IndexConstraint, splitter)
 }