kvserver: switch to btreemap for store replica map

Epic: none
Release note: None

Author: RaduBerinde

pkg/kv/kvserver/BUILD.bazel

@@ -260,10 +260,10 @@ go_library(
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_gogo_protobuf//proto",
-        "@com_github_google_btree//:btree",
         "@com_github_kr_pretty//:pretty",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_prometheus_client_model//go",
+        "@com_github_raduberinde_btreemap//:btreemap",
         "@io_opentelemetry_go_otel//attribute",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_x_time//rate",
@@ -572,12 +572,12 @@ go_test(
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_dustin_go_humanize//:go-humanize",
         "@com_github_gogo_protobuf//proto",
-        "@com_github_google_btree//:btree",
         "@com_github_jackc_pgx_v5//:pgx",
         "@com_github_kr_pretty//:pretty",
         "@com_github_lib_pq//:pq",
         "@com_github_olekukonko_tablewriter//:tablewriter",
         "@com_github_prometheus_common//expfmt",
+        "@com_github_raduberinde_btreemap//:btreemap",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_grpc//:grpc",

pkg/kv/kvserver/replica_placeholder.go

@@ -84,17 +84,11 @@ type ReplicaPlaceholder struct {
 	tainted   int32 // atomic
 }
 
-var _ rangeKeyItem = (*ReplicaPlaceholder)(nil)
-
 // Desc returns the range Placeholder's descriptor.
 func (r *ReplicaPlaceholder) Desc() *roachpb.RangeDescriptor {
 	return &r.rangeDesc
 }
 
-func (r *ReplicaPlaceholder) key() roachpb.RKey {
-	return r.Desc().StartKey
-}
-
 func (r *ReplicaPlaceholder) String() string {
 	tainted := ""
 	if atomic.LoadInt32(&r.tainted) != 0 {

pkg/kv/kvserver/scanner_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/RaduBerinde/btreemap"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -21,7 +22,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
-	"github.com/google/btree"
 )
 
 func makeAmbCtx() log.AmbientContext {
@@ -31,16 +31,15 @@ func makeAmbCtx() log.AmbientContext {
 // Test implementation of a range set backed by btree.BTree.
 type testRangeSet struct {
 	syncutil.Mutex
-	replicasByKey *btree.BTreeG[*Replica]
+	replicasByKey *btreemap.BTreeMap[roachpb.RKey, *Replica]
 	visited       int
 }
 
 // newTestRangeSet creates a new range set that has the count number of ranges.
 func newTestRangeSet(count int, t *testing.T) *testRangeSet {
-	lessFn := func(a, b *Replica) bool {
-		return a.startKey.Less(b.startKey)
+	rs := &testRangeSet{
+		replicasByKey: btreemap.New[roachpb.RKey, *Replica](64 /* degree */, roachpb.RKey.Compare),
 	}
-	rs := &testRangeSet{replicasByKey: btree.NewG[*Replica](64 /* degree */, lessFn)}
 	for i := 0; i < count; i++ {
 		desc := &roachpb.RangeDescriptor{
 			RangeID:  roachpb.RangeID(i),
@@ -59,7 +58,7 @@ func newTestRangeSet(count int, t *testing.T) *testRangeSet {
 		}
 		repl.shMu.state.Desc = desc
 		repl.startKey = desc.StartKey // actually used by replicasByKey
-		if exRngItem, _ := rs.replicasByKey.ReplaceOrInsert(repl); exRngItem != nil {
+		if _, exRngItem, _ := rs.replicasByKey.ReplaceOrInsert(repl.startKey, repl); exRngItem != nil {
 			t.Fatalf("failed to insert range %s", repl)
 		}
 	}
@@ -70,12 +69,15 @@ func (rs *testRangeSet) Visit(visitor func(*Replica) bool) {
 	rs.Lock()
 	defer rs.Unlock()
 	rs.visited = 0
-	rs.replicasByKey.Ascend(func(r *Replica) bool {
-		rs.visited++
-		rs.Unlock()
-		defer rs.Lock()
-		return visitor(r)
-	})
+	rs.replicasByKey.AscendFunc(
+		btreemap.Min[roachpb.RKey](), btreemap.Max[roachpb.RKey](),
+		func(_ roachpb.RKey, r *Replica) bool {
+			rs.visited++
+			rs.Unlock()
+			defer rs.Lock()
+			return visitor(r)
+		},
+	)
 }
 
 func (rs *testRangeSet) EstimatedCount() int {
@@ -93,7 +95,7 @@ func (rs *testRangeSet) remove(index int, t *testing.T) *Replica {
 	endKey := roachpb.RKey(fmt.Sprintf("%03d", index+1))
 	rs.Lock()
 	defer rs.Unlock()
-	repl, _ := rs.replicasByKey.Delete(&Replica{startKey: endKey})
+	_, repl, _ := rs.replicasByKey.Delete(endKey)
 	if repl == nil {
 		t.Fatalf("failed to delete range of end key %s", endKey)
 	}

pkg/kv/kvserver/store_create_replica.go

@@ -264,13 +264,13 @@ func (s *Store) addToReplicasByKeyLocked(repl *Replica, desc *roachpb.RangeDescr
 	if got := s.GetReplicaIfExists(repl.RangeID); got != repl { // NB: got can be nil too
 		return errors.Errorf("%s: replica %s not in replicasByRangeID; got %s", s, repl, got)
 	}
-	if it := s.getOverlappingKeyRangeLocked(desc); it.item != nil {
+	if it := s.getOverlappingKeyRangeLocked(desc); !it.isEmpty() {
 		return errors.Errorf(
 			"%s: cannot add to replicasByKey: range %s overlaps with %s", s, repl, it.Desc())
 	}
-	if it := s.mu.replicasByKey.ReplaceOrInsertReplica(context.Background(), repl); it.item != nil {
+	if it := s.mu.replicasByKey.ReplaceOrInsertReplica(context.Background(), repl); !it.isEmpty() {
 		return errors.Errorf(
-			"%s: cannot add to replicasByKey: key %v already exists in the btree", s, it.item.key())
+			"%s: cannot add to replicasByKey: key %v already exists in the btree", s, it.key())
 	}
 	return nil
 }
@@ -279,8 +279,8 @@ func (s *Store) addToReplicasByKeyLocked(repl *Replica, desc *roachpb.RangeDescr
 // and the raftMu of the replica whose place is being held are locked.
 func (s *Store) addPlaceholderLocked(placeholder *ReplicaPlaceholder) error {
 	rangeID := placeholder.Desc().RangeID
-	if it := s.mu.replicasByKey.ReplaceOrInsertPlaceholder(context.Background(), placeholder); it.item != nil {
-		return errors.Errorf("%s overlaps with existing replicaOrPlaceholder %+v in replicasByKey btree", placeholder, it.item)
+	if it := s.mu.replicasByKey.ReplaceOrInsertPlaceholder(context.Background(), placeholder); !it.isEmpty() {
+		return errors.Errorf("%s overlaps with existing replicaOrPlaceholder %+v in replicasByKey btree", placeholder, it.item())
 	}
 	if exRng, ok := s.mu.replicaPlaceholders[rangeID]; ok {
 		return errors.Errorf("%s has ID collision with placeholder %+v", placeholder, exRng)

pkg/kv/kvserver/store_merge.go

@@ -62,7 +62,7 @@ func (s *Store) maybeAssertNoHole(ctx context.Context, from, to roachpb.RKey) fu
 						// TODO(tbg): this deadlocks, see #74384. By disabling this branch,
 						// the only check that we perform is the one that the monitored
 						// keyspace isn't all a gap.
-						if last.item != nil {
+						if !last.isEmpty() {
 							gapStart, gapEnd := last.Desc().EndKey, cur.Desc().StartKey
 							if !gapStart.Equal(gapEnd) {
 								return errors.AssertionFailedf(
@@ -77,7 +77,7 @@ func (s *Store) maybeAssertNoHole(ctx context.Context, from, to roachpb.RKey) fu
 				if err != nil {
 					log.Fatalf(ctx, "%v", err)
 				}
-				if last.item == nil {
+				if last.isEmpty() {
 					log.Fatalf(ctx, "found hole in keyspace [%s,%s), during:\n%s", from, to, caller)
 				}
 				runtime.Gosched()

pkg/kv/kvserver/store_remove_replica.go

@@ -157,7 +157,7 @@ func (s *Store) removeInitializedReplicaRaftMuLocked(
 		// this far. This method will need some changes when we introduce GC of
 		// uninitialized replicas.
 		s.mu.Unlock()
-		log.Fatalf(ctx, "replica %+v unexpectedly overlapped by %+v", rep, it.item)
+		log.Fatalf(ctx, "replica %+v unexpectedly overlapped by %+v", rep, it.item())
 	}
 	// Adjust stats before calling Destroy. This can be called before or after
 	// Destroy, but this configuration helps avoid races in stat verification
@@ -202,7 +202,7 @@ func (s *Store) removeInitializedReplicaRaftMuLocked(
 		if it := s.mu.replicasByKey.ReplaceOrInsertPlaceholder(ctx, ph); it.repl != rep {
 			// We already checked that our replica was present in replicasByKey
 			// above. Nothing should have been able to change that.
-			log.Fatalf(ctx, "replica %+v unexpectedly overlapped by %+v", rep, it.item)
+			log.Fatalf(ctx, "replica %+v unexpectedly overlapped by %+v", rep, it.item())
 		}
 		if exPH, ok := s.mu.replicaPlaceholders[desc.RangeID]; ok {
 			log.Fatalf(ctx, "cannot insert placeholder %s, already have %s", ph, exPH)
@@ -218,8 +218,8 @@ func (s *Store) removeInitializedReplicaRaftMuLocked(
 
 		s.mu.replicasByKey.DeletePlaceholder(ctx, ph)
 		delete(s.mu.replicaPlaceholders, desc.RangeID)
-		if it := s.getOverlappingKeyRangeLocked(desc); it.item != nil && it.item != ph {
-			log.Fatalf(ctx, "corrupted replicasByKey map: %s and %s overlapped", rep, it.item)
+		if it := s.getOverlappingKeyRangeLocked(desc); !it.isEmpty() && it.item() != ph {
+			log.Fatalf(ctx, "corrupted replicasByKey map: %s and %s overlapped", rep, it.item())
 		}
 		return nil
 	}()
@@ -400,8 +400,8 @@ func (s *Store) removePlaceholderLocked(
 		return false, errors.AssertionFailedf("placeholder %+v not found, got %+v", placeholder, it)
 	}
 	delete(s.mu.replicaPlaceholders, rngID)
-	if it := s.getOverlappingKeyRangeLocked(&placeholder.rangeDesc); it.item != nil {
-		return false, errors.AssertionFailedf("corrupted replicasByKey map: %+v and %+v overlapped", it.ph, it.item)
+	if it := s.getOverlappingKeyRangeLocked(&placeholder.rangeDesc); !it.isEmpty() {
+		return false, errors.AssertionFailedf("corrupted replicasByKey map: %+v and %+v overlapped", it.ph, it.item())
 	}
 	switch typ {
 	case removePlaceholderDropped: