Add NewBatchWithSize to DB interface

CHANGELOG.md

@@ -4,3 +4,4 @@
 
 - added bloom filter:  https://github.com/cosmos/cosmos-db/pull/42/files
 - Removed Badger & Boltdb
+- Add `NewBatchWithSize` to `DB` interface: https://github.com/cosmos/cosmos-db/pull/64

goleveldb.go

@@ -187,6 +187,11 @@ func (db *GoLevelDB) NewBatch() Batch {
 	return newGoLevelDBBatch(db)
 }
 
+// NewBatchWithSize implements DB.
+func (db *GoLevelDB) NewBatchWithSize(size int) Batch {
+	return newGoLevelDBBatchWithSize(db, size)
+}
+
 // Iterator implements DB.
 func (db *GoLevelDB) Iterator(start, end []byte) (Iterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {

goleveldb_batch.go

@@ -19,6 +19,13 @@ func newGoLevelDBBatch(db *GoLevelDB) *goLevelDBBatch {
 	}
 }
 
+func newGoLevelDBBatchWithSize(db *GoLevelDB, size int) *goLevelDBBatch {
+	return &goLevelDBBatch{
+		db:    db,
+		batch: leveldb.MakeBatch(size),
+	}
+}
+
 // Set implements Batch.
 func (b *goLevelDBBatch) Set(key, value []byte) error {
 	if len(key) == 0 {

memdb.go

@@ -173,6 +173,12 @@ func (db *MemDB) NewBatch() Batch {
 	return newMemDBBatch(db)
 }
 
+// NewBatchWithSize implements DB.
+// It does the same thing as NewBatch because we can't pre-allocate memDBBatch
+func (db *MemDB) NewBatchWithSize(size int) Batch {
+	return newMemDBBatch(db)
+}
+
 // Iterator implements DB.
 // Takes out a read-lock on the database until the iterator is closed.
 func (db *MemDB) Iterator(start, end []byte) (Iterator, error) {

pebble.go

@@ -240,6 +240,12 @@ func (db *PebbleDB) NewBatch() Batch {
 	return newPebbleDBBatch(db)
 }
 
+// NewBatchWithSize implements DB.
+// It does the same thing as NewBatch because we can't pre-allocate pebbleDBBatch
+func (db *PebbleDB) NewBatchWithSize(size int) Batch {
+	return newPebbleDBBatch(db)
+}
+
 // Iterator implements DB.
 func (db *PebbleDB) Iterator(start, end []byte) (Iterator, error) {
 	// fmt.Println("PebbleDB.Iterator")

prefixdb.go

@@ -143,6 +143,11 @@ func (pdb *PrefixDB) NewBatch() Batch {
 	return newPrefixBatch(pdb.prefix, pdb.db.NewBatch())
 }
 
+// NewBatchWithSize implements DB.
+func (pdb *PrefixDB) NewBatchWithSize(size int) Batch {
+	return newPrefixBatch(pdb.prefix, pdb.db.NewBatchWithSize(size))
+}
+
 // Close implements DB.
 func (pdb *PrefixDB) Close() error {
 	pdb.mtx.Lock()

rocksdb.go

@@ -182,6 +182,12 @@ func (db *RocksDB) NewBatch() Batch {
 	return newRocksDBBatch(db)
 }
 
+// NewBatchWithSize implements DB.
+// It does the same thing as NewBatch because we can't pre-allocate rocksDBBatch
+func (db *RocksDB) NewBatchWithSize(size int) Batch {
+	return newRocksDBBatch(db)
+}
+
 // Iterator implements DB.
 func (db *RocksDB) Iterator(start, end []byte) (Iterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {

types.go

@@ -63,6 +63,10 @@ type DB interface {
 	// NewBatch creates a batch for atomic updates. The caller must call Batch.Close.
 	NewBatch() Batch
 
+	// NewBatchWithSize create a new batch for atomic updates, but with pre-allocated size.
+	// This will does the same thing as NewBatch if the batch implementation doesn't support pre-allocation.
+	NewBatchWithSize(int) Batch
+
 	// Print is used for debugging.
 	Print() error