aggmetric: switch to new btree

Use `BtreeG[MetricItem]` in the `aggmetric` code.

Informs: #144504
Release note: None

Author: RaduBerinde

pkg/util/metric/aggmetric/BUILD.bazel

@@ -18,8 +18,8 @@ go_library(
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_gogo_protobuf//proto",
+        "@com_github_google_btree//:btree",
         "@com_github_prometheus_client_model//go",
-        "@com_github_raduberinde_btree//:btree",
     ],
 )
 

pkg/util/metric/aggmetric/agg_metric.go

@@ -13,11 +13,11 @@ import (
 	"strings"
 	"sync/atomic"
 
-	"github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/util/cache"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
+	"github.com/google/btree"
 	io_prometheus_client "github.com/prometheus/client_model/go"
 )
 
@@ -86,8 +86,21 @@ type childSet struct {
 
 func (cs *childSet) initWithBTreeStorageType(labels []string) {
 	cs.labels = labels
+
+	lessFn := func(a, b MetricItem) bool {
+		av, bv := a.labelValues(), b.labelValues()
+		if len(av) != len(bv) {
+			panic(errors.AssertionFailedf("mismatch in label values lengths %v vs %v", av, bv))
+		}
+		for i := range av {
+			if cmp := strings.Compare(av[i], bv[i]); cmp != 0 {
+				return cmp < 0
+			}
+		}
+		return false
+	}
 	cs.mu.children = &BtreeWrapper{
-		tree: btree.New(8),
+		tree: btree.NewG[MetricItem](8, lessFn),
 	}
 }
 
@@ -280,11 +293,6 @@ type MetricItem interface {
 	labelValuer
 }
 
-type BtreeMetricItem interface {
-	btree.Item
-	MetricItem
-}
-
 type CacheMetricItem interface {
 	MetricItem
 }
@@ -364,34 +372,33 @@ func (ucw *UnorderedCacheWrapper) Clear() {
 }
 
 type BtreeWrapper struct {
-	tree *btree.BTree
+	tree *btree.BTreeG[MetricItem]
 }
 
 func (b BtreeWrapper) Get(labelVals ...string) (ChildMetric, bool) {
 	key := labelValuesSlice(labelVals)
-	cm := b.tree.Get(&key)
-	if cm == nil {
+	cm, ok := b.tree.Get(&key)
+	if !ok {
 		return nil, false
 	}
 	return cm.(ChildMetric), true
 }
 
 func (b BtreeWrapper) Add(metric ChildMetric) {
-	if b.tree.Has(metric.(BtreeMetricItem)) {
+	if b.tree.Has(metric) {
 		panic(errors.AssertionFailedf("child %v already exists", metric.labelValues()))
 	}
-	b.tree.ReplaceOrInsert(metric.(BtreeMetricItem))
+	b.tree.ReplaceOrInsert(metric)
 }
 
 func (b BtreeWrapper) Del(metric ChildMetric) {
-	if existing := b.tree.Delete(metric.(btree.Item)); existing == nil {
-		panic(errors.AssertionFailedf(
-			"child %v does not exists", metric.labelValues()))
+	if _, ok := b.tree.Delete(metric); !ok {
+		panic(errors.AssertionFailedf("child %v does not exist", metric.labelValues()))
 	}
 }
 
 func (b BtreeWrapper) ForEach(f func(metric ChildMetric)) {
-	b.tree.Ascend(func(i btree.Item) bool {
+	b.tree.Ascend(func(i MetricItem) bool {
 		f(i.(ChildMetric))
 		return true
 	})
@@ -400,17 +407,3 @@ func (b BtreeWrapper) ForEach(f func(metric ChildMetric)) {
 func (b BtreeWrapper) Clear() {
 	b.tree.Clear(false)
 }
-
-func (lv *labelValuesSlice) Less(o btree.Item) bool {
-	ov := o.(labelValuer).labelValues()
-	if len(ov) != len(*lv) {
-		panic(errors.AssertionFailedf("mismatch in label values lengths %v vs %v",
-			ov, *lv))
-	}
-	for i := range ov {
-		if cmp := strings.Compare((*lv)[i], ov[i]); cmp != 0 {
-			return cmp < 0
-		}
-	}
-	return false // eq
-}