Merge #144604

144604: kv: update uses of btree r=RaduBerinde a=RaduBerinde

Use the newer btree implementation with generics.

Informs: #144504
Release note: None

Co-authored-by: Radu Berinde <[email protected]>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/BUILD.bazel

@@ -98,7 +98,7 @@ go_library(
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_gogo_protobuf//proto",
-        "@com_github_raduberinde_btree//:btree",
+        "@com_github_google_btree//:btree",
         "@io_opentelemetry_go_otel//attribute",
     ],
 )

pkg/kv/kvclient/kvcoord/txn_interceptor_pipeliner.go

@@ -10,7 +10,6 @@ import (
 	"math"
 	"sort"
 
-	gbtree "github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -22,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
+	gbtree "github.com/google/btree"
 )
 
 // The degree of the inFlightWrites btree.
@@ -1048,13 +1048,12 @@ func makeInFlightWrite(key roachpb.Key, seq enginepb.TxnSeq, str lock.Strength)
 	}}
 }
 
-// Less implements the gbtree.Item interface.
+// less is the ordering function used by the B tree.
 //
 // inFlightWrites are ordered by Key, then by Sequence, then by Strength. Two
 // inFlightWrites with the same Key but different Sequences and/or Strengths are
 // not considered equal and are maintained separately in the inFlightWritesSet.
-func (a *inFlightWrite) Less(bItem gbtree.Item) bool {
-	b := bItem.(*inFlightWrite)
+func (a *inFlightWrite) less(b *inFlightWrite) bool {
 	kCmp := a.Key.Compare(b.Key)
 	if kCmp != 0 {
 		// Different Keys.
@@ -1077,7 +1076,7 @@ func (a *inFlightWrite) Less(bItem gbtree.Item) bool {
 // writes, O(log n) removal of existing in-flight writes, and O(m + log n)
 // retrieval over m in-flight writes that overlap with a given key.
 type inFlightWriteSet struct {
-	t     *gbtree.BTree
+	t     *gbtree.BTreeG[*inFlightWrite]
 	bytes int64
 
 	// Avoids allocs.
@@ -1090,15 +1089,15 @@ type inFlightWriteSet struct {
 func (s *inFlightWriteSet) insert(key roachpb.Key, seq enginepb.TxnSeq, str lock.Strength) {
 	if s.t == nil {
 		// Lazily initialize btree.
-		s.t = gbtree.New(txnPipelinerBtreeDegree)
+		s.t = gbtree.NewG[*inFlightWrite](txnPipelinerBtreeDegree, (*inFlightWrite).less)
 	}
 
 	w := s.alloc.alloc(key, seq, str)
-	delItem := s.t.ReplaceOrInsert(w)
+	delItem, _ := s.t.ReplaceOrInsert(w)
 	if delItem != nil {
 		// An in-flight write with the same key and sequence already existed in the
 		// set. We replaced it with an identical in-flight write.
-		*delItem.(*inFlightWrite) = inFlightWrite{} // for GC
+		*delItem = inFlightWrite{} // for GC
 	} else {
 		s.bytes += keySize(key)
 	}
@@ -1114,12 +1113,12 @@ func (s *inFlightWriteSet) remove(key roachpb.Key, seq enginepb.TxnSeq, str lock
 
 	// Delete the write from the in-flight writes set.
 	s.tmp1 = makeInFlightWrite(key, seq, str)
-	delItem := s.t.Delete(&s.tmp1)
+	delItem, _ := s.t.Delete(&s.tmp1)
 	if delItem == nil {
 		// The write was already proven or the txn epoch was incremented.
 		return
 	}
-	*delItem.(*inFlightWrite) = inFlightWrite{} // for GC
+	*delItem = inFlightWrite{} // for GC
 	s.bytes -= keySize(key)
 
 	// Assert that the byte accounting is believable.
@@ -1136,8 +1135,8 @@ func (s *inFlightWriteSet) ascend(f func(w *inFlightWrite)) {
 		// Set is empty.
 		return
 	}
-	s.t.Ascend(func(i gbtree.Item) bool {
-		f(i.(*inFlightWrite))
+	s.t.Ascend(func(i *inFlightWrite) bool {
+		f(i)
 		return true
 	})
 }
@@ -1157,8 +1156,8 @@ func (s *inFlightWriteSet) ascendRange(start, end roachpb.Key, f func(w *inFligh
 		// Range lookup.
 		s.tmp2 = makeInFlightWrite(end, 0, 0)
 	}
-	s.t.AscendRange(&s.tmp1, &s.tmp2, func(i gbtree.Item) bool {
-		f(i.(*inFlightWrite))
+	s.t.AscendRange(&s.tmp1, &s.tmp2, func(i *inFlightWrite) bool {
+		f(i)
 		return true
 	})
 }

pkg/kv/kvclient/kvcoord/txn_interceptor_pipeliner_test.go

@@ -208,7 +208,7 @@ func TestTxnPipelinerTrackInFlightWrites(t *testing.T) {
 	require.NotNil(t, br)
 	require.Equal(t, 1, tp.ifWrites.len())
 
-	w := tp.ifWrites.t.Min().(*inFlightWrite)
+	w, _ := tp.ifWrites.t.Min()
 	require.Equal(t, putArgs.Key, w.Key)
 	require.Equal(t, putArgs.Sequence, w.Sequence)
 
@@ -266,10 +266,10 @@ func TestTxnPipelinerTrackInFlightWrites(t *testing.T) {
 	require.Nil(t, pErr)
 	require.Equal(t, 4, tp.ifWrites.len())
 
-	wMin := tp.ifWrites.t.Min().(*inFlightWrite)
+	wMin, _ := tp.ifWrites.t.Min()
 	require.Equal(t, cputArgs.Key, wMin.Key)
 	require.Equal(t, cputArgs.Sequence, wMin.Sequence)
-	wMax := tp.ifWrites.t.Max().(*inFlightWrite)
+	wMax, _ := tp.ifWrites.t.Max()
 	require.Equal(t, delArgs.Key, wMax.Key)
 	require.Equal(t, delArgs.Sequence, wMax.Sequence)
 
@@ -384,10 +384,10 @@ func TestTxnPipelinerTrackInFlightWritesPaginatedResponse(t *testing.T) {
 	require.NotNil(t, br)
 	require.Equal(t, 2, tp.ifWrites.len())
 
-	w := tp.ifWrites.t.Min().(*inFlightWrite)
+	w, _ := tp.ifWrites.t.Min()
 	require.Equal(t, putArgs1.Key, w.Key)
 	require.Equal(t, putArgs1.Sequence, w.Sequence)
-	w = tp.ifWrites.t.Max().(*inFlightWrite)
+	w, _ = tp.ifWrites.t.Max()
 	require.Equal(t, putArgs2.Key, w.Key)
 	require.Equal(t, putArgs2.Sequence, w.Sequence)
 
@@ -434,7 +434,7 @@ func TestTxnPipelinerTrackInFlightWritesPaginatedResponse(t *testing.T) {
 	require.Nil(t, pErr)
 	require.Equal(t, 1, tp.ifWrites.len())
 
-	w = tp.ifWrites.t.Min().(*inFlightWrite)
+	w, _ = tp.ifWrites.t.Min()
 	require.Equal(t, putArgs2.Key, w.Key)
 	require.Equal(t, putArgs2.Sequence, w.Sequence)
 }

pkg/kv/kvserver/BUILD.bazel

@@ -263,7 +263,7 @@ go_library(
         "@com_github_kr_pretty//:pretty",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_prometheus_client_model//go",
-        "@com_github_raduberinde_btree//:btree",
+        "@com_github_raduberinde_btreemap//:btreemap",
         "@io_opentelemetry_go_otel//attribute",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_x_time//rate",
@@ -577,7 +577,7 @@ go_test(
         "@com_github_lib_pq//:pq",
         "@com_github_olekukonko_tablewriter//:tablewriter",
         "@com_github_prometheus_common//expfmt",
-        "@com_github_raduberinde_btree//:btree",
+        "@com_github_raduberinde_btreemap//:btreemap",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_grpc//:grpc",

pkg/kv/kvserver/asim/state/BUILD.bazel

@@ -42,7 +42,7 @@ go_library(
         "//pkg/util/metric",
         "//pkg/util/stop",
         "//pkg/util/timeutil",
-        "@com_github_raduberinde_btree//:btree",
+        "@com_github_google_btree//:btree",
         "@org_golang_google_protobuf//proto",
     ],
 )

pkg/kv/kvserver/asim/state/change.go

@@ -9,10 +9,10 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/google/btree"
 )
 
 // Change is a state change for a range, to a target store that has some delay.
@@ -404,7 +404,7 @@ func (rc *ReplicaChange) Blocking() bool {
 // pushed before a change.
 type replicaChanger struct {
 	lastTicket     int
-	completeAt     *btree.BTree
+	completeAt     *btree.BTreeG[*pendingChange]
 	pendingTickets map[int]Change
 	pendingTarget  map[StoreID]time.Time
 	pendingRange   map[RangeID]int
@@ -414,7 +414,7 @@ type replicaChanger struct {
 // replica changes.
 func NewReplicaChanger() Changer {
 	return &replicaChanger{
-		completeAt:     btree.New(8),
+		completeAt:     btree.NewG[*pendingChange](8, (*pendingChange).Less),
 		pendingTickets: make(map[int]Change),
 		pendingTarget:  make(map[StoreID]time.Time),
 		pendingRange:   make(map[RangeID]int),
@@ -426,11 +426,10 @@ type pendingChange struct {
 	completeAt time.Time
 }
 
-// Less is part of the btree.Item interface.
-func (pc *pendingChange) Less(than btree.Item) bool {
-	// Targettal order on (completeAt, ticket)
-	return pc.completeAt.Before(than.(*pendingChange).completeAt) ||
-		(pc.completeAt.Equal(than.(*pendingChange).completeAt) && pc.ticket < than.(*pendingChange).ticket)
+func (pc *pendingChange) Less(than *pendingChange) bool {
+	// Order on (completeAt, ticket).
+	return pc.completeAt.Before(than.completeAt) ||
+		(pc.completeAt.Equal(than.completeAt) && pc.ticket < than.ticket)
 }
 
 // Push appends a state change to occur. There must not be more than one
@@ -477,8 +476,7 @@ func (rc *replicaChanger) Tick(tick time.Time, state State) {
 	// NB: Add the smallest unit of time, in order to find all items in
 	// [smallest, tick].
 	pivot := &pendingChange{completeAt: tick.Add(time.Nanosecond)}
-	rc.completeAt.AscendLessThan(pivot, func(i btree.Item) bool {
-		nextChange, _ := i.(*pendingChange)
+	rc.completeAt.AscendLessThan(pivot, func(nextChange *pendingChange) bool {
 		changeList = append(changeList, nextChange)
 		return true
 	})

pkg/kv/kvserver/asim/state/impl.go

@@ -17,7 +17,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
@@ -34,6 +33,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigreporter"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/google/btree"
 )
 
 type state struct {
@@ -88,28 +88,26 @@ func newState(settings *config.SimulationSettings) *state {
 type rmap struct {
 	// NB: Both rangeTree and rangeMap hold references to ranges. They must
 	// both be updated on insertion and deletion to maintain consistent state.
-	rangeTree *btree.BTree
+	rangeTree *btree.BTreeG[*rng]
 	rangeMap  map[RangeID]*rng
 
 	// Unique ID generator for Ranges.
 	rangeSeqGen RangeID
 }
 
 func newRMap() *rmap {
+	lessFn := func(a, b *rng) bool {
+		return a.startKey < b.startKey
+	}
 	rmap := &rmap{
-		rangeTree: btree.New(8),
+		rangeTree: btree.NewG[*rng](8, lessFn),
 		rangeMap:  make(map[RangeID]*rng),
 	}
 
 	rmap.initFirstRange()
 	return rmap
 }
 
-// Less is part of the btree.Item interface.
-func (r *rng) Less(than btree.Item) bool {
-	return r.startKey < than.(*rng).startKey
-}
-
 // initFirstRange initializes the first range within the rangemap, with
 // [MinKey, MaxKey) start and end key. All other ranges are split from this.
 func (rm *rmap) initFirstRange() {
@@ -164,8 +162,7 @@ func (s *state) String() string {
 	builder := &strings.Builder{}
 
 	orderedRanges := []*rng{}
-	s.ranges.rangeTree.Ascend(func(i btree.Item) bool {
-		r := i.(*rng)
+	s.ranges.rangeTree.Ascend(func(r *rng) bool {
 		orderedRanges = append(orderedRanges, r)
 		return !r.desc.EndKey.Equal(MaxKey.ToRKey())
 	})
@@ -330,8 +327,8 @@ func (s *state) rangeFor(key Key) *rng {
 	var r *rng
 	// If keyToFind equals to MinKey of the range, we found the right range, if
 	// the range is less than keyToFind then this is the right range also.
-	s.ranges.rangeTree.DescendLessOrEqual(keyToFind, func(i btree.Item) bool {
-		r = i.(*rng)
+	s.ranges.rangeTree.DescendLessOrEqual(keyToFind, func(i *rng) bool {
+		r = i
 		return false
 	})
 	return r
@@ -688,8 +685,7 @@ func (s *state) SetSpanConfig(span roachpb.Span, config *roachpb.SpanConfig) {
 	//   [f, z)         - keeps old span config from [c,z)
 
 	splitsRequired := []Key{}
-	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: startKey}, func(i btree.Item) bool {
-		cur, _ := i.(*rng)
+	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: startKey}, func(cur *rng) bool {
 		rStart := cur.startKey
 		// There are two cases we handle:
 		// (1) rStart == startKey: We don't need to split.
@@ -702,8 +698,7 @@ func (s *state) SetSpanConfig(span roachpb.Span, config *roachpb.SpanConfig) {
 		return false
 	})
 
-	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: endKey}, func(i btree.Item) bool {
-		cur, _ := i.(*rng)
+	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: endKey}, func(cur *rng) bool {
 		rEnd := cur.endKey
 		rStart := cur.startKey
 		if rStart == endKey {
@@ -732,8 +727,7 @@ func (s *state) SetSpanConfig(span roachpb.Span, config *roachpb.SpanConfig) {
 	}
 
 	// Apply the span config to all the ranges affected.
-	s.ranges.rangeTree.AscendGreaterOrEqual(&rng{startKey: startKey}, func(i btree.Item) bool {
-		cur, _ := i.(*rng)
+	s.ranges.rangeTree.AscendGreaterOrEqual(&rng{startKey: startKey}, func(cur *rng) bool {
 		if cur.startKey == endKey {
 			return false
 		}
@@ -800,16 +794,15 @@ func (s *state) SplitRange(splitKey Key) (Range, Range, bool) {
 	endKey := Key(math.MaxInt32)
 	failed := false
 	// Find the sucessor range in the range map, to determine the endkey.
-	ranges.rangeTree.AscendGreaterOrEqual(r, func(i btree.Item) bool {
+	ranges.rangeTree.AscendGreaterOrEqual(r, func(i *rng) bool {
 		// The min key already exists in the range map, we cannot return a new
 		// range.
-		if !r.Less(i) {
+		if r.startKey == i.startKey {
 			failed = true
 			return false
 		}
 
-		successorRange, _ := i.(*rng)
-		endKey = successorRange.startKey
+		endKey = i.startKey
 		return false
 	})
 
@@ -822,10 +815,10 @@ func (s *state) SplitRange(splitKey Key) (Range, Range, bool) {
 	var predecessorRange *rng
 	// Find the predecessor range, to update it's endkey to the new range's min
 	// key.
-	ranges.rangeTree.DescendLessOrEqual(r, func(i btree.Item) bool {
+	ranges.rangeTree.DescendLessOrEqual(r, func(i *rng) bool {
 		// The case where the min key already exists cannot occur here, as the
 		// failed flag will have been set above.
-		predecessorRange, _ = i.(*rng)
+		predecessorRange = i
 		return false
 	})
 
@@ -998,8 +991,7 @@ func (s *state) ApplyLoad(lb workload.LoadBatch) {
 	// that range is not larger than the any key of the remaining load events.
 	iter := n - 1
 	max := &rng{startKey: Key(lb[iter].Key)}
-	s.ranges.rangeTree.DescendLessOrEqual(max, func(i btree.Item) bool {
-		next, _ := i.(*rng)
+	s.ranges.rangeTree.DescendLessOrEqual(max, func(next *rng) bool {
 		for iter > -1 && lb[iter].Key >= int64(next.startKey) {
 			s.applyLoad(next, lb[iter])
 			iter--

pkg/kv/kvserver/asim/state/state_test.go

@@ -84,7 +84,9 @@ func TestRangeMap(t *testing.T) {
 	// Assert that the first range is correctly initialized upon creation of a
 	// new state.
 	require.Len(t, s.ranges.rangeMap, 1)
-	require.Equal(t, s.ranges.rangeTree.Max(), s.ranges.rangeTree.Min())
+	maxRes, _ := s.ranges.rangeTree.Max()
+	minRes, _ := s.ranges.rangeTree.Min()
+	require.Equal(t, maxRes, minRes)
 	firstRange := s.ranges.rangeMap[1]
 	require.Equal(t, s.rangeFor(MinKey), firstRange)
 	require.Equal(t, firstRange.startKey, MinKey)