From fdd9e49d32482d875e80b02bb2374be4fcb7dbbc Mon Sep 17 00:00:00 2001 From: Spongecaptain Date: Sat, 24 Sep 2022 10:35:54 +0800 Subject: [PATCH] upate:use google/btree in the genric way --- server/storage/mvcc/index.go | 47 +++++++++++++------------------ server/storage/mvcc/index_test.go | 13 +++++---- server/storage/mvcc/key_index.go | 5 ++-- 3 files changed, 29 insertions(+), 36 deletions(-) diff --git a/server/storage/mvcc/index.go b/server/storage/mvcc/index.go index 954121757ee2..68f3228aebe1 100644 --- a/server/storage/mvcc/index.go +++ b/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,8 +86,8 @@ 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 } @@ -94,11 +95,11 @@ func (ti *treeIndex) keyIndex(keyi *keyIndex) *keyIndex { 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 diff --git a/server/storage/mvcc/index_test.go b/server/storage/mvcc/index_test.go index 8fd5f6c741d5..7d947670d765 100644 --- a/server/storage/mvcc/index_test.go +++ b/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) } diff --git a/server/storage/mvcc/key_index.go b/server/storage/mvcc/key_index.go index fadb6938a03f..e7aac273c9e1 100644 --- a/server/storage/mvcc/key_index.go +++ b/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 {