Calling non-Close() methods on closed or written batches panic

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Breaking Changes
 
+- [\#71](https://github.com/tendermint/tm-db/pull/71) Closed or written batches can no longer be reused, all non-`Close()` calls will panic
+
 - [memdb] [\#56](https://github.com/tendermint/tm-db/pull/56) Removed some exported methods that were mainly meant for internal use: `Mutex()`, `SetNoLock()`, `SetNoLockSync()`, `DeleteNoLock()`, and `DeleteNoLockSync()`
 
 ### Improvements
@@ -14,11 +16,9 @@
 
 ### Bug Fixes
 
-- [boltdb] Properly handle blank keys in iterators
-
-- [cleveldb] Fix handling of empty keys as iterator endpoints
+- [boltdb] [\#69](https://github.com/tendermint/tm-db/pull/69) Properly handle blank keys in iterators
 
-- [goleveldb] [\#58](https://github.com/tendermint/tm-db/pull/58) Make `Batch.Close()` actually remove the batch contents
+- [cleveldb] [\#65](https://github.com/tendermint/tm-db/pull/65) Fix handling of empty keys as iterator endpoints
 
 ## 0.4.1
 

backend_test.go

@@ -439,41 +439,14 @@ func testDBBatch(t *testing.T, backend BackendType) {
 	require.NoError(t, err)
 	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}, "c": {3}})
 
-	// the batch still keeps these values internally, so changing values and rewriting batch
-	// should set the values again
-	err = db.Set([]byte("a"), []byte{9})
-	require.NoError(t, err)
-	err = db.Delete([]byte("c"))
-	require.NoError(t, err)
-	err = batch.WriteSync()
-	require.NoError(t, err)
-	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}, "c": {3}})
-
-	// but when we close, it should no longer set the values
-	batch.Close()
-	err = db.Delete([]byte("c"))
-	require.NoError(t, err)
-	err = batch.Write()
-	require.NoError(t, err)
-	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}})
-
-	// it should be possible to re-close the batch
-	batch.Close()
-
-	// it should also be possible to reuse a closed batch as if it were a new one
-	batch.Set([]byte("c"), []byte{3})
-	err = batch.Write()
-	require.NoError(t, err)
-	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}, "c": {3}})
-	batch.Close()
-
-	batch.Delete([]byte("c"))
-	err = batch.WriteSync()
-	require.NoError(t, err)
-	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}})
+	// trying to modify or rewrite a written batch should panic, but closing it should work
+	require.Panics(t, func() { batch.Set([]byte("a"), []byte{9}) })
+	require.Panics(t, func() { batch.Delete([]byte("a")) })
+	require.Panics(t, func() { batch.Write() })
+	require.Panics(t, func() { batch.WriteSync() })
 	batch.Close()
 
-	// batches should also write changes in order
+	// batches should write changes in order
 	batch = db.NewBatch()
 	batch.Delete([]byte("a"))
 	batch.Set([]byte("a"), []byte{1})
@@ -486,13 +459,38 @@ func testDBBatch(t *testing.T, backend BackendType) {
 	batch.Close()
 	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}})
 
-	// and writing an empty batch should not fail
+	// writing nil keys and values should be the same as empty keys and values
+	// FIXME CLevelDB panics here: https://github.com/jmhodges/levigo/issues/55
+	if backend != CLevelDBBackend {
+		batch = db.NewBatch()
+		batch.Set(nil, nil)
+		err = batch.WriteSync()
+		require.NoError(t, err)
+		assertKeyValues(t, db, map[string][]byte{"": {}, "a": {1}, "b": {2}})
+
+		batch = db.NewBatch()
+		batch.Delete(nil)
+		err = batch.Write()
+		require.NoError(t, err)
+		assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}})
+	}
+
+	// it should be possible to write an empty batch
 	batch = db.NewBatch()
 	err = batch.Write()
 	require.NoError(t, err)
-	err = batch.WriteSync()
-	require.NoError(t, err)
 	assertKeyValues(t, db, map[string][]byte{"a": {1}, "b": {2}})
+
+	// it should be possible to close an empty batch, and to re-close a closed batch
+	batch = db.NewBatch()
+	batch.Close()
+	batch.Close()
+
+	// all other operations on a closed batch should panic
+	require.Panics(t, func() { batch.Set([]byte("a"), []byte{9}) })
+	require.Panics(t, func() { batch.Delete([]byte("a")) })
+	require.Panics(t, func() { batch.Write() })
+	require.Panics(t, func() { batch.WriteSync() })
 }
 
 func assertKeyValues(t *testing.T, db DB, expect map[string][]byte) {

boltdb.go

@@ -169,10 +169,7 @@ func (bdb *BoltDB) Stats() map[string]string {
 
 // NewBatch implements DB.
 func (bdb *BoltDB) NewBatch() Batch {
-	return &boltDBBatch{
-		ops: nil,
-		db:  bdb,
-	}
+	return newBoltDBBatch(bdb)
 }
 
 // WARNING: Any concurrent writes or reads will block until the iterator is

boltdb_batch.go

@@ -12,19 +12,35 @@ type boltDBBatch struct {
 
 var _ Batch = (*boltDBBatch)(nil)
 
+func newBoltDBBatch(db *BoltDB) *boltDBBatch {
+	return &boltDBBatch{
+		db:  db,
+		ops: make([]operation, 0, 1),
+	}
+}
+
+func (b *boltDBBatch) assertOpen() {
+	if b.ops == nil {
+		panic("batch has been written or closed")
+	}
+}
+
 // Set implements Batch.
 func (b *boltDBBatch) Set(key, value []byte) {
+	b.assertOpen()
 	b.ops = append(b.ops, operation{opTypeSet, key, value})
 }
 
 // Delete implements Batch.
 func (b *boltDBBatch) Delete(key []byte) {
+	b.assertOpen()
 	b.ops = append(b.ops, operation{opTypeDelete, key, nil})
 }
 
 // Write implements Batch.
 func (b *boltDBBatch) Write() error {
-	return b.db.db.Batch(func(tx *bbolt.Tx) error {
+	b.assertOpen()
+	err := b.db.db.Batch(func(tx *bbolt.Tx) error {
 		bkt := tx.Bucket(bucket)
 		for _, op := range b.ops {
 			key := nonEmptyKey(nonNilBytes(op.key))
@@ -41,6 +57,12 @@ func (b *boltDBBatch) Write() error {
 		}
 		return nil
 	})
+	if err != nil {
+		return err
+	}
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for errors.
+	b.Close()
+	return nil
 }
 
 // WriteSync implements Batch.
@@ -49,4 +71,6 @@ func (b *boltDBBatch) WriteSync() error {
 }
 
 // Close implements Batch.
-func (b *boltDBBatch) Close() {}
+func (b *boltDBBatch) Close() {
+	b.ops = nil
+}

cleveldb.go

@@ -159,8 +159,7 @@ func (db *CLevelDB) Stats() map[string]string {
 
 // NewBatch implements DB.
 func (db *CLevelDB) NewBatch() Batch {
-	batch := levigo.NewWriteBatch()
-	return &cLevelDBBatch{db, batch}
+	return newCLevelDBBatch(db)
 }
 
 // Iterator implements DB.

cleveldb_batch.go

@@ -10,33 +10,59 @@ type cLevelDBBatch struct {
 	batch *levigo.WriteBatch
 }
 
+func newCLevelDBBatch(db *CLevelDB) *cLevelDBBatch {
+	return &cLevelDBBatch{
+		db:    db,
+		batch: levigo.NewWriteBatch(),
+	}
+}
+
+func (b *cLevelDBBatch) assertOpen() {
+	if b.batch == nil {
+		panic("batch has been written or closed")
+	}
+}
+
 // Set implements Batch.
 func (b *cLevelDBBatch) Set(key, value []byte) {
-	b.batch.Put(key, value)
+	b.assertOpen()
+	b.batch.Put(nonNilBytes(key), nonNilBytes(value))
 }
 
 // Delete implements Batch.
 func (b *cLevelDBBatch) Delete(key []byte) {
-	b.batch.Delete(key)
+	b.assertOpen()
+	b.batch.Delete(nonNilBytes(key))
 }
 
 // Write implements Batch.
 func (b *cLevelDBBatch) Write() error {
-	if err := b.db.db.Write(b.db.wo, b.batch); err != nil {
+	b.assertOpen()
+	err := b.db.db.Write(b.db.wo, b.batch)
+	if err != nil {
 		return err
 	}
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for errors.
+	b.Close()
 	return nil
 }
 
 // WriteSync implements Batch.
 func (b *cLevelDBBatch) WriteSync() error {
-	if err := b.db.db.Write(b.db.woSync, b.batch); err != nil {
+	b.assertOpen()
+	err := b.db.db.Write(b.db.woSync, b.batch)
+	if err != nil {
 		return err
 	}
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for errors.
+	b.Close()
 	return nil
 }
 
 // Close implements Batch.
 func (b *cLevelDBBatch) Close() {
-	b.batch.Close()
+	if b.batch != nil {
+		b.batch.Close()
+		b.batch = nil
+	}
 }

goleveldb.go

@@ -153,8 +153,7 @@ func (db *GoLevelDB) Stats() map[string]string {
 
 // NewBatch implements DB.
 func (db *GoLevelDB) NewBatch() Batch {
-	batch := new(leveldb.Batch)
-	return &goLevelDBBatch{db, batch}
+	return newGoLevelDBBatch(db)
 }
 
 // Iterator implements DB.

goleveldb_batch.go

@@ -12,35 +12,56 @@ type goLevelDBBatch struct {
 
 var _ Batch = (*goLevelDBBatch)(nil)
 
+func newGoLevelDBBatch(db *GoLevelDB) *goLevelDBBatch {
+	return &goLevelDBBatch{
+		db:    db,
+		batch: new(leveldb.Batch),
+	}
+}
+
+func (b *goLevelDBBatch) assertOpen() {
+	if b.batch == nil {
+		panic("batch has been written or closed")
+	}
+}
+
 // Set implements Batch.
 func (b *goLevelDBBatch) Set(key, value []byte) {
+	b.assertOpen()
 	b.batch.Put(key, value)
 }
 
 // Delete implements Batch.
 func (b *goLevelDBBatch) Delete(key []byte) {
+	b.assertOpen()
 	b.batch.Delete(key)
 }
 
 // Write implements Batch.
 func (b *goLevelDBBatch) Write() error {
-	err := b.db.db.Write(b.batch, &opt.WriteOptions{Sync: false})
-	if err != nil {
-		return err
-	}
-	return nil
+	return b.write(false)
 }
 
 // WriteSync implements Batch.
 func (b *goLevelDBBatch) WriteSync() error {
-	err := b.db.db.Write(b.batch, &opt.WriteOptions{Sync: true})
+	return b.write(true)
+}
+
+func (b *goLevelDBBatch) write(sync bool) error {
+	b.assertOpen()
+	err := b.db.db.Write(b.batch, &opt.WriteOptions{Sync: sync})
 	if err != nil {
 		return err
 	}
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for errors.
+	b.Close()
 	return nil
 }
 
 // Close implements Batch.
 func (b *goLevelDBBatch) Close() {
-	b.batch.Reset()
+	if b.batch != nil {
+		b.batch.Reset()
+		b.batch = nil
+	}
 }

memdb.go

@@ -150,7 +150,7 @@ func (db *MemDB) Stats() map[string]string {
 
 // NewBatch implements DB.
 func (db *MemDB) NewBatch() Batch {
-	return &memDBBatch{db, nil}
+	return newMemDBBatch(db)
 }
 
 // Iterator implements DB.

memdb_batch.go

@@ -24,18 +24,35 @@ type memDBBatch struct {
 
 var _ Batch = (*memDBBatch)(nil)
 
+// newMemDBBatch creates a new memDBBatch
+func newMemDBBatch(db *MemDB) *memDBBatch {
+	return &memDBBatch{
+		db:  db,
+		ops: make([]operation, 0, 1),
+	}
+}
+
+func (b *memDBBatch) assertOpen() {
+	if b.ops == nil {
+		panic("batch has been written or closed")
+	}
+}
+
 // Set implements Batch.
 func (b *memDBBatch) Set(key, value []byte) {
+	b.assertOpen()
 	b.ops = append(b.ops, operation{opTypeSet, key, value})
 }
 
 // Delete implements Batch.
 func (b *memDBBatch) Delete(key []byte) {
+	b.assertOpen()
 	b.ops = append(b.ops, operation{opTypeDelete, key, nil})
 }
 
 // Write implements Batch.
 func (b *memDBBatch) Write() error {
+	b.assertOpen()
 	b.db.mtx.Lock()
 	defer b.db.mtx.Unlock()
 
@@ -49,6 +66,9 @@ func (b *memDBBatch) Write() error {
 			return errors.Errorf("unknown operation type %v (%v)", op.opType, op)
 		}
 	}
+
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for errors.
+	b.Close()
 	return nil
 }