update iterator test

Author: roseduan

index/btree.go

@@ -177,6 +177,7 @@ func newMemoryBTreeIterator(tree *btree.BTree, reverse bool) *memoryBTreeIterato
 		}
 		valid = true
 	}
+
 	return &memoryBTreeIterator{
 		tree:    tree.Clone(),
 		reverse: reverse,
@@ -189,17 +190,20 @@ func (it *memoryBTreeIterator) Rewind() {
 	if it.tree == nil || it.tree.Len() == 0 {
 		return
 	}
+
 	if it.reverse {
 		it.current = it.tree.Max().(*item)
 	} else {
 		it.current = it.tree.Min().(*item)
 	}
+	it.valid = true
 }
 
 func (it *memoryBTreeIterator) Seek(key []byte) {
 	if it.tree == nil || !it.valid {
 		return
 	}
+
 	seekItem := &item{key: key}
 	it.valid = false
 	if it.reverse {
@@ -221,6 +225,7 @@ func (it *memoryBTreeIterator) Next() {
 	if it.tree == nil || !it.valid {
 		return
 	}
+
 	it.valid = false
 	if it.reverse {
 		it.tree.DescendLessOrEqual(it.current, func(i btree.Item) bool {
@@ -241,6 +246,7 @@ func (it *memoryBTreeIterator) Next() {
 			return false
 		})
 	}
+
 	if !it.valid {
 		it.current = nil
 	}

iterator.go

@@ -14,8 +14,10 @@ type Item struct {
 	Value []byte
 }
 
-// Iterator represents a database-level iterator that provides methods to traverse over the key/value pairs in the database.
-// It wraps the index iterator and adds functionality to retrieve the actual values from the database.
+// Iterator represents a database-level iterator that
+// provides methods to traverse over the key/value pairs in the database.
+// It wraps the index iterator and adds functionality to
+// retrieve the actual values from the database.
 type Iterator struct {
 	indexIter index.IndexIterator // index iterator for traversing keys
 	db        *DB                 // database instance for retrieving values
@@ -76,6 +78,7 @@ func (it *Iterator) Item() *Item {
 	if it.db == nil || it.indexIter == nil || !it.Valid() {
 		return nil
 	}
+
 	record := it.skipToNext()
 	if record == nil {
 		return nil
@@ -91,6 +94,7 @@ func (it *Iterator) Close() {
 	if it.db == nil || it.indexIter == nil {
 		return
 	}
+
 	it.indexIter.Close()
 	it.indexIter = nil
 	it.db = nil
@@ -119,13 +123,13 @@ func (it *Iterator) skipToNext() *LogRecord {
 			}
 		}
 
-		// Check if the key is expired
 		position := it.indexIter.Value()
 		if position == nil {
 			it.indexIter.Next()
 			continue
 		}
 
+		// read the record from data file
 		chunk, err := it.db.dataFiles.Read(position)
 		if err != nil {
 			it.lastError = err
@@ -140,7 +144,8 @@ func (it *Iterator) skipToNext() *LogRecord {
 
 		// Skip if record is deleted or expired
 		record := decodeLogRecord(chunk)
-		if record.Type == LogRecordDeleted || record.IsExpired(time.Now().UnixNano()) {
+		now := time.Now().UnixNano()
+		if record.Type == LogRecordDeleted || record.IsExpired(now) {
 			it.indexIter.Next()
 			continue
 		}

iterator_test.go

@@ -246,3 +246,53 @@ func TestIterator_Error(t *testing.T) {
 	assert.False(t, iter.Valid())
 	iter.Close()
 }
+
+func TestIterator_UseTwice(t *testing.T) {
+	options := DefaultOptions
+	db, err := Open(options)
+	assert.Nil(t, err)
+	defer destroyDB(db)
+
+	db.Put([]byte("aceds"), []byte("value1"))
+	db.Put([]byte("eedsq"), []byte("value2"))
+	db.Put([]byte("sedas"), []byte("value3"))
+	db.Put([]byte("efeds"), []byte("value4"))
+	db.Put([]byte("bbdes"), []byte("value5"))
+
+	iteratorOptions := DefaultIteratorOptions
+	iter := db.NewIterator(iteratorOptions)
+
+	for iter.Seek([]byte("bbe")); iter.Valid(); iter.Next() {
+		item := iter.Item()
+		assert.NotNil(t, item)
+	}
+
+	// rewind and iterate again
+	iter.Rewind()
+	assert.Equal(t, iter.Valid(), true)
+	for iter.Seek([]byte("bbe")); iter.Valid(); iter.Next() {
+		item := iter.Item()
+		assert.NotNil(t, item)
+	}
+}
+
+func TestIterator_UseAfterClose(t *testing.T) {
+	options := DefaultOptions
+	db, err := Open(options)
+	assert.Nil(t, err)
+	defer destroyDB(db)
+
+	db.Put([]byte("aceds"), []byte("value1"))
+	db.Put([]byte("eedsq"), []byte("value2"))
+
+	iteratorOptions := DefaultIteratorOptions
+	iter := db.NewIterator(iteratorOptions)
+
+	for iter.Seek([]byte("bbe")); iter.Valid(); iter.Next() {
+		item := iter.Item()
+		assert.NotNil(t, item)
+	}
+
+	iter.Close()
+	assert.Equal(t, iter.Valid(), false)
+}