diff --git a/ledis/tx.go b/ledis/tx.go index f49cd34..5bd1a1a 100644 --- a/ledis/tx.go +++ b/ledis/tx.go @@ -27,7 +27,7 @@ func newTx(l *Ledis) *tx { } func (t *tx) Close() { - t.wb.Close() + t.wb = nil } func (t *tx) Put(key []byte, value []byte) { diff --git a/store/driver/driver.go b/store/driver/driver.go index b39e700..a557d3e 100644 --- a/store/driver/driver.go +++ b/store/driver/driver.go @@ -40,8 +40,6 @@ type IIterator interface { } type IWriteBatch interface { - Close() error - Put(key []byte, value []byte) Delete(key []byte) Commit() error diff --git a/store/goleveldb/batch.go b/store/goleveldb/batch.go index f3ff67c..b17e85e 100644 --- a/store/goleveldb/batch.go +++ b/store/goleveldb/batch.go @@ -9,10 +9,6 @@ type WriteBatch struct { wbatch *leveldb.Batch } -func (w *WriteBatch) Close() error { - return nil -} - func (w *WriteBatch) Put(key, value []byte) { w.wbatch.Put(key, value) } diff --git a/store/leveldb/batch.go b/store/leveldb/batch.go index 87600cb..08a7d46 100644 --- a/store/leveldb/batch.go +++ b/store/leveldb/batch.go @@ -17,6 +17,8 @@ type WriteBatch struct { func (w *WriteBatch) Close() error { C.leveldb_writebatch_destroy(w.wbatch) + w.wbatch = nil + return nil } diff --git a/store/leveldb/db.go b/store/leveldb/db.go index f85e04b..049c9f4 100644 --- a/store/leveldb/db.go +++ b/store/leveldb/db.go @@ -13,6 +13,7 @@ import "C" import ( "github.com/siddontang/ledisdb/store/driver" "os" + "runtime" "unsafe" ) @@ -187,6 +188,10 @@ func (db *DB) NewWriteBatch() driver.IWriteBatch { db: db, wbatch: C.leveldb_writebatch_create(), } + + runtime.SetFinalizer(wb, func(w *WriteBatch) { + w.Close() + }) return wb } diff --git a/store/mdb/batch.go b/store/mdb/batch.go index 1638435..bab56ec 100644 --- a/store/mdb/batch.go +++ b/store/mdb/batch.go @@ -14,10 +14,6 @@ type WriteBatch struct { wb []Write } -func (w *WriteBatch) Close() error { - return nil -} - func (w *WriteBatch) Put(key, value []byte) { w.wb = append(w.wb, Write{key, value}) } diff --git a/store/rocksdb/batch.go b/store/rocksdb/batch.go index 46e700b..b69c383 100644 --- a/store/rocksdb/batch.go +++ b/store/rocksdb/batch.go @@ -17,6 +17,7 @@ type WriteBatch struct { func (w *WriteBatch) Close() error { C.rocksdb_writebatch_destroy(w.wbatch) + w.wbatch = nil return nil } diff --git a/store/rocksdb/db.go b/store/rocksdb/db.go index 28a0544..2d8a56d 100644 --- a/store/rocksdb/db.go +++ b/store/rocksdb/db.go @@ -207,6 +207,11 @@ func (db *DB) NewWriteBatch() driver.IWriteBatch { db: db, wbatch: C.rocksdb_writebatch_create(), } + + runtime.SetFinalizer(wb, func(w *WriteBatch) { + w.Close() + }) + return wb } diff --git a/store/store_test.go b/store/store_test.go index f37c496..4c6c74d 100644 --- a/store/store_test.go +++ b/store/store_test.go @@ -49,7 +49,6 @@ func testBatch(db *DB, t *testing.T) { db.Put(key2, value) wb := db.NewWriteBatch() - defer wb.Close() wb.Delete(key2) wb.Put(key1, []byte("hello world2"))