lockspanset: introduce LockSpanSets

Today we use `spanset.SpanSet` to declare both lock and latch spans.
However, `SpanSets` do not generalize to more lock strengths. This patch
introduces `LockSpanSets` in preperation for shared locks. We don't make
use of them in the `concurrency` package yet; that will happen in an
upcoming commit.

The main motivations for the new `LockSpanSets` structure are:
- A desire to store/declare lock spans using lock strengths, not span
access.
- There isn't a need to store MVCC spans -- lock strengths do not have
associated timestamps.
- There is no need to store global/local lock keys separately, like we
do for latches.

Furthermore, we only port over methods on `SpanSet` that lock spans made
use of in this patch.

Informs: cockroachdb#102008

Release note: None

pkg/kv/kvserver/concurrency/lock/locking.go

@@ -53,12 +53,20 @@ var ExclusiveLocksBlockNonLockingReads = settings.RegisterBoolSetting(
 // MaxDurability is the maximum value in the Durability enum.
 const MaxDurability = Unreplicated
 
+// MaxStrength is the maximum value of the Strength enum.
+const MaxStrength = Intent
+
 func init() {
 	for v := range Durability_name {
 		if d := Durability(v); d > MaxDurability {
 			panic(fmt.Sprintf("Durability (%s) with value larger than MaxDurability", d))
 		}
 	}
+	for v := range Strength_name {
+		if st := Strength(v); st > MaxStrength {
+			panic(fmt.Sprintf("Strength (%s) with value larger than MaxDurability", st))
+		}
+	}
 }
 
 // Conflicts returns whether the supplied lock modes conflict with each other.

pkg/kv/kvserver/lockspanset/BUILD.bazel

@@ -0,0 +1,25 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "lockspanset",
+    srcs = ["lockspanset.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/lockspanset",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/kv/kvserver/concurrency/lock",
+        "//pkg/roachpb",
+    ],
+)
+
+go_test(
+    name = "lockspanset_test",
+    srcs = ["lockspanset_test.go"],
+    embed = [":lockspanset"],
+    deps = [
+        "//pkg/kv/kvserver/concurrency/lock",
+        "//pkg/roachpb",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/kv/kvserver/lockspanset/lockspanset.go

@@ -0,0 +1,97 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package lockspanset
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+)
+
+const NumLockStrength = lock.MaxStrength + 1
+
+type LockSpanSet struct {
+	spans [NumLockStrength][]roachpb.Span
+}
+
+var lockSpanSetPool = sync.Pool{
+	New: func() interface{} { return new(LockSpanSet) },
+}
+
+// New creates a new empty LockSpanSet.
+func New() *LockSpanSet {
+	return lockSpanSetPool.Get().(*LockSpanSet)
+}
+
+// GetSpans returns a slice of spans with the given access.
+func (l *LockSpanSet) GetSpans(access lock.Strength) []roachpb.Span {
+	return l.spans[access]
+}
+
+// Add adds the supplied span to the lock span set to be accessed with the given
+// lock strength.
+func (l *LockSpanSet) Add(access lock.Strength, span roachpb.Span) {
+	l.spans[access] = append(l.spans[access], span)
+}
+
+// SortAndDeDup sorts the spans in the SpanSet and removes any duplicates.
+func (l *LockSpanSet) SortAndDeDup() {
+	for st := lock.Strength(0); st < NumLockStrength; st++ {
+		l.spans[st], _ /* distinct */ = roachpb.MergeSpans(&l.spans[st])
+	}
+}
+
+// Release releases the LockSpanSet and its underlying slices. The receiver
+// should not be used after being released.
+func (l *LockSpanSet) Release() {
+	for st := lock.Strength(0); st < NumLockStrength; st++ {
+		// Recycle slice if capacity below threshold.
+		const maxRecycleCap = 8
+		var recycle []roachpb.Span
+		if sl := l.spans[st]; cap(sl) <= maxRecycleCap {
+			for i := range sl {
+				sl[i] = roachpb.Span{}
+			}
+			recycle = sl[:0]
+		}
+		l.spans[st] = recycle
+	}
+	lockSpanSetPool.Put(l)
+}
+
+// Empty returns whether the set contains any spans across all lock strengths.
+func (l *LockSpanSet) Empty() bool {
+	return l.Len() == 0
+}
+
+// String prints a string representation of the LockSpanSet.
+func (l *LockSpanSet) String() string {
+	var buf strings.Builder
+	for st := lock.Strength(0); st < NumLockStrength; st++ {
+		for _, span := range l.GetSpans(st) {
+			fmt.Fprintf(&buf, "%s: %s\n",
+				st, span)
+		}
+	}
+	return buf.String()
+}
+
+// Len returns the total number of spans tracked across all strengths.
+func (l *LockSpanSet) Len() int {
+	var count int
+	for st := lock.Strength(0); st < NumLockStrength; st++ {
+		count += len(l.GetSpans(st))
+	}
+	return count
+}

pkg/kv/kvserver/lockspanset/lockspanset_test.go

@@ -0,0 +1,113 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package lockspanset
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+// Test that spans are properly classified according to the lock strength they
+// are added with.
+func TestGetSpansStrength(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	var lss LockSpanSet
+	spA := roachpb.Span{Key: roachpb.Key("a")}
+	spBC := roachpb.Span{Key: roachpb.Key("b"), EndKey: roachpb.Key("c")}
+	spD := roachpb.Span{Key: roachpb.Key("d")}
+	spE := roachpb.Span{Key: roachpb.Key("e")}
+	spF := roachpb.Span{Key: roachpb.Key("f")}
+	spGH := roachpb.Span{Key: roachpb.Key("g"), EndKey: roachpb.Key("h")}
+
+	lss.Add(lock.None, spA)
+	lss.Add(lock.Shared, spBC)
+	lss.Add(lock.Update, spD)
+
+	lss.Add(lock.Exclusive, spE)
+	lss.Add(lock.Exclusive, spGH)
+
+	lss.Add(lock.Intent, spF)
+
+	spans := lss.GetSpans(lock.None)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spA}))
+
+	spans = lss.GetSpans(lock.Shared)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spBC}))
+
+	spans = lss.GetSpans(lock.Update)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spD}))
+
+	spans = lss.GetSpans(lock.Exclusive)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spE, spGH}))
+
+	spans = lss.GetSpans(lock.Intent)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spF}))
+}
+
+// TestLockSpanSetSortAndDeDup ensures that spans in a lock span set are sorted
+// and de-duplicated correctly. Spans should be sorted and de-duplicated within
+// a particular lock strength but not amongst different lock strengths.
+func TestLockSpanSetSortAndDeDup(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	makeSpan := func(start, end string) roachpb.Span {
+		var endKey roachpb.Key
+		if end != "" {
+			endKey = roachpb.Key(end)
+		}
+		return roachpb.Span{Key: roachpb.Key(start), EndKey: endKey}
+	}
+
+	spA := makeSpan("a", "")
+	spB := makeSpan("b", "")
+	spCF := makeSpan("c", "f")
+	spEH := makeSpan("e", "h")
+	spGJ := makeSpan("g", "j")
+	spIL := makeSpan("i", "l")
+	spXZ := makeSpan("x", "z")
+
+	var lss LockSpanSet
+	lss.Add(lock.None, spA)
+	lss.Add(lock.None, spA)  // duplicate
+	lss.Add(lock.None, spA)  // duplicate
+	lss.Add(lock.None, spA)  // duplicate
+	lss.Add(lock.None, spCF) // overlapping
+	lss.Add(lock.None, spEH) // overlapping
+	lss.Add(lock.None, spGJ) // overlapping
+	lss.Add(lock.None, spB)  // out of order
+
+	// Shared.
+	lss.Add(lock.Shared, spXZ) // out of order
+	lss.Add(lock.Shared, spA)  // should not be considered a duplicate
+	lss.Add(lock.Shared, spIL) // should not overlap
+
+	lss.SortAndDeDup()
+
+	spans := lss.GetSpans(lock.None)
+	require.Len(t, spans, 3)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spA, spB, makeSpan("c", "j")}))
+
+	spans = lss.GetSpans(lock.Shared)
+	require.True(t, reflect.DeepEqual(spans, []roachpb.Span{spA, spIL, spXZ}))
+
+	require.Len(t, lss.GetSpans(lock.Update), 0)
+	require.Len(t, lss.GetSpans(lock.Exclusive), 0)
+	require.Len(t, lss.GetSpans(lock.Intent), 0)
+}