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/mutation.go

@@ -1134,3 +1134,26 @@ func unwrapProjectExprs(input memo.RelExpr) memo.RelExpr {
 	}
 	return input
 }
+
+func (b *Builder) buildLock(lock *memo.LockExpr) (execPlan, error) {
+	// Don't bother creating the lookup join if we don't need it.
+	locking, err := b.buildLocking(lock.Locking)
+	if err != nil {
+		return execPlan{}, err
+	}
+	if !locking.IsLocking() {
+		return b.buildRelational(lock.Input)
+	}
+
+	private := &memo.LookupJoinPrivate{
+		JoinType:              opt.SemiJoinOp,
+		Table:                 lock.Table,
+		Index:                 cat.PrimaryIndex,
+		KeyCols:               lock.KeyCols,
+		Cols:                  lock.Cols,
+		LookupColsAreTableKey: true,
+		Locking:               locking,
+	}
+	join := lock.Memo().MemoizeLookupJoin(lock.Input, nil /* on */, private)
+	return b.buildLookupJoin(join.(*memo.LookupJoinExpr))
+}

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)
 
@@ -3626,7 +3629,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()
 	}