Merge pull request #44 from sei-protocol/PebbleDBIteratorFix

[SeiDB] Fix PebbleDB Iterator Issue

Author: Kbhat1

ss/pebbledb/comparator.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/cockroachdb/pebble"
+	"github.com/sei-protocol/sei-db/common/utils"
 )
 
 // MVCCComparer returns a PebbleDB Comparer with encoding and decoding routines
@@ -137,20 +138,25 @@ func (f mvccKeyFormatter) Format(s fmt.State, verb rune) {
 
 // SplitMVCCKey accepts an MVCC key and returns the "user" key, the MVCC version,
 // and a boolean indicating if the provided key is an MVCC key.
+//
+// Note, internally, we must make a copy of the provided mvccKey argument, which
+// typically comes from the Key() method as it's not safe.
 func SplitMVCCKey(mvccKey []byte) (key, version []byte, ok bool) {
 	if len(mvccKey) == 0 {
 		return nil, nil, false
 	}
 
-	n := len(mvccKey) - 1
-	tsLen := int(mvccKey[n])
+	mvccKeyCopy := utils.Clone(mvccKey)
+
+	n := len(mvccKeyCopy) - 1
+	tsLen := int(mvccKeyCopy[n])
 	if n < tsLen {
 		return nil, nil, false
 	}
 
-	key = mvccKey[:n-tsLen]
+	key = mvccKeyCopy[:n-tsLen]
 	if tsLen > 0 {
-		version = mvccKey[n-tsLen+1 : len(mvccKey)-1]
+		version = mvccKeyCopy[n-tsLen+1 : len(mvccKeyCopy)-1]
 	}
 
 	return key, version, true

ss/pebbledb/iterator.go

@@ -116,15 +116,15 @@ func (itr *iterator) Value() []byte {
 }
 
 func (itr *iterator) Next() {
+	currKey, _, ok := SplitMVCCKey(itr.source.Key())
+	if !ok {
+		// XXX: This should not happen as that would indicate we have a malformed
+		// MVCC key.
+		panic(fmt.Sprintf("invalid PebbleDB MVCC key: %s", itr.source.Key()))
+	}
+
 	var next bool
 	if itr.reverse {
-		currKey, _, ok := SplitMVCCKey(itr.source.Key())
-		if !ok {
-			// XXX: This should not happen as that would indicate we have a malformed
-			// MVCC key.
-			panic(fmt.Sprintf("invalid PebbleDB MVCC key: %s", itr.source.Key()))
-		}
-
 		// Since PebbleDB has no PrevPrefix API, we must manually seek to the next
 		// key that is lexicographically less than the current key.
 		next = itr.source.SeekLT(MVCCEncode(currKey, 0))
@@ -154,6 +154,25 @@ func (itr *iterator) Next() {
 		// append the current iterator key to the prefix and seek to that key.
 		itr.valid = itr.source.SeekLT(MVCCEncode(nextKey, itr.version+1))
 
+		tmpKey, _, ok := SplitMVCCKey(itr.source.Key())
+		if !ok {
+			// XXX: This should not happen as that would indicate we have a malformed
+			// MVCC key.
+			itr.valid = false
+			return
+		}
+
+		// There exists cases where the SeekLT() call moved us back to the same key
+		// we started at, so we must move to next key, i.e. two keys forward.
+		if bytes.Equal(tmpKey, currKey) {
+			if itr.source.NextPrefix() {
+				itr.Next()
+			} else {
+				itr.valid = false
+				return
+			}
+		}
+
 		// The cursor might now be pointing at a key/value pair that is tombstoned.
 		// If so, we must move the cursor.
 		if itr.valid && itr.cursorTombstoned() {

ss/test/storage_test_suite.go

@@ -369,6 +369,51 @@ func (s *StorageTestSuite) TestDatabaseIteratorMultiVersion() {
 	s.Require().NoError(itr.Error())
 }
 
+// Tests bug where iterator loops continuously
+func (s *StorageTestSuite) TestDatabaseIteratorLooping() {
+	db, err := s.NewDB(s.T().TempDir())
+	s.Require().NoError(err)
+	defer db.Close()
+
+	// Forward Iteration
+	// Less than iterator version
+	s.Require().NoError(DBApplyChangeset(db, 58827506, storeKey1, [][]byte{[]byte("keyC")}, [][]byte{[]byte("value003")}))
+	// Both below are greater
+	s.Require().NoError(DBApplyChangeset(db, 58833605, storeKey1, [][]byte{[]byte("keyC")}, [][]byte{[]byte("value004")}))
+	s.Require().NoError(DBApplyChangeset(db, 58833606, storeKey1, [][]byte{[]byte("keyD")}, [][]byte{[]byte("value006")}))
+	s.Require().NoError(DBApplyChangeset(db, 58827507, storeKey1, [][]byte{[]byte("keyE")}, [][]byte{[]byte("value006")}))
+	s.Require().NoError(DBApplyChangeset(db, 58827508, storeKey1, [][]byte{[]byte("keyF")}, [][]byte{[]byte("value007")}))
+	// Another case of a loop
+	s.Require().NoError(DBApplyChangeset(db, 58827509, storeKey1, [][]byte{[]byte("keyG")}, [][]byte{[]byte("value008")}))
+	s.Require().NoError(DBApplyChangeset(db, 58831545, storeKey1, [][]byte{[]byte("keyG")}, [][]byte{[]byte("value009")}))
+	s.Require().NoError(DBApplyChangeset(db, 58831565, storeKey1, [][]byte{[]byte("keyH")}, [][]byte{[]byte("value010")}))
+
+	itr, err := db.Iterator(storeKey1, 58831525, []byte("keyA"), nil)
+	s.Require().NoError(err)
+
+	defer itr.Close()
+
+	count := 0
+	for ; itr.Valid(); itr.Next() {
+		count++
+	}
+
+	s.Require().Equal(4, count)
+
+	// Reverse Iteration
+	itr, err = db.Iterator(storeKey1, 58831525, nil, []byte("keyZ"))
+	s.Require().NoError(err)
+
+	defer itr.Close()
+
+	count = 0
+	for ; itr.Valid(); itr.Next() {
+		count++
+	}
+
+	s.Require().Equal(4, count)
+}
+
 func (s *StorageTestSuite) TestDatabaseIteratorNoDomain() {
 	db, err := s.NewDB(s.T().TempDir())
 	s.Require().NoError(err)