implement NewBatchWithSize method on DB interface to pre-allocate buffer

goleveldb.go

@@ -172,7 +172,12 @@ func (db *GoLevelDB) ForceCompact(start, limit []byte) error {
 
 // NewBatch implements DB.
 func (db *GoLevelDB) NewBatch() Batch {
-	return newGoLevelDBBatch(db)
+	return newGoLevelDBBatchWithSize(db, 0)
+}
+
+// NewBatchWithSize implements DB.
+func (db *GoLevelDB) NewBatchWithSize(size int) Batch {
+	return newGoLevelDBBatchWithSize(db, size)
 }
 
 // Iterator implements DB.

goleveldb_batch.go

@@ -12,10 +12,10 @@ type goLevelDBBatch struct {
 
 var _ Batch = (*goLevelDBBatch)(nil)
 
-func newGoLevelDBBatch(db *GoLevelDB) *goLevelDBBatch {
+func newGoLevelDBBatchWithSize(db *GoLevelDB, size int) *goLevelDBBatch {
 	return &goLevelDBBatch{
 		db:    db,
-		batch: new(leveldb.Batch),
+		batch: leveldb.MakeBatch(size),
 	}
 }
 

memdb.go

@@ -170,7 +170,12 @@ func (db *MemDB) Stats() map[string]string {
 
 // NewBatch implements DB.
 func (db *MemDB) NewBatch() Batch {
-	return newMemDBBatch(db)
+	return newMemDBBatchWithSize(db, 0)
+}
+
+// NewBatchWithSize implements DB.
+func (db *MemDB) NewBatchWithSize(size int) Batch {
+	return newMemDBBatchWithSize(db, size)
 }
 
 // Iterator implements DB.

memdb_batch.go

@@ -24,11 +24,11 @@ type memDBBatch struct {
 
 var _ Batch = (*memDBBatch)(nil)
 
-// newMemDBBatch creates a new memDBBatch
-func newMemDBBatch(db *MemDB) *memDBBatch {
+// newMemDBBatchWithSize creates a new memDBBatch
+func newMemDBBatchWithSize(db *MemDB, size int) *memDBBatch {
 	return &memDBBatch{
 		db:  db,
-		ops: []operation{},
+		ops: make([]operation, 0, size),
 	}
 }
 

prefixdb.go

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

remotedb/batch.go

@@ -17,10 +17,10 @@ type batch struct {
 
 var _ db.Batch = (*batch)(nil)
 
-func newBatch(rdb *RemoteDB) *batch {
+func newBatchWithSize(rdb *RemoteDB, size int) *batch {
 	return &batch{
 		db:  rdb,
-		ops: []*protodb.Operation{},
+		ops: make([]*protodb.Operation, 0, size),
 	}
 }
 

remotedb/remotedb.go

@@ -97,7 +97,11 @@ func (rd *RemoteDB) ReverseIterator(start, end []byte) (db.Iterator, error) {
 }
 
 func (rd *RemoteDB) NewBatch() db.Batch {
-	return newBatch(rd)
+	return newBatchWithSize(rd, 0)
+}
+
+func (rd *RemoteDB) NewBatchWithSize(size int) db.Batch {
+	return newBatchWithSize(rd, size)
 }
 
 // TODO: Implement Print when db.DB implements a method

types.go

@@ -63,6 +63,9 @@ type DB interface {
 	// NewBatch creates a batch for atomic updates. The caller must call Batch.Close.
 	NewBatch() Batch
 
+	// NewBatch creates a batch for atomic updates with pre-allocated buffer. The caller must call Batch.Close.
+	NewBatchWithSize(size int) Batch
+
 	// Print is used for debugging.
 	Print() error