writebatch with capacity (#58)

librocksdb_sys/crocksdb/c.cc

@@ -1145,6 +1145,12 @@ crocksdb_writebatch_t* crocksdb_writebatch_create() {
   return new crocksdb_writebatch_t;
 }
 
+crocksdb_writebatch_t* crocksdb_writebatch_create_with_capacity(size_t reserved_bytes) {
+  crocksdb_writebatch_t* b = new crocksdb_writebatch_t;
+  b->rep = WriteBatch(reserved_bytes);
+  return b;
+}
+
 crocksdb_writebatch_t* crocksdb_writebatch_create_from(const char* rep,
                                                      size_t size) {
   crocksdb_writebatch_t* b = new crocksdb_writebatch_t;

librocksdb_sys/crocksdb/rocksdb/c.h

@@ -404,6 +404,8 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_iter_get_error(
 /* Write batch */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_writebatch_t* crocksdb_writebatch_create();
+extern C_ROCKSDB_LIBRARY_API crocksdb_writebatch_t*
+    crocksdb_writebatch_create_with_capacity(size_t reserved_bytes);
 extern C_ROCKSDB_LIBRARY_API crocksdb_writebatch_t* crocksdb_writebatch_create_from(
     const char* rep, size_t size);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_writebatch_destroy(

librocksdb_sys/src/lib.rs

@@ -222,7 +222,8 @@ extern "C" {
     pub fn crocksdb_options_optimize_for_point_lookup(options: *mut DBOptions,
                                                       block_cache_size_mb: u64);
     pub fn crocksdb_options_set_table_cache_numshardbits(options: *mut DBOptions, bits: c_int);
-    pub fn crocksdb_options_set_writable_file_max_buffer_size(options: *mut DBOptions, nbytes: c_int);
+    pub fn crocksdb_options_set_writable_file_max_buffer_size(options: *mut DBOptions,
+                                                              nbytes: c_int);
     pub fn crocksdb_options_set_max_write_buffer_number(options: *mut DBOptions, bufno: c_int);
     pub fn crocksdb_options_set_min_write_buffer_number_to_merge(options: *mut DBOptions,
                                                                  bufno: c_int);
@@ -475,6 +476,7 @@ extern "C" {
                           batch: *mut DBWriteBatch,
                           err: *mut *mut c_char);
     pub fn crocksdb_writebatch_create() -> *mut DBWriteBatch;
+    pub fn crocksdb_writebatch_create_with_capacity(cap: size_t) -> *mut DBWriteBatch;
     pub fn crocksdb_writebatch_create_from(rep: *const u8, size: size_t) -> *mut DBWriteBatch;
     pub fn crocksdb_writebatch_destroy(batch: *mut DBWriteBatch);
     pub fn crocksdb_writebatch_clear(batch: *mut DBWriteBatch);

src/rocksdb.rs

@@ -1103,6 +1103,10 @@ impl WriteBatch {
         WriteBatch::default()
     }
 
+    pub fn with_capacity(cap: usize) -> WriteBatch {
+        WriteBatch { inner: unsafe { crocksdb_ffi::crocksdb_writebatch_create_with_capacity(cap) } }
+    }
+
     pub fn count(&self) -> usize {
         unsafe { crocksdb_ffi::crocksdb_writebatch_count(self.inner) as usize }
     }
@@ -1516,6 +1520,17 @@ mod test {
         assert!(r.unwrap().is_none());
         let r = db.get(b"k13");
         assert!(r.unwrap().is_none());
+
+        // test with capacity
+        let batch = WriteBatch::with_capacity(1024);
+        batch.put(b"kc1", b"v1").unwrap();
+        batch.put(b"kc2", b"v2").unwrap();
+        let p = db.write(batch);
+        assert!(p.is_ok());
+        let r = db.get(b"kc1");
+        assert!(r.unwrap().is_some());
+        let r = db.get(b"kc2");
+        assert!(r.unwrap().is_some());
     }
 
     #[test]