Merge #154885 #154988

154885: concurrency: make the lock table limit configurable r=stevendanna a=arulajmani

Previously, this was hardcoded to 10K. We now introduce a new cluster setting to change this, if needed. We go through all ranges to update the limit on all exisiting lock table's in response to the cluster setting change.

Release note: None

Epic: none

154988: bench/rttanalysis: Revert "bench/rttanalysis: shard TestBenchmarkExpectation to avoid timeouts" r=fqazi a=fqazi

This reverts commit 9fecc53.

Fixes: #154317
Fixes: #154084
Fixes: #154559

Release note: None

Co-authored-by: Arul Ajmani <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 3 people

pkg/bench/rttanalysis/BUILD.bazel

@@ -13,8 +13,6 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/base",
-        "//pkg/jobs",
-        "//pkg/jobs/jobspb",
         "//pkg/kv/kvclient/kvcoord",
         "//pkg/sql",
         "//pkg/sql/parser",
@@ -58,9 +56,9 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":rttanalysis"],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 4,
     deps = [
         "//pkg/base",
+        "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
@@ -71,7 +69,6 @@ go_test(
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/testcluster",
-        "//pkg/util/envutil",
         "//pkg/util/protoutil",
         "//pkg/util/randutil",
     ],

pkg/bench/rttanalysis/registry.go

@@ -9,8 +9,6 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/cockroachdb/cockroach/pkg/jobs"
-	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
@@ -53,66 +51,15 @@ func (r *Registry) Run(b *testing.B) {
 // benchmarks can be filtered by passing the usual test filters underneath
 // this test's name.
 //
-// It takes a long time and thus is skipped under duress and short.
+// It takes a long time and thus is skipped under stress, race
+// and short.
 func (r *Registry) RunExpectations(t *testing.T) {
-	r.RunExpectationsSharded(t, 1, 1)
-}
-
-// RunExpectationsSharded runs all the benchmarks for one iteration
-// and validates that the number of RPCs meets the expectation. If run
-// with the --rewrite flag, it will rewrite the run benchmarks. The
-// benchmarks can be filtered by passing the usual test filters underneath
-// this test's name.
-//
-// It takes a long time and thus is skipped under duress and short.
-//
-// When shard and totalShards are provided (> 1), only a subset of benchmarks
-// assigned to the specific shard will be run, enabling parallel execution.
-// Test groups are distributed across shards using round-robin assignment.
-func (r *Registry) RunExpectationsSharded(t *testing.T, shard, totalShards int) {
-	defer jobs.TestingSetIDsToIgnore(map[jobspb.JobID]struct{}{3001: {}, 3002: {}})()
-	skip.UnderDuress(t)
+	skip.UnderStress(t)
+	skip.UnderRace(t)
 	skip.UnderShort(t)
+	skip.UnderDeadlock(t)
 
-	// If totalShards is 1, run all tests; otherwise shard them
-	var registryToUse *Registry
-	if totalShards <= 1 {
-		// Run all test groups
-		registryToUse = r
-	} else {
-		// Create a registry with only the test groups assigned to this shard
-		shardRegistry := &Registry{
-			numNodes: r.numNodes,
-			cc:       r.cc,
-			r:        make(map[string][]RoundTripBenchTestCase),
-		}
-
-		// Distribute test groups across shards using round-robin assignment
-		// First, get all group names and sort them for consistent ordering
-		groupNames := make([]string, 0, len(r.r))
-		for groupName := range r.r {
-			groupNames = append(groupNames, groupName)
-		}
-		// Sort for deterministic assignment across runs
-		for i := 0; i < len(groupNames); i++ {
-			for j := i + 1; j < len(groupNames); j++ {
-				if groupNames[i] > groupNames[j] {
-					groupNames[i], groupNames[j] = groupNames[j], groupNames[i]
-				}
-			}
-		}
-
-		// Assign groups to shards using round-robin
-		for i, groupName := range groupNames {
-			assignedShard := (i % totalShards) + 1
-			if assignedShard == shard {
-				shardRegistry.r[groupName] = r.r[groupName]
-			}
-		}
-		registryToUse = shardRegistry
-	}
-
-	runBenchmarkExpectationTests(t, registryToUse)
+	runBenchmarkExpectationTests(t, r)
 }
 
 // Register registers a set of test cases to a given benchmark name. It is

pkg/bench/rttanalysis/validate_benchmark_data_test.go

@@ -6,44 +6,13 @@
 package rttanalysis
 
 import (
-	"strconv"
 	"testing"
 
-	"github.com/cockroachdb/cockroach/pkg/util/envutil"
+	"github.com/cockroachdb/cockroach/pkg/jobs"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 )
 
-// NOTE: If you change the number of shards, you must also update the
-// shard_count in BUILD.bazel to match.
-const shardCount = 4
-
-// Validate that shardCount matches TEST_TOTAL_SHARDS environment variable at init time
-var _ = func() int {
-	totalShardsStr, found := envutil.ExternalEnvString("TEST_TOTAL_SHARDS", 1)
-	if totalShardsStr == "" || !found {
-		return 0
-	}
-	totalShards, err := strconv.Atoi(totalShardsStr)
-	if err != nil {
-		return 0
-	}
-	if totalShards != shardCount {
-		panic("shardCount mismatch: update shard_count in pkg/bench/rttanalysis/BUILD.bazel to match shardCount constant")
-	}
-	return 0
-}()
-
-func TestBenchmarkExpectationShard1(t *testing.T) {
-	reg.RunExpectationsSharded(t, 1, shardCount)
-}
-
-func TestBenchmarkExpectationShard2(t *testing.T) {
-	reg.RunExpectationsSharded(t, 2, shardCount)
-}
-
-func TestBenchmarkExpectationShard3(t *testing.T) {
-	reg.RunExpectationsSharded(t, 3, shardCount)
-}
-
-func TestBenchmarkExpectationShard4(t *testing.T) {
-	reg.RunExpectationsSharded(t, 4, shardCount)
+func TestBenchmarkExpectation(t *testing.T) {
+	defer jobs.TestingSetIDsToIgnore(map[jobspb.JobID]struct{}{3001: {}, 3002: {}})()
+	reg.RunExpectations(t)
 }

pkg/kv/kvserver/concurrency/concurrency_control.go

@@ -271,6 +271,11 @@ type LockManager interface {
 	// ExportUnreplicatedLocks runs exporter on each held, unreplicated lock
 	// in the given span until the exporter returns false.
 	ExportUnreplicatedLocks(span roachpb.Span, exporter func(*roachpb.LockAcquisition) bool)
+
+	// SetMaxLockTableSize updates the lock table's maximum size limit. It may
+	// be used to dynamically adjust the lock table's size after it has been
+	// initialized.
+	SetMaxLockTableSize(maxLocks int64)
 }
 
 // TransactionManager is concerned with tracking transactions that have their
@@ -354,10 +359,6 @@ type TestingAccessor interface {
 
 	// TestingTxnWaitQueue returns the concurrency manager's txnWaitQueue.
 	TestingTxnWaitQueue() *txnwait.Queue
-
-	// TestingSetMaxLocks updates the locktable's lock limit. This can be used to
-	// force the locktable to exceed its limit and clear locks.
-	TestingSetMaxLocks(n int64)
 }
 
 ///////////////////////////////////
@@ -789,9 +790,10 @@ type lockTable interface {
 	// String returns a debug string representing the state of the lockTable.
 	String() string
 
-	// TestingSetMaxLocks updates the locktable's lock limit. This can be used to
-	// force the locktable to exceed its limit and clear locks.
-	TestingSetMaxLocks(maxLocks int64)
+	// SetMaxLockTableSize updates the lock table's maximum size limit. It may
+	// be used to dynamically adjust the lock table's size after it has been
+	// initialized.
+	SetMaxLockTableSize(maxLocks int64)
 }
 
 // lockTableGuard is a handle to a request as it waits on conflicting locks in a

pkg/kv/kvserver/concurrency/concurrency_manager.go

@@ -103,6 +103,17 @@ var DiscoveredLocksThresholdToConsultTxnStatusCache = settings.RegisterIntSettin
 	settings.NonNegativeInt,
 )
 
+// DefaultLockTableSize controls the default upper bound on the number of locks
+// in a lock table.
+var DefaultLockTableSize = settings.RegisterIntSetting(
+	settings.SystemOnly,
+	"kv.lock_table.default_size",
+	"the default upper bound on the number of locks in a lock table. This setting "+
+		"controls the maximum number of locks that can be held in memory by the lock table "+
+		"before it starts evicting locks to manage memory pressure.",
+	10000,
+)
+
 // BatchPushedLockResolution controls whether the lock table should allow
 // non-locking readers to defer and batch the resolution of conflicting locks
 // whose holder is known to be pending and have been pushed above the reader's
@@ -199,7 +210,7 @@ type Config struct {
 
 func (c *Config) initDefaults() {
 	if c.MaxLockTableSize == 0 {
-		c.MaxLockTableSize = defaultLockTableSize
+		c.MaxLockTableSize = DefaultLockTableSize.Get(&c.Settings.SV)
 	}
 }
 
@@ -759,9 +770,9 @@ func (m *managerImpl) TestingTxnWaitQueue() *txnwait.Queue {
 	return m.twq.(*txnwait.Queue)
 }
 
-// TestingSetMaxLocks implements the TestingAccessor interface.
-func (m *managerImpl) TestingSetMaxLocks(maxLocks int64) {
-	m.lt.TestingSetMaxLocks(maxLocks)
+// SetMaxLockTableSize implements the LockManager interface.
+func (m *managerImpl) SetMaxLockTableSize(maxLocks int64) {
+	m.lt.SetMaxLockTableSize(maxLocks)
 }
 
 func (r *Request) isSingle(m kvpb.Method) bool {

pkg/kv/kvserver/concurrency/concurrency_manager_test.go

@@ -641,7 +641,7 @@ func TestConcurrencyManagerBasic(t *testing.T) {
 			case "debug-set-max-locks":
 				var n int
 				d.ScanArgs(t, "n", &n)
-				m.TestingSetMaxLocks(int64(n))
+				m.SetMaxLockTableSize(int64(n))
 				return ""
 
 			case "reset":