boltdb: handle empty key placeholders in iterators

CHANGELOG.md

@@ -14,6 +14,10 @@
 
 ### Bug Fixes
 
+- [boltdb] Properly handle blank keys in iterators
+
+- [boltdb] Properly close batches
+
 - [goleveldb] [\#58](https://github.com/tendermint/tm-db/pull/58) Make `Batch.Close()` actually remove the batch contents
 
 ## 0.4.1

boltdb.go

@@ -11,7 +11,10 @@ import (
 	"github.com/pkg/errors"
 )
 
-var bucket = []byte("tm")
+var (
+	bucket         = []byte("tm")
+	boltDBEmptyKey = []byte("nil")
+)
 
 func init() {
 	registerDBCreator(BoltDBBackend, func(name, dir string) (DB, error) {
@@ -196,7 +199,7 @@ func (bdb *BoltDB) ReverseIterator(start, end []byte) (Iterator, error) {
 // WARNING: this may collude with "nil" user key!
 func nonEmptyKey(key []byte) []byte {
 	if len(key) == 0 {
-		return []byte("nil")
+		return boltDBEmptyKey
 	}
 	return key
 }

boltdb_batch.go

@@ -49,4 +49,6 @@ func (b *boltDBBatch) WriteSync() error {
 }
 
 // Close implements Batch.
-func (b *boltDBBatch) Close() {}
+func (b *boltDBBatch) Close() {
+	b.ops = nil
+}

boltdb_iterator.go

@@ -13,9 +13,10 @@ import (
 type boltDBIterator struct {
 	tx *bbolt.Tx
 
-	itr   *bbolt.Cursor
-	start []byte
-	end   []byte
+	itr           *bbolt.Cursor
+	start         []byte
+	end           []byte
+	emptyKeyValue []byte // Tracks the value of the empty key, if it exists
 
 	currentKey   []byte
 	currentValue []byte
@@ -28,33 +29,58 @@ var _ Iterator = (*boltDBIterator)(nil)
 
 // newBoltDBIterator creates a new boltDBIterator.
 func newBoltDBIterator(tx *bbolt.Tx, start, end []byte, isReverse bool) *boltDBIterator {
+	// We can check for empty key at the start, because we use a read/write transaction that blocks
+	// the entire database for writes while the iterator exists. If we change to a read-only txn
+	// that supports concurrency we'll need to rewrite this logic.
+	emptyKeyValue := tx.Bucket(bucket).Get(boltDBEmptyKey)
 	itr := tx.Bucket(bucket).Cursor()
 
 	var ck, cv []byte
 	if isReverse {
-		if end == nil {
+		switch {
+		case end == nil:
 			ck, cv = itr.Last()
-		} else {
+		case len(end) == 0:
+			// If end is the blank key, then we don't return any keys by definition
+			ck = nil
+			cv = nil
+		default:
 			_, _ = itr.Seek(end) // after key
 			ck, cv = itr.Prev()  // return to end key
 		}
+		// If we're currently positioned at the placeholder for the empty key, skip it (handle later)
+		if emptyKeyValue != nil && bytes.Equal(ck, boltDBEmptyKey) {
+			ck, cv = itr.Prev()
+		}
+		// If we didn't find any initial key, but there's a placeholder for the empty key at the
+		// end that we've skipped, then the initial key should be the empty one (the final one).
+		if emptyKeyValue != nil && ck == nil && (end == nil || len(end) > 0) {
+			ck = []byte{}
+			cv = emptyKeyValue
+			emptyKeyValue = nil // ensure call to Next() skips this
+		}
 	} else {
-		if start == nil {
+		switch {
+		case (start == nil || len(start) == 0) && emptyKeyValue != nil:
+			ck = []byte{}
+			cv = emptyKeyValue
+		case (start == nil || len(start) == 0) && emptyKeyValue == nil:
 			ck, cv = itr.First()
-		} else {
+		default:
 			ck, cv = itr.Seek(start)
 		}
 	}
 
 	return &boltDBIterator{
-		tx:           tx,
-		itr:          itr,
-		start:        start,
-		end:          end,
-		currentKey:   ck,
-		currentValue: cv,
-		isReverse:    isReverse,
-		isInvalid:    false,
+		tx:            tx,
+		itr:           itr,
+		start:         start,
+		end:           end,
+		emptyKeyValue: emptyKeyValue,
+		currentKey:    ck,
+		currentValue:  cv,
+		isReverse:     isReverse,
+		isInvalid:     false,
 	}
 }
 
@@ -70,7 +96,7 @@ func (itr *boltDBIterator) Valid() bool {
 	}
 
 	// iterated to the end of the cursor
-	if len(itr.currentKey) == 0 {
+	if itr.currentKey == nil {
 		itr.isInvalid = true
 		return false
 	}
@@ -96,8 +122,24 @@ func (itr *boltDBIterator) Next() {
 	itr.assertIsValid()
 	if itr.isReverse {
 		itr.currentKey, itr.currentValue = itr.itr.Prev()
+		if itr.emptyKeyValue != nil && itr.currentKey == nil {
+			// If we reached the end, but there exists an empty key whose placeholder we skipped,
+			// we should set up the empty key and its value as the final pair.
+			itr.currentKey = []byte{}
+			itr.currentValue = itr.emptyKeyValue
+			itr.emptyKeyValue = nil // This ensures the next call to Next() terminates
+		}
 	} else {
-		itr.currentKey, itr.currentValue = itr.itr.Next()
+		if len(itr.currentKey) == 0 {
+			// If the first key was the empty key, then we need to move to the first non-empty key
+			itr.currentKey, itr.currentValue = itr.itr.First()
+		} else {
+			itr.currentKey, itr.currentValue = itr.itr.Next()
+		}
+	}
+	// If we encounter the placeholder for the empty key, skip it
+	if itr.emptyKeyValue != nil && bytes.Equal(itr.currentKey, boltDBEmptyKey) {
+		itr.Next()
 	}
 }