Fix Rocksdb MVCC read timestamp lifetime for iterators

sei-db/db_engine/pebbledb/mvcc/iterator.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"sync"
 
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
@@ -32,6 +33,8 @@ type iterator struct {
 	reverse            bool
 	iterationCount     int64
 	storeKey           string
+
+	closeSync sync.Once
 }
 
 func newPebbleDBIterator(src *pebble.Iterator, prefix, mvccStart, mvccEnd []byte, version int64, earliestVersion int64, reverse bool, storeKey string) *iterator {
@@ -336,20 +339,21 @@ func (itr *iterator) Error() error {
 }
 
 func (itr *iterator) Close() error {
-	_ = itr.source.Close()
-	itr.source = nil
-	itr.valid = false
-
-	// Record the number of iterations performed by this iterator
-	otelMetrics.iteratorIterations.Record(
-		context.Background(),
-		float64(itr.iterationCount),
-		metric.WithAttributes(
-			attribute.Bool("reverse", itr.reverse),
-			attribute.String("store", itr.storeKey),
-		),
-	)
+	itr.closeSync.Do(func() {
+		_ = itr.source.Close()
+		itr.source = nil
+		itr.valid = false
 
+		// Record the number of iterations performed by this iterator
+		otelMetrics.iteratorIterations.Record(
+			context.Background(),
+			float64(itr.iterationCount),
+			metric.WithAttributes(
+				attribute.Bool("reverse", itr.reverse),
+				attribute.String("store", itr.storeKey),
+			),
+		)
+	})
 	return nil
 }
 

sei-db/db_engine/rocksdb/mvcc/db.go

@@ -128,8 +128,11 @@ func OpenDB(dataDir string, config config.StateStoreConfig) (*Database, error) {
 }
 
 func (db *Database) getSlice(storeKey string, version int64, key []byte) (*grocksdb.Slice, error) {
+	readOpts := newTSReadOptions(version)
+	defer readOpts.Destroy()
+
 	return db.storage.GetCF(
-		newTSReadOptions(version),
+		readOpts,
 		db.cfHandle,
 		prependStoreKey(storeKey, key),
 	)
@@ -323,8 +326,9 @@ func (db *Database) Iterator(storeKey string, version int64, start, end []byte)
 	prefix := storePrefix(storeKey)
 	start, end = util.IterateWithPrefix(prefix, start, end)
 
-	itr := db.storage.NewIteratorCF(newTSReadOptions(version), db.cfHandle)
-	return NewRocksDBIterator(itr, prefix, start, end, version, db.earliestVersion, false), nil
+	readOpts := newTSReadOptions(version)
+	itr := db.storage.NewIteratorCF(readOpts, db.cfHandle)
+	return NewRocksDBIterator(itr, readOpts, prefix, start, end, version, db.earliestVersion, false), nil
 }
 
 func (db *Database) ReverseIterator(storeKey string, version int64, start, end []byte) (types.DBIterator, error) {
@@ -339,8 +343,9 @@ func (db *Database) ReverseIterator(storeKey string, version int64, start, end [
 	prefix := storePrefix(storeKey)
 	start, end = util.IterateWithPrefix(prefix, start, end)
 
-	itr := db.storage.NewIteratorCF(newTSReadOptions(version), db.cfHandle)
-	return NewRocksDBIterator(itr, prefix, start, end, version, db.earliestVersion, true), nil
+	readOpts := newTSReadOptions(version)
+	itr := db.storage.NewIteratorCF(readOpts, db.cfHandle)
+	return NewRocksDBIterator(itr, readOpts, prefix, start, end, version, db.earliestVersion, true), nil
 }
 
 // Import loads the initial version of the state in parallel with numWorkers goroutines
@@ -401,20 +406,19 @@ func (db *Database) RawIterate(storeKey string, fn func(key []byte, value []byte
 		return false, err
 	}
 
-	var startTs [TimestampSize]byte
-	binary.LittleEndian.PutUint64(startTs[:], uint64(0))
+	startTs := make([]byte, TimestampSize)
+	binary.LittleEndian.PutUint64(startTs, uint64(0))
 
-	var endTs [TimestampSize]byte
-	binary.LittleEndian.PutUint64(endTs[:], uint64(latestVersion))
+	endTs := make([]byte, TimestampSize)
+	binary.LittleEndian.PutUint64(endTs, uint64(latestVersion))
 
 	// Set timestamp lower and upper bound to iterate over all keys in db
 	readOpts := grocksdb.NewDefaultReadOptions()
-	readOpts.SetIterStartTimestamp(startTs[:])
-	readOpts.SetTimestamp(endTs[:])
-	defer readOpts.Destroy()
+	readOpts.SetIterStartTimestamp(startTs)
+	readOpts.SetTimestamp(endTs)
 
 	itr := db.storage.NewIteratorCF(readOpts, db.cfHandle)
-	rocksItr := NewRocksDBIterator(itr, prefix, start, end, latestVersion, 1, false)
+	rocksItr := NewRocksDBIterator(itr, readOpts, prefix, start, end, latestVersion, 1, false)
 	defer func() { _ = rocksItr.Close() }()
 
 	for rocksItr.Valid() {
@@ -443,11 +447,11 @@ func (db *Database) GetLatestMigratedModule() (string, error) {
 
 // newTSReadOptions returns ReadOptions used in the RocksDB column family read.
 func newTSReadOptions(version int64) *grocksdb.ReadOptions {
-	var ts [TimestampSize]byte
-	binary.LittleEndian.PutUint64(ts[:], uint64(version))
+	ts := make([]byte, TimestampSize)
+	binary.LittleEndian.PutUint64(ts, uint64(version))
 
 	readOpts := grocksdb.NewDefaultReadOptions()
-	readOpts.SetTimestamp(ts[:])
+	readOpts.SetTimestamp(ts)
 
 	return readOpts
 }

sei-db/db_engine/rocksdb/mvcc/iterator.go

@@ -5,6 +5,7 @@ package mvcc
 
 import (
 	"bytes"
+	"sync"
 
 	"github.com/linxGnu/grocksdb"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/ss/types"
@@ -14,23 +15,26 @@ var _ types.DBIterator = (*iterator)(nil)
 
 type iterator struct {
 	source             *grocksdb.Iterator
+	readOpts           *grocksdb.ReadOptions
 	prefix, start, end []byte
 	version            int64
 	reverse            bool
 	invalid            bool
+	closeOnce          sync.Once
 }
 
-func NewRocksDBIterator(source *grocksdb.Iterator, prefix, start, end []byte, version int64, earliestVersion int64, reverse bool) *iterator {
+func NewRocksDBIterator(source *grocksdb.Iterator, readOpts *grocksdb.ReadOptions, prefix, start, end []byte, version int64, earliestVersion int64, reverse bool) *iterator {
 	// Return invalid iterator if requested iterator height is lower than earliest version after pruning
 	if version < earliestVersion {
 		return &iterator{
-			source:  source,
-			prefix:  prefix,
-			start:   start,
-			end:     end,
-			version: version,
-			reverse: reverse,
-			invalid: true,
+			source:   source,
+			readOpts: readOpts,
+			prefix:   prefix,
+			start:    start,
+			end:      end,
+			version:  version,
+			reverse:  reverse,
+			invalid:  true,
 		}
 	}
 
@@ -58,13 +62,14 @@ func NewRocksDBIterator(source *grocksdb.Iterator, prefix, start, end []byte, ve
 	}
 
 	return &iterator{
-		source:  source,
-		prefix:  prefix,
-		start:   start,
-		end:     end,
-		version: version,
-		reverse: reverse,
-		invalid: !source.Valid(),
+		source:   source,
+		readOpts: readOpts,
+		prefix:   prefix,
+		start:    start,
+		end:      end,
+		version:  version,
+		reverse:  reverse,
+		invalid:  !source.Valid(),
 	}
 }
 
@@ -138,7 +143,7 @@ func (itr *iterator) Value() []byte {
 	return copyAndFreeSlice(itr.source.Value())
 }
 
-func (itr iterator) Next() {
+func (itr *iterator) Next() {
 	if itr.invalid {
 		return
 	}
@@ -155,9 +160,20 @@ func (itr *iterator) Error() error {
 }
 
 func (itr *iterator) Close() error {
-	itr.source.Close()
-	itr.source = nil
-	itr.invalid = true
+	itr.closeOnce.Do(func() {
+		src := itr.source
+		ro := itr.readOpts
+		itr.source = nil
+		itr.readOpts = nil
+		itr.invalid = true
+
+		if src != nil {
+			src.Close()
+		}
+		if ro != nil {
+			ro.Destroy()
+		}
+	})
 	return nil
 }
 

sei-db/state_db/ss/test/storage_test_suite.go

@@ -306,7 +306,8 @@ func (s *StorageTestSuite) TestDatabaseIteratorClose() {
 	s.Require().NoError(iter.Close())
 
 	s.Require().False(iter.Valid())
-	s.Require().Panics(func() { _ = iter.Close() })
+	// Close is idempotent
+	s.Require().NotPanics(func() { _ = iter.Close() })
 }
 
 func (s *StorageTestSuite) TestDatabaseIteratorDomain() {