diff --git a/pkg/kv/kvserver/concurrency/lock/locking.go b/pkg/kv/kvserver/concurrency/lock/locking.go index 111393ebad96..19b00cf2edc7 100644 --- a/pkg/kv/kvserver/concurrency/lock/locking.go +++ b/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. diff --git a/pkg/kv/kvserver/lockspanset/BUILD.bazel b/pkg/kv/kvserver/lockspanset/BUILD.bazel new file mode 100644 index 000000000000..99b9ab653aa4 --- /dev/null +++ b/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", + ], +) diff --git a/pkg/kv/kvserver/lockspanset/lockspanset.go b/pkg/kv/kvserver/lockspanset/lockspanset.go new file mode 100644 index 000000000000..a35ef9c012cd --- /dev/null +++ b/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 +} diff --git a/pkg/kv/kvserver/lockspanset/lockspanset_test.go b/pkg/kv/kvserver/lockspanset/lockspanset_test.go new file mode 100644 index 000000000000..44274abe8177 --- /dev/null +++ b/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) +}