From b295cebc05062a152253d1313cfd6a0d0465fa30 Mon Sep 17 00:00:00 2001
From: SimFG <bang.fu@zilliz.com>
Date: Wed, 1 Jun 2022 17:48:13 +0800
Subject: [PATCH] mvcc: improve the use of locks in index.go

To make the meaning of the RangeSince function in treeIndex clearer.
---
 server/storage/mvcc/index.go | 48 +++++++++++++++++++-----------------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/server/storage/mvcc/index.go b/server/storage/mvcc/index.go
index be817c5a6..eb8b82d2a 100644
--- a/server/storage/mvcc/index.go
+++ b/server/storage/mvcc/index.go
@@ -67,9 +67,13 @@ func (ti *treeIndex) Put(key []byte, rev revision) {
 }
 
 func (ti *treeIndex) Get(key []byte, atRev int64) (modified, created revision, ver int64, err error) {
-	keyi := &keyIndex{key: key}
 	ti.RLock()
 	defer ti.RUnlock()
+	return ti.unsafeGet(key, atRev)
+}
+
+func (ti *treeIndex) unsafeGet(key []byte, atRev int64) (modified, created revision, ver int64, err error) {
+	keyi := &keyIndex{key: key}
 	if keyi = ti.keyIndex(keyi); keyi == nil {
 		return revision{}, revision{}, 0, ErrRevisionNotFound
 	}
@@ -89,12 +93,9 @@ func (ti *treeIndex) keyIndex(keyi *keyIndex) *keyIndex {
 	return nil
 }
 
-func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex) bool) {
+func (ti *treeIndex) unsafeVisit(key, end []byte, f func(ki *keyIndex) bool) {
 	keyi, endi := &keyIndex{key: key}, &keyIndex{key: end}
 
-	ti.RLock()
-	defer ti.RUnlock()
-
 	ti.tree.AscendGreaterOrEqual(keyi, func(item btree.Item) bool {
 		if len(endi.key) > 0 && !item.Less(endi) {
 			return false
@@ -107,14 +108,17 @@ func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex) bool) {
 }
 
 func (ti *treeIndex) Revisions(key, end []byte, atRev int64, limit int) (revs []revision, total int) {
+	ti.RLock()
+	defer ti.RUnlock()
+
 	if end == nil {
-		rev, _, _, err := ti.Get(key, atRev)
+		rev, _, _, err := ti.unsafeGet(key, atRev)
 		if err != nil {
 			return nil, 0
 		}
 		return []revision{rev}, 1
 	}
-	ti.visit(key, end, func(ki *keyIndex) bool {
+	ti.unsafeVisit(key, end, func(ki *keyIndex) bool {
 		if rev, _, _, err := ki.get(ti.lg, atRev); err == nil {
 			if limit <= 0 || len(revs) < limit {
 				revs = append(revs, rev)
@@ -127,15 +131,18 @@ func (ti *treeIndex) Revisions(key, end []byte, atRev int64, limit int) (revs []
 }
 
 func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int {
+	ti.RLock()
+	defer ti.RUnlock()
+
 	if end == nil {
-		_, _, _, err := ti.Get(key, atRev)
+		_, _, _, err := ti.unsafeGet(key, atRev)
 		if err != nil {
 			return 0
 		}
 		return 1
 	}
 	total := 0
-	ti.visit(key, end, func(ki *keyIndex) bool {
+	ti.unsafeVisit(key, end, func(ki *keyIndex) bool {
 		if _, _, _, err := ki.get(ti.lg, atRev); err == nil {
 			total++
 		}
@@ -145,14 +152,17 @@ func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int {
 }
 
 func (ti *treeIndex) Range(key, end []byte, atRev int64) (keys [][]byte, revs []revision) {
+	ti.RLock()
+	defer ti.RUnlock()
+
 	if end == nil {
-		rev, _, _, err := ti.Get(key, atRev)
+		rev, _, _, err := ti.unsafeGet(key, atRev)
 		if err != nil {
 			return nil, nil
 		}
 		return [][]byte{key}, []revision{rev}
 	}
-	ti.visit(key, end, func(ki *keyIndex) bool {
+	ti.unsafeVisit(key, end, func(ki *keyIndex) bool {
 		if rev, _, _, err := ki.get(ti.lg, atRev); err == nil {
 			revs = append(revs, rev)
 			keys = append(keys, ki.key)
@@ -180,28 +190,20 @@ func (ti *treeIndex) Tombstone(key []byte, rev revision) error {
 // at or after the given rev. The returned slice is sorted in the order
 // of revision.
 func (ti *treeIndex) RangeSince(key, end []byte, rev int64) []revision {
-	keyi := &keyIndex{key: key}
-
 	ti.RLock()
 	defer ti.RUnlock()
 
 	if end == nil {
-		item := ti.tree.Get(keyi)
-		if item == nil {
+		keyi := ti.keyIndex(&keyIndex{key: key})
+		if keyi == nil {
 			return nil
 		}
-		keyi = item.(*keyIndex)
 		return keyi.since(ti.lg, rev)
 	}
 
-	endi := &keyIndex{key: end}
 	var revs []revision
-	ti.tree.AscendGreaterOrEqual(keyi, func(item btree.Item) bool {
-		if len(endi.key) > 0 && !item.Less(endi) {
-			return false
-		}
-		curKeyi := item.(*keyIndex)
-		revs = append(revs, curKeyi.since(ti.lg, rev)...)
+	ti.unsafeVisit(key, end, func(ki *keyIndex) bool {
+		revs = append(revs, ki.since(ti.lg, rev)...)
 		return true
 	})
 	sort.Sort(revisions(revs))