upate:use google/btree in the genric way

server/storage/mvcc/index.go

@@ -38,14 +38,16 @@ type index interface {
 
 type treeIndex struct {
 	sync.RWMutex
-	tree *btree.BTree
+	tree *btree.BTreeG[*keyIndex]
 	lg   *zap.Logger
 }
 
 func newTreeIndex(lg *zap.Logger) index {
 	return &treeIndex{
-		tree: btree.New(32),
-		lg:   lg,
+		tree: btree.NewG(32, func(aki *keyIndex, bki *keyIndex) bool {
+			return aki.Less(bki)
+		}),
+		lg: lg,
 	}
 }
 
@@ -54,13 +56,12 @@ func (ti *treeIndex) Put(key []byte, rev revision) {
 
 	ti.Lock()
 	defer ti.Unlock()
-	item := ti.tree.Get(keyi)
-	if item == nil {
+	okeyi, _ := ti.tree.Get(keyi)
+	if okeyi == nil {
 		keyi.put(ti.lg, rev.main, rev.sub)
 		ti.tree.ReplaceOrInsert(keyi)
 		return
 	}
-	okeyi := item.(*keyIndex)
 	okeyi.put(ti.lg, rev.main, rev.sub)
 }
 
@@ -85,20 +86,20 @@ func (ti *treeIndex) KeyIndex(keyi *keyIndex) *keyIndex {
 }
 
 func (ti *treeIndex) keyIndex(keyi *keyIndex) *keyIndex {
-	if item := ti.tree.Get(keyi); item != nil {
-		return item.(*keyIndex)
+	if ki, ok := ti.tree.Get(keyi); ok {
+		return ki
 	}
 	return nil
 }
 
 func (ti *treeIndex) unsafeVisit(key, end []byte, f func(ki *keyIndex) bool) {
 	keyi, endi := &keyIndex{key: key}, &keyIndex{key: end}
 
-	ti.tree.AscendGreaterOrEqual(keyi, func(item btree.Item) bool {
+	ti.tree.AscendGreaterOrEqual(keyi, func(item *keyIndex) bool {
 		if len(endi.key) > 0 && !item.Less(endi) {
 			return false
 		}
-		if !f(item.(*keyIndex)) {
+		if !f(item) {
 			return false
 		}
 		return true
@@ -180,12 +181,11 @@ func (ti *treeIndex) Tombstone(key []byte, rev revision) error {
 
 	ti.Lock()
 	defer ti.Unlock()
-	item := ti.tree.Get(keyi)
-	if item == nil {
+	ki, ok := ti.tree.Get(keyi)
+	if !ok {
 		return ErrRevisionNotFound
 	}
 
-	ki := item.(*keyIndex)
 	return ki.tombstone(ti.lg, rev.main, rev.sub)
 }
 
@@ -196,15 +196,14 @@ func (ti *treeIndex) Compact(rev int64) map[revision]struct{} {
 	clone := ti.tree.Clone()
 	ti.Unlock()
 
-	clone.Ascend(func(item btree.Item) bool {
-		keyi := item.(*keyIndex)
+	clone.Ascend(func(keyi *keyIndex) bool {
 		// Lock is needed here to prevent modification to the keyIndex while
 		// compaction is going on or revision added to empty before deletion
 		ti.Lock()
 		keyi.compact(ti.lg, rev, available)
 		if keyi.isEmpty() {
-			item := ti.tree.Delete(keyi)
-			if item == nil {
+			_, ok := ti.tree.Delete(keyi)
+			if !ok {
 				ti.lg.Panic("failed to delete during compaction")
 			}
 		}
@@ -219,8 +218,7 @@ func (ti *treeIndex) Keep(rev int64) map[revision]struct{} {
 	available := make(map[revision]struct{})
 	ti.RLock()
 	defer ti.RUnlock()
-	ti.tree.Ascend(func(i btree.Item) bool {
-		keyi := i.(*keyIndex)
+	ti.tree.Ascend(func(keyi *keyIndex) bool {
 		keyi.keep(rev, available)
 		return true
 	})
@@ -236,15 +234,8 @@ func (ti *treeIndex) Equal(bi index) bool {
 
 	equal := true
 
-	ti.tree.Ascend(func(item btree.Item) bool {
-		var aki, bki *keyIndex
-		var ok bool
-		if aki, ok = item.(*keyIndex); !ok {
-			return false
-		}
-		if bki, ok = b.tree.Get(item).(*keyIndex); !ok {
-			return false
-		}
+	ti.tree.Ascend(func(aki *keyIndex) bool {
+		bki, _ := b.tree.Get(aki)
 		if !aki.equal(bki) {
 			equal = false
 			return false

server/storage/mvcc/index_test.go

@@ -270,7 +270,9 @@ func TestIndexCompactAndKeep(t *testing.T) {
 		if !(reflect.DeepEqual(am, keep)) {
 			t.Errorf("#%d: compact keep %v != Keep keep %v", i, am, keep)
 		}
-		wti := &treeIndex{tree: btree.New(32)}
+		wti := &treeIndex{tree: btree.NewG(32, func(aki *keyIndex, bki *keyIndex) bool {
+			return aki.Less(bki)
+		})}
 		for _, tt := range tests {
 			if _, ok := am[tt.rev]; ok || tt.rev.GreaterThan(revision{main: i}) {
 				if tt.remove {
@@ -300,7 +302,9 @@ func TestIndexCompactAndKeep(t *testing.T) {
 		if !(reflect.DeepEqual(am, keep)) {
 			t.Errorf("#%d: compact keep %v != Keep keep %v", i, am, keep)
 		}
-		wti := &treeIndex{tree: btree.New(32)}
+		wti := &treeIndex{tree: btree.NewG(32, func(aki *keyIndex, bki *keyIndex) bool {
+			return aki.Less(bki)
+		})}
 		for _, tt := range tests {
 			if _, ok := am[tt.rev]; ok || tt.rev.GreaterThan(revision{main: i}) {
 				if tt.remove {
@@ -321,12 +325,11 @@ func restore(ti *treeIndex, key []byte, created, modified revision, ver int64) {
 
 	ti.Lock()
 	defer ti.Unlock()
-	item := ti.tree.Get(keyi)
-	if item == nil {
+	okeyi, _ := ti.tree.Get(keyi)
+	if okeyi == nil {
 		keyi.restore(ti.lg, created, modified, ver)
 		ti.tree.ReplaceOrInsert(keyi)
 		return
 	}
-	okeyi := item.(*keyIndex)
 	okeyi.put(ti.lg, modified.main, modified.sub)
 }

server/storage/mvcc/key_index.go

@@ -19,7 +19,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/google/btree"
 	"go.uber.org/zap"
 )
 
@@ -305,8 +304,8 @@ func (ki *keyIndex) findGeneration(rev int64) *generation {
 	return nil
 }
 
-func (ki *keyIndex) Less(b btree.Item) bool {
-	return bytes.Compare(ki.key, b.(*keyIndex).key) == -1
+func (ki *keyIndex) Less(bki *keyIndex) bool {
+	return bytes.Compare(ki.key, bki.key) == -1
 }
 
 func (ki *keyIndex) equal(b *keyIndex) bool {