write batch add Close

store/driver/batch.go

@@ -21,6 +21,11 @@ type WriteBatch struct {
 	w  BatchPuter
 }
 
+func (wb *WriteBatch) Close() {
+	wb.d.Reset()
+	wb.wb = wb.wb[0:0]
+}
+
 func (wb *WriteBatch) Put(key, value []byte) {
 	if value == nil {
 		value = []byte{}

store/driver/driver.go

@@ -59,6 +59,7 @@ type IWriteBatch interface {
 	SyncCommit() error
 	Rollback() error
 	Data() []byte
+	Close()
 }
 
 type Tx interface {

store/goleveldb/batch.go

@@ -30,6 +30,10 @@ func (w *WriteBatch) Rollback() error {
 	return nil
 }
 
+func (w *WriteBatch) Close() {
+	w.wbatch.Reset()
+}
+
 func (w *WriteBatch) Data() []byte {
 	return w.wbatch.Dump()
 }

store/leveldb/batch.go

@@ -34,13 +34,13 @@ func newWriteBatch(db *DB) *WriteBatch {
 	return w
 }
 
-func (w *WriteBatch) Close() error {
-	C.leveldb_writebatch_destroy(w.wbatch)
-	w.wbatch = nil
+func (w *WriteBatch) Close() {
+	if w.wbatch != nil {
+		C.leveldb_writebatch_destroy(w.wbatch)
+		w.wbatch = nil
+	}
 
 	w.gbatch = nil
-
-	return nil
 }
 
 func (w *WriteBatch) Put(key, value []byte) {

store/leveldb/db.go

@@ -14,6 +14,7 @@ import (
 	"github.com/siddontang/ledisdb/config"
 	"github.com/siddontang/ledisdb/store/driver"
 	"os"
+	"runtime"
 	"unsafe"
 )
 
@@ -182,7 +183,13 @@ func (db *DB) SyncDelete(key []byte) error {
 }
 
 func (db *DB) NewWriteBatch() driver.IWriteBatch {
-	return newWriteBatch(db)
+	wb := newWriteBatch(db)
+
+	runtime.SetFinalizer(wb, func(w *WriteBatch) {
+		w.Close()
+	})
+
+	return wb
 }
 
 func (db *DB) NewIterator() driver.IIterator {

store/rocksdb/batch.go

@@ -17,10 +17,11 @@ type WriteBatch struct {
 	commitOk bool
 }
 
-func (w *WriteBatch) Close() error {
-	C.rocksdb_writebatch_destroy(w.wbatch)
-	w.wbatch = nil
-	return nil
+func (w *WriteBatch) Close() {
+	if w.wbatch != nil {
+		C.rocksdb_writebatch_destroy(w.wbatch)
+		w.wbatch = nil
+	}
 }
 
 func (w *WriteBatch) Put(key, value []byte) {

store/rocksdb/db.go

@@ -15,6 +15,7 @@ import (
 	"github.com/siddontang/ledisdb/config"
 	"github.com/siddontang/ledisdb/store/driver"
 	"os"
+	"runtime"
 	"unsafe"
 )
 
@@ -215,6 +216,10 @@ func (db *DB) NewWriteBatch() driver.IWriteBatch {
 		wbatch: C.rocksdb_writebatch_create(),
 	}
 
+	runtime.SetFinalizer(wb, func(w *WriteBatch) {
+		w.Close()
+	})
+
 	return wb
 }
 

store/writebatch.go

@@ -16,6 +16,10 @@ type WriteBatch struct {
 	db        *DB
 }
 
+func (wb *WriteBatch) Close() {
+	wb.wb.Close()
+}
+
 func (wb *WriteBatch) Put(key []byte, value []byte) {
 	wb.putNum++
 	wb.wb.Put(key, value)