txnlock: implement lock synthesis and sorting within txn (cockroachdb#164390)

txnlock: implement lock synthesis and sorting within txn

Author: trunk-io[bot]

pkg/BUILD.bazel

@@ -205,6 +205,7 @@ ALL_TESTS = [
     "//pkg/crosscluster/ldrrandgen:ldrrandgen_test",
     "//pkg/crosscluster/logical/ldrdecoder:ldrdecoder_test",
     "//pkg/crosscluster/logical/sqlwriter:sqlwriter_test",
+    "//pkg/crosscluster/logical/txnlock:txnlock_test",
     "//pkg/crosscluster/logical/txnwriter:txnwriter_test",
     "//pkg/crosscluster/logical:logical_disallowed_imports_test",
     "//pkg/crosscluster/logical:logical_test",
@@ -1527,6 +1528,8 @@ GO_TARGETS = [
     "//pkg/crosscluster/logical/ldrdecoder:ldrdecoder_test",
     "//pkg/crosscluster/logical/sqlwriter:sqlwriter",
     "//pkg/crosscluster/logical/sqlwriter:sqlwriter_test",
+    "//pkg/crosscluster/logical/txnlock:txnlock",
+    "//pkg/crosscluster/logical/txnlock:txnlock_test",
     "//pkg/crosscluster/logical/txnwriter:txnwriter",
     "//pkg/crosscluster/logical/txnwriter:txnwriter_test",
     "//pkg/crosscluster/logical:logical",

pkg/crosscluster/logical/txnlock/BUILD.bazel

@@ -0,0 +1,62 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "txnlock",
+    srcs = [
+        "column_set.go",
+        "doc.go",
+        "lock_synthesis.go",
+        "schema.go",
+        "sort.go",
+    ],
+    importpath = "github.com/cockroachdb/cockroach/pkg/crosscluster/logical/txnlock",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/crosscluster/logical/ldrdecoder",
+        "//pkg/crosscluster/logical/sqlwriter",
+        "//pkg/sql/catalog",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/lease",
+        "//pkg/sql/rowenc",
+        "//pkg/sql/sem/eval",
+        "//pkg/sql/sem/tree",
+        "//pkg/util/hlc",
+        "@com_github_cockroachdb_errors//:errors",
+    ],
+)
+
+go_test(
+    name = "txnlock_test",
+    srcs = [
+        "column_set_test.go",
+        "lock_synthesis_bench_test.go",
+        "lock_synthesis_test.go",
+        "main_test.go",
+        "schema_test.go",
+    ],
+    embed = [":txnlock"],
+    deps = [
+        "//pkg/base",
+        "//pkg/ccl",
+        "//pkg/ccl/changefeedccl/cdctest",
+        "//pkg/crosscluster/logical/ldrdecoder",
+        "//pkg/repstream/streampb",
+        "//pkg/security/securityassets",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/descs",
+        "//pkg/sql/catalog/desctestutils",
+        "//pkg/sql/catalog/lease",
+        "//pkg/sql/randgen",
+        "//pkg/sql/sem/eval",
+        "//pkg/sql/sem/tree",
+        "//pkg/sql/types",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/sqlutils",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/randutil",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/crosscluster/logical/txnlock/column_set.go

@@ -0,0 +1,100 @@
+// Copyright 2026 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package txnlock
+
+import (
+	"context"
+	"encoding/binary"
+	"hash/fnv"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/errors"
+)
+
+func uniqueIndexMixin(tableID descpb.ID, ucID descpb.ConstraintID) (uint64, error) {
+	h := fnv.New64a()
+	var buf [8]byte
+	binary.BigEndian.PutUint32(buf[:4], uint32(tableID))
+	binary.BigEndian.PutUint32(buf[4:], uint32(ucID))
+	if _, err := h.Write(buf[:]); err != nil {
+		return 0, errors.Wrap(err, "hashing unique index mixin")
+	}
+	return h.Sum64(), nil
+}
+
+// A columnSet is a collection of columns that are relevant for an individual
+// constraint.
+type columnSet struct {
+	columns []int32
+	// mixin is an integer that is combined with the hash. It is used to ensure
+	// different tables or unique constraints produce different hashes.
+	mixin uint64
+}
+
+// hash computes the hash that is used as the lock. hash(rowA) != hash(rowB)
+// implies !equal(rowA, rowB).
+func (c *columnSet) hash(ctx context.Context, row tree.Datums) (LockHash, error) {
+	h := fnv.New64a()
+	var prefixBytes [8]byte
+	binary.BigEndian.PutUint64(prefixBytes[:], c.mixin)
+	if _, err := h.Write(prefixBytes[:]); err != nil {
+		return 0, errors.Wrap(err, "hashing mixin for lock derivation")
+	}
+	for _, idx := range c.columns {
+		datum := row[idx]
+		ed := rowenc.EncDatum{Datum: datum}
+		encoded, err := ed.Fingerprint(
+			ctx,
+			datum.ResolvedType(),
+			&tree.DatumAlloc{},
+			nil, /* appendTo */
+			nil, /* acc */
+		)
+		if err != nil {
+			return 0, errors.Wrap(err, "hashing datum for lock derivation")
+		}
+		if _, err := h.Write(encoded); err != nil {
+			return 0, errors.Wrap(err, "hashing encoded datum for lock derivation")
+		}
+	}
+	return LockHash(h.Sum64()), nil
+}
+
+// null returns true if any of the columns are null.
+func (c *columnSet) null(row tree.Datums) bool {
+	if len(row) == 0 {
+		return true
+	}
+	for _, idx := range c.columns {
+		if row[idx] == tree.DNull {
+			return true
+		}
+	}
+	return false
+}
+
+// equal returns true if the columns are equal. equal(rowA, rowB) implies
+// hash(rowA) == hash(rowB).
+func (c *columnSet) equal(
+	ctx context.Context, evalCtx *eval.Context, rowA, rowB tree.Datums,
+) (bool, error) {
+	if c.null(rowA) || c.null(rowB) {
+		return false, nil
+	}
+	for _, idx := range c.columns {
+		cmp, err := rowA[idx].Compare(ctx, evalCtx, rowB[idx])
+		if err != nil {
+			return false, errors.Wrap(err, "comparing datums for lock derivation")
+		}
+		if cmp != 0 {
+			return false, nil
+		}
+	}
+	return true, nil
+}

pkg/crosscluster/logical/txnlock/column_set_test.go

@@ -0,0 +1,97 @@
+// Copyright 2026 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package txnlock
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/randgen"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestColumnSet(t *testing.T) {
+	rng := rand.New(rand.NewSource(0))
+
+	randomTypes := func(numCols int) []*types.T {
+		types := make([]*types.T, numCols)
+		for i := range types {
+			types[i] = randgen.RandColumnType(rng)
+		}
+		return types
+	}
+
+	randomRow := func(types []*types.T) tree.Datums {
+		row := make(tree.Datums, len(types))
+		for i, typ := range types {
+			row[i] = randgen.RandDatum(rng, typ, false /* nullOk */)
+		}
+		return row
+	}
+
+	test := func(t *testing.T, types []*types.T) {
+		ctx := context.Background()
+		evalCtx := eval.Context{}
+		colSet := columnSet{
+			columns: make([]int32, len(types)),
+		}
+		for i := range len(types) {
+			colSet.columns[i] = int32(i)
+		}
+		for range 1000 {
+			a := randomRow(types)
+			b := randomRow(types)
+
+			hashA1, err := colSet.hash(ctx, a)
+			require.NoError(t, err)
+			hashA2, err := colSet.hash(ctx, a)
+			require.NoError(t, err)
+			if hashA1 != hashA2 {
+				t.Errorf("hash(a) != hash(a)")
+			}
+
+			if !colSet.null(a) {
+				eq, err := colSet.equal(ctx, &evalCtx, a, a)
+				require.NoError(t, err)
+				if !eq {
+					t.Errorf("!equal(a, a) when !null(a)")
+				}
+			}
+
+			eqAB, err := colSet.equal(ctx, &evalCtx, a, b)
+			require.NoError(t, err)
+			if eqAB {
+				hashA, err := colSet.hash(ctx, a)
+				require.NoError(t, err)
+				hashB, err := colSet.hash(ctx, b)
+				require.NoError(t, err)
+				if hashA != hashB {
+					t.Errorf("equal(a, b) but hash(a) != hash(b)")
+				}
+			}
+
+			eqBA, err := colSet.equal(ctx, &evalCtx, b, a)
+			require.NoError(t, err)
+			if eqAB != eqBA {
+				t.Errorf("equal(a, b) != equal(b, a)")
+			}
+
+			if (colSet.null(a) || colSet.null(b)) && eqAB {
+				t.Errorf("null(a) or null(b) but equal(a, b) is true")
+			}
+		}
+	}
+
+	for range 100 {
+		types := randomTypes(rand.Intn(5) + 1)
+		t.Logf("types: %v", types)
+		test(t, types)
+	}
+}

pkg/crosscluster/logical/txnlock/doc.go

@@ -0,0 +1,73 @@
+// Copyright 2026 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+// Package txnlock synthesizes replication locks for transactional logical data
+// replication.
+//
+// # Replication Locks
+//
+// A replication lock is a (mode, kind, lock_id) triple:
+//
+//   - mode is one of {R, W}.
+//   - kind is one of {primary_key, foreign_key_constraint,
+//     unique_constraint}.
+//   - lock_id is a datum tuple whose values depend on the lock's kind.
+//
+// Two rows or transactions conflict when they hold locks with the same
+// lock_id and at least one of the locks is a write.
+//
+// # Primary Key Locks
+//
+// For a primary key lock, the lock_id is the primary key values of the
+// row inserted, updated, or deleted. The primary key values are derived
+// from the (prev_value, new_value) decoded from a rangefeed event.
+// Because both values always share the same primary key, each row
+// produces exactly one primary key lock.
+//
+// # Foreign Key Constraint Locks
+//
+// Foreign key constraint locks are derived from the (prev_value,
+// new_value) decoded from a rangefeed event.
+//
+// If table_parent has tables that reference it via foreign key
+// constraints, it generates a write lock: (W, foreign_key_constraint,
+// referenced_columns). The write lock can be omitted when (prev_value,
+// new_value) shows the referenced_columns did not change. If the
+// replication event does not modify the referenced columns, the
+// primary_key lock already orders it relative to any transaction that
+// does modify the referenced columns, which is sufficient to ensure
+// referential integrity.
+//
+// If table_child has a foreign key reference on table_parent, it
+// generates a read lock: (R, foreign_key_constraint, reference_key),
+// where reference_key is the referenced datums for the table_parent
+// row referenced by table_child. A replicated row may generate two
+// foreign key read locks for the same constraint if prev_value and
+// new_value have different non-null reference_keys. The read lock can
+// be omitted when prev_value and new_value have the same
+// reference_key.
+//
+// If any column in the reference_key is NULL, no lock is generated for
+// that constraint.
+//
+// # Unique Constraint Locks
+//
+// Unique constraint locks are derived from the (prev_value, new_value)
+// decoded from a rangefeed event. The lock_id is the tuple of column
+// values that form the unique constraint.
+//
+// A unique constraint lock is only generated when the constraint
+// column values differ between prev_value and new_value. If both the
+// old and new constraint values are non-null and different, the
+// replication event generates two unique constraint locks: one for the
+// old constraint values being released and one for the new constraint
+// values being acquired.
+//
+// # Replication Locks vs Transaction Locks
+//
+// An alternative to lock inference is to record the locks acquired by
+// the original transaction. This does not work because there are cases
+// where the SQL engine validates a constraint using a non-locking read.
+package txnlock