opt: add lock operator

Add a new implementation of `SELECT FOR UPDATE` and `SELECT FOR SHARE`
statements. Instead of locking during the initial row fetch, this new
implementation constructs a `Lock` operator on the top of the query plan
which performs the locking phase using a locking semi-join lookup.

During optbuilder we build plans with both `Lock` operators and
initial-row-fetch locking. During execbuilder we decide which
implementation to use based on the isolation level and whether
`optimizer_use_lock_op_for_serializable` is set. If the new
implementation is chosen, `Lock` operators become locking
semi-LookupJoins.

In some cases these new plans will have superfluous lookup joins. A
future PR will optimize away some of these superfluous lookup joins.

Fixes: cockroachdb#57031, cockroachdb#75457

Epic: CRDB-25322

Release note (sql change): Add a new session variable,
`optimizer_use_lock_op_for_serializable`, which when set enables a new
implementation of `SELECT FOR UPDATE`. This new implementation of
`SELECT FOR UPDATE` acquires row locks *after* any joins and filtering,
and always acquires row locks on the primary index of the table being
locked. This more closely matches `SELECT FOR UPDATE` behavior in
PostgreSQL, but at the cost of more round trips from gateway node to
replica leaseholder.

Under read committed isolation (and other isolation levels weaker than
serializable) we will always use this new implementation of `SELECT FOR
UPDATE` regardless of the value of
`optimizer_use_lock_op_for_serializable` to ensure correctness.

pkg/sql/exec_util.go

@@ -3667,6 +3667,10 @@ func (m *sessionDataMutator) SetUnsafeSettingInterlockKey(val string) {
 	m.data.UnsafeSettingInterlockKey = val
 }
 
+func (m *sessionDataMutator) SetOptimizerUseLockOpForSerializable(val bool) {
+	m.data.OptimizerUseLockOpForSerializable = val
+}
+
 // Utility functions related to scrubbing sensitive information on SQL Stats.
 
 // quantizeCounts ensures that the Count field in the

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -5568,6 +5568,7 @@ optimizer_use_improved_disjunction_stats                   on
 optimizer_use_improved_join_elimination                    on
 optimizer_use_improved_split_disjunction_for_joins         on
 optimizer_use_limit_ordering_for_streaming_group_by        on
+optimizer_use_lock_op_for_serializable                     off
 optimizer_use_multicol_stats                               on
 optimizer_use_not_visible_indexes                          off
 override_multi_region_zone_config                          off

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -2886,6 +2886,7 @@ optimizer_use_improved_disjunction_stats                   on                  N
 optimizer_use_improved_join_elimination                    on                  NULL      NULL        NULL        string
 optimizer_use_improved_split_disjunction_for_joins         on                  NULL      NULL        NULL        string
 optimizer_use_limit_ordering_for_streaming_group_by        on                  NULL      NULL        NULL        string
+optimizer_use_lock_op_for_serializable                     off                 NULL      NULL        NULL        string
 optimizer_use_multicol_stats                               on                  NULL      NULL        NULL        string
 optimizer_use_not_visible_indexes                          off                 NULL      NULL        NULL        string
 override_multi_region_zone_config                          off                 NULL      NULL        NULL        string
@@ -3048,6 +3049,7 @@ optimizer_use_improved_disjunction_stats                   on                  N
 optimizer_use_improved_join_elimination                    on                  NULL  user     NULL      on                  on
 optimizer_use_improved_split_disjunction_for_joins         on                  NULL  user     NULL      on                  on
 optimizer_use_limit_ordering_for_streaming_group_by        on                  NULL  user     NULL      on                  on
+optimizer_use_lock_op_for_serializable                     off                 NULL  user     NULL      off                 off
 optimizer_use_multicol_stats                               on                  NULL  user     NULL      on                  on
 optimizer_use_not_visible_indexes                          off                 NULL  user     NULL      off                 off
 override_multi_region_zone_config                          off                 NULL  user     NULL      off                 off
@@ -3209,6 +3211,7 @@ optimizer_use_improved_disjunction_stats                   NULL    NULL     NULL
 optimizer_use_improved_join_elimination                    NULL    NULL     NULL     NULL        NULL
 optimizer_use_improved_split_disjunction_for_joins         NULL    NULL     NULL     NULL        NULL
 optimizer_use_limit_ordering_for_streaming_group_by        NULL    NULL     NULL     NULL        NULL
+optimizer_use_lock_op_for_serializable                     NULL    NULL     NULL     NULL        NULL
 optimizer_use_multicol_stats                               NULL    NULL     NULL     NULL        NULL
 optimizer_use_not_visible_indexes                          NULL    NULL     NULL     NULL        NULL
 override_multi_region_zone_config                          NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -124,6 +124,7 @@ optimizer_use_improved_disjunction_stats                   on
 optimizer_use_improved_join_elimination                    on
 optimizer_use_improved_split_disjunction_for_joins         on
 optimizer_use_limit_ordering_for_streaming_group_by        on
+optimizer_use_lock_op_for_serializable                     off
 optimizer_use_multicol_stats                               on
 optimizer_use_not_visible_indexes                          off
 override_multi_region_zone_config                          off

pkg/sql/opt/exec/execbuilder/builder.go

@@ -14,6 +14,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
@@ -96,6 +97,12 @@ type Builder struct {
 
 	allowInsertFastPath bool
 
+	// onlyForcedLocking is passed through to factory methods for operators that
+	// use locking. When set to true, it silences all locking that is not
+	// forced. This is used to ignore initial-row-fetch locking in favor of the
+	// Lock operator when running under read committed isolation.
+	onlyForcedLocking bool
+
 	// forceForUpdateLocking is conditionally passed through to factory methods
 	// for scan operators that serve as the input for mutation operators. When
 	// set to true, it ensures that a FOR UPDATE row-level locking mode is used
@@ -314,6 +321,9 @@ func (b *Builder) build(e opt.Expr) (_ execPlan, err error) {
 	//    auto-commit).
 	b.allowInsertFastPath = b.allowInsertFastPath && canAutoCommit
 
+	b.onlyForcedLocking = b.evalCtx.TxnIsoLevel != isolation.Serializable ||
+		b.evalCtx.SessionData().OptimizerUseLockOpForSerializable
+
 	return b.buildRelational(rel)
 }
 

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -14,6 +14,8 @@ import (
 	"bytes"
 	"fmt"
 
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/lexbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
@@ -156,6 +158,11 @@ func (b *Builder) tryBuildFastPathInsert(ins *memo.InsertExpr) (_ execPlan, ok b
 		return execPlan{}, false, nil
 	}
 
+	if b.onlyForcedLocking {
+		b.onlyForcedLocking = false
+		defer func() { b.onlyForcedLocking = true }()
+	}
+
 	md := b.mem.Metadata()
 	tab := md.Table(ins.Table)
 
@@ -763,6 +770,10 @@ func (b *Builder) checkContainsLocking(mainContainsLocking bool) {
 func (b *Builder) buildUniqueChecks(checks memo.UniqueChecksExpr) error {
 	defer b.checkContainsLocking(b.ContainsNonDefaultKeyLocking)
 	b.ContainsNonDefaultKeyLocking = false
+	if b.onlyForcedLocking {
+		b.onlyForcedLocking = false
+		defer func() { b.onlyForcedLocking = true }()
+	}
 	md := b.mem.Metadata()
 	for i := range checks {
 		c := &checks[i]
@@ -795,6 +806,10 @@ func (b *Builder) buildUniqueChecks(checks memo.UniqueChecksExpr) error {
 func (b *Builder) buildFKChecks(checks memo.FKChecksExpr) error {
 	defer b.checkContainsLocking(b.ContainsNonDefaultKeyLocking)
 	b.ContainsNonDefaultKeyLocking = false
+	if b.onlyForcedLocking {
+		b.onlyForcedLocking = false
+		defer func() { b.onlyForcedLocking = true }()
+	}
 	md := b.mem.Metadata()
 	for i := range checks {
 		c := &checks[i]
@@ -1134,3 +1149,35 @@ func unwrapProjectExprs(input memo.RelExpr) memo.RelExpr {
 	}
 	return input
 }
+
+func (b *Builder) buildLock(lock *memo.LockExpr) (execPlan, error) {
+	locking := lock.Locking
+	locking.Forced = true
+
+	// Don't bother creating the lookup join if we don't need it.
+	if !locking.IsLocking() ||
+		!b.evalCtx.Settings.Version.IsActive(b.ctx, clusterversion.V23_2) ||
+		(b.evalCtx.TxnIsoLevel == isolation.Serializable &&
+			!b.evalCtx.SessionData().OptimizerUseLockOpForSerializable) ||
+		((locking.Strength == tree.ForShare || locking.Strength == tree.ForKeyShare) &&
+			b.evalCtx.TxnIsoLevel == isolation.Serializable &&
+			!b.evalCtx.SessionData().SharedLockingForSerializable) {
+		return b.buildRelational(lock.Input)
+	}
+
+	join := &memo.LookupJoinExpr{
+		Input: lock.Input,
+		LookupJoinPrivate: memo.LookupJoinPrivate{
+			JoinType:              opt.SemiJoinOp,
+			Table:                 lock.Table,
+			Index:                 cat.PrimaryIndex,
+			KeyCols:               lock.KeyCols,
+			Cols:                  lock.Cols,
+			LookupColsAreTableKey: true,
+			Locking:               locking,
+		},
+	}
+	memo.CopyLockGroupIntoLookupJoin(lock, join)
+
+	return b.buildLookupJoin(join)
+}

pkg/sql/opt/exec/execbuilder/relational.go

@@ -292,6 +292,9 @@ func (b *Builder) buildRelational(e memo.RelExpr) (execPlan, error) {
 	case *memo.DeleteExpr:
 		ep, err = b.buildDelete(t)
 
+	case *memo.LockExpr:
+		ep, err = b.buildLock(t)
+
 	case *memo.CreateTableExpr:
 		ep, err = b.buildCreateTable(t)
 
@@ -2906,6 +2909,7 @@ func (b *Builder) buildZigzagJoin(join *memo.ZigzagJoinExpr) (execPlan, error) {
 func (b *Builder) buildLocking(locking opt.Locking) (opt.Locking, error) {
 	if b.forceForUpdateLocking {
 		locking = locking.Max(forUpdateLocking)
+		locking.Forced = true
 	}
 	if locking.IsLocking() {
 		// Raise error if row-level locking is part of a read-only transaction.
@@ -2914,6 +2918,10 @@ func (b *Builder) buildLocking(locking opt.Locking) (opt.Locking, error) {
 				"cannot execute SELECT %s in a read-only transaction", locking.Strength.String(),
 			)
 		}
+		// Silence initial-row-fetch locking when using the Lock operator.
+		if b.onlyForcedLocking && !locking.Forced {
+			return opt.Locking{}, nil
+		}
 		if locking.Form == tree.LockPredicate {
 			return opt.Locking{}, unimplemented.NewWithIssuef(
 				110873, "explicit unique checks are not yet supported under read committed isolation",
@@ -3639,7 +3647,8 @@ func (b *Builder) statementTag(expr memo.RelExpr) string {
 	switch expr.Op() {
 	case opt.OpaqueRelOp, opt.OpaqueMutationOp, opt.OpaqueDDLOp:
 		return expr.Private().(*memo.OpaqueRelPrivate).Metadata.String()
-
+	case opt.LockOp:
+		return "SELECT " + expr.Private().(*memo.LockPrivate).Locking.Strength.String()
 	default:
 		return expr.Op().SyntaxTag()
 	}