From 8fb200819c0f0f27acfc301f9ab7db22775402f2 Mon Sep 17 00:00:00 2001
From: siddontang <siddontang@gmail.com>
Date: Mon, 10 Nov 2014 21:43:25 +0800
Subject: [PATCH] update leveldb batch data

---
 store/leveldb/batch.go       | 41 +++++++++++++++++++-----------------
 store/leveldb/leveldb_ext.cc |  7 ++++++
 store/leveldb/leveldb_ext.h  |  1 +
 3 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/store/leveldb/batch.go b/store/leveldb/batch.go
index b588735..4e2f1c1 100644
--- a/store/leveldb/batch.go
+++ b/store/leveldb/batch.go
@@ -4,6 +4,7 @@ package leveldb
 
 // #cgo LDFLAGS: -lleveldb
 // #include "leveldb/c.h"
+// #include "leveldb_ext.h"
 import "C"
 
 import (
@@ -11,13 +12,6 @@ import (
 	"unsafe"
 )
 
-/*
-	It's not easy to let leveldb support Data() function like rocksdb,
-	so here, we will use goleveldb batch instead
-
-	later optimize
-*/
-
 type WriteBatch struct {
 	db     *DB
 	wbatch *C.leveldb_writebatch_t
@@ -44,8 +38,6 @@ func (w *WriteBatch) Close() {
 }
 
 func (w *WriteBatch) Put(key, value []byte) {
-	w.gbatch.Put(key, value)
-
 	var k, v *C.char
 	if len(key) != 0 {
 		k = (*C.char)(unsafe.Pointer(&key[0]))
@@ -61,29 +53,21 @@ func (w *WriteBatch) Put(key, value []byte) {
 }
 
 func (w *WriteBatch) Delete(key []byte) {
-	w.gbatch.Delete(key)
-
 	C.leveldb_writebatch_delete(w.wbatch,
 		(*C.char)(unsafe.Pointer(&key[0])), C.size_t(len(key)))
 }
 
 func (w *WriteBatch) Commit() error {
-	w.gbatch.Reset()
-
 	return w.commit(w.db.writeOpts)
 }
 
 func (w *WriteBatch) SyncCommit() error {
-	w.gbatch.Reset()
-
 	return w.commit(w.db.syncOpts)
 }
 
 func (w *WriteBatch) Rollback() error {
 	C.leveldb_writebatch_clear(w.wbatch)
 
-	w.gbatch.Reset()
-
 	return nil
 }
 
@@ -96,6 +80,25 @@ func (w *WriteBatch) commit(wb *WriteOptions) error {
 	return nil
 }
 
-func (w *WriteBatch) Data() []byte {
-	return w.gbatch.Dump()
+//export leveldb_writebatch_iterate_put
+func leveldb_writebatch_iterate_put(p unsafe.Pointer, k *C.char, klen C.size_t, v *C.char, vlen C.size_t) {
+	b := (*leveldb.Batch)(p)
+	key := slice(unsafe.Pointer(k), int(klen))
+	value := slice(unsafe.Pointer(v), int(vlen))
+	b.Put(key, value)
+}
+
+//export leveldb_writebatch_iterate_delete
+func leveldb_writebatch_iterate_delete(p unsafe.Pointer, k *C.char, klen C.size_t) {
+	b := (*leveldb.Batch)(p)
+	key := slice(unsafe.Pointer(k), int(klen))
+	b.Delete(key)
+}
+
+func (w *WriteBatch) Data() []byte {
+	w.gbatch.Reset()
+	C.leveldb_writebatch_iterate_ext(w.wbatch,
+		unsafe.Pointer(w.gbatch))
+	b := w.gbatch.Dump()
+	return b
 }
diff --git a/store/leveldb/leveldb_ext.cc b/store/leveldb/leveldb_ext.cc
index a362ab5..540b739 100644
--- a/store/leveldb/leveldb_ext.cc
+++ b/store/leveldb/leveldb_ext.cc
@@ -84,5 +84,12 @@ unsigned char leveldb_iter_prev_ext(leveldb_iterator_t* iter) {
     return leveldb_iter_valid(iter);
 }
 
+extern void leveldb_writebatch_iterate_put(void*, const char* k, size_t klen, const char* v, size_t vlen);
+extern void leveldb_writebatch_iterate_delete(void*, const char* k, size_t klen);
+
+void leveldb_writebatch_iterate_ext(leveldb_writebatch_t* w, void *p) {
+    leveldb_writebatch_iterate(w, p, 
+        leveldb_writebatch_iterate_put, leveldb_writebatch_iterate_delete);
+}
 
 }
\ No newline at end of file
diff --git a/store/leveldb/leveldb_ext.h b/store/leveldb/leveldb_ext.h
index 1c5f986..3eed41b 100644
--- a/store/leveldb/leveldb_ext.h
+++ b/store/leveldb/leveldb_ext.h
@@ -32,6 +32,7 @@ extern unsigned char leveldb_iter_seek_ext(leveldb_iterator_t*, const char* k, s
 extern unsigned char leveldb_iter_next_ext(leveldb_iterator_t*);
 extern unsigned char leveldb_iter_prev_ext(leveldb_iterator_t*);
 
+extern void leveldb_writebatch_iterate_ext(leveldb_writebatch_t*, void* p);
 
 #ifdef __cplusplus
 }