implement with_capacity for WriteBatch (#947)

* implement `with_capacity` for `WriteBatch`

* add tests for WriteBatch::with_capacity

* update changelog
 - remove `clippy.toml` because we already have `msrv` in `Cargo.toml`

Author: 0xdeafbeef

clippy.toml

@@ -89,6 +89,22 @@ unsafe extern "C" fn writebatch_delete_callback(state: *mut c_void, k: *const c_
 }
 
 impl<const TRANSACTION: bool> WriteBatchWithTransaction<TRANSACTION> {
+    /// Create a new `WriteBatch` without allocating memory.
+    pub fn new() -> Self {
+        Self {
+            inner: unsafe { ffi::rocksdb_writebatch_create() },
+        }
+    }
+
+    /// Creates `WriteBatch` with the specified `capacity` in bytes. Allocates immediately.
+    pub fn with_capacity_bytes(capacity_bytes: usize) -> Self {
+        Self {
+            // zeroes from default constructor
+            // https://github.com/facebook/rocksdb/blob/0f35db55d86ea8699ea936c9e2a4e34c82458d6b/include/rocksdb/write_batch.h#L66
+            inner: unsafe { ffi::rocksdb_writebatch_create_with_params(capacity_bytes, 0, 0, 0) },
+        }
+    }
+
     /// Construct with a reference to a byte array serialized by [`WriteBatch`].
     pub fn from_data(data: &[u8]) -> Self {
         unsafe {
@@ -354,9 +370,7 @@ impl WriteBatchWithTransaction<false> {
 
 impl<const TRANSACTION: bool> Default for WriteBatchWithTransaction<TRANSACTION> {
     fn default() -> Self {
-        Self {
-            inner: unsafe { ffi::rocksdb_writebatch_create() },
-        }
+        Self::new()
     }
 }
 

src/write_batch.rs

@@ -535,29 +535,31 @@ fn test_get_updates_since_multiple_batches() {
 
 #[test]
 fn test_get_updates_since_one_batch() {
-    let path = DBPath::new("_rust_rocksdb_test_get_updates_since_one_batch");
-    let db = DB::open_default(&path).unwrap();
-    db.put(b"key2", b"value2").unwrap();
-    // some puts and deletes in a single batch,
-    // verify 1 put and 1 delete were done
-    let seq1 = db.latest_sequence_number();
-    assert_eq!(seq1, 1);
-    let mut batch = WriteBatch::default();
-    batch.put(b"key1", b"value1");
-    batch.delete(b"key2");
-    db.write(batch).unwrap();
-    assert_eq!(db.latest_sequence_number(), 3);
-    let mut iter = db.get_updates_since(seq1).unwrap();
-    let mut counts = OperationCounts {
-        puts: 0,
-        deletes: 0,
-    };
-    let (seq, batch) = iter.next().unwrap().unwrap();
-    assert_eq!(seq, 2);
-    batch.iterate(&mut counts);
-    assert!(iter.next().is_none());
-    assert_eq!(counts.puts, 1);
-    assert_eq!(counts.deletes, 1);
+    let batches = [WriteBatch::default(), WriteBatch::with_capacity_bytes(13)];
+    for mut batch in batches {
+        let path = DBPath::new("_rust_rocksdb_test_get_updates_since_one_batch");
+        let db = DB::open_default(&path).unwrap();
+        db.put(b"key2", b"value2").unwrap();
+        // some puts and deletes in a single batch,
+        // verify 1 put and 1 delete were done
+        let seq1 = db.latest_sequence_number();
+        assert_eq!(seq1, 1);
+        batch.put(b"key1", b"value1");
+        batch.delete(b"key2");
+        db.write(batch).unwrap();
+        assert_eq!(db.latest_sequence_number(), 3);
+        let mut iter = db.get_updates_since(seq1).unwrap();
+        let mut counts = OperationCounts {
+            puts: 0,
+            deletes: 0,
+        };
+        let (seq, batch) = iter.next().unwrap().unwrap();
+        assert_eq!(seq, 2);
+        batch.iterate(&mut counts);
+        assert!(iter.next().is_none());
+        assert_eq!(counts.puts, 1);
+        assert_eq!(counts.deletes, 1);
+    }
 }
 
 #[test]
@@ -953,134 +955,140 @@ fn prefix_extract_and_iterate_test() {
 
 #[test]
 fn get_with_cache_and_bulkload_test() {
-    let path = DBPath::new("_rust_rocksdb_get_with_cache_and_bulkload_test");
-    let log_path = DBPath::new("_rust_rocksdb_log_path_test");
-
-    // create options
-    let mut opts = Options::default();
-    opts.create_if_missing(true);
-    opts.create_missing_column_families(true);
-    opts.set_wal_bytes_per_sync(8 << 10); // 8KB
-    opts.set_writable_file_max_buffer_size(512 << 10); // 512KB
-    opts.set_enable_write_thread_adaptive_yield(true);
-    opts.set_unordered_write(true);
-    opts.set_max_subcompactions(2);
-    opts.set_max_background_jobs(4);
-    opts.set_use_adaptive_mutex(true);
-    opts.set_db_log_dir(&log_path);
-    opts.set_memtable_whole_key_filtering(true);
-    opts.set_dump_malloc_stats(true);
-    opts.set_level_compaction_dynamic_level_bytes(false);
-
-    // trigger all sst files in L1/2 instead of L0
-    opts.set_max_bytes_for_level_base(64 << 10); // 64KB
-
-    {
-        // set block based table and cache
-        let cache = Cache::new_lru_cache(512 << 10);
-        assert_eq!(cache.get_usage(), 0);
-        let mut block_based_opts = BlockBasedOptions::default();
-        block_based_opts.set_block_cache(&cache);
-        block_based_opts.set_cache_index_and_filter_blocks(true);
-        opts.set_block_based_table_factory(&block_based_opts);
-
-        // open db
-        let db = DB::open(&opts, &path).unwrap();
-
-        // write a lot
-        let mut batch = WriteBatch::default();
-        for i in 0..10_000 {
-            batch.put(format!("{i:0>4}").as_bytes(), b"v");
-        }
-        assert!(db.write(batch).is_ok());
+    let batches = [
+        WriteBatch::default(),
+        WriteBatch::with_capacity_bytes(13),
+        WriteBatch::with_capacity_bytes(100_000),
+    ];
+    for mut batch in batches {
+        let path = DBPath::new("_rust_rocksdb_get_with_cache_and_bulkload_test");
+        let log_path = DBPath::new("_rust_rocksdb_log_path_test");
+
+        // create options
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+        opts.set_wal_bytes_per_sync(8 << 10); // 8KB
+        opts.set_writable_file_max_buffer_size(512 << 10); // 512KB
+        opts.set_enable_write_thread_adaptive_yield(true);
+        opts.set_unordered_write(true);
+        opts.set_max_subcompactions(2);
+        opts.set_max_background_jobs(4);
+        opts.set_use_adaptive_mutex(true);
+        opts.set_db_log_dir(&log_path);
+        opts.set_memtable_whole_key_filtering(true);
+        opts.set_dump_malloc_stats(true);
+        opts.set_level_compaction_dynamic_level_bytes(false);
 
-        // flush memory table to sst and manual compaction
-        assert!(db.flush().is_ok());
-        db.compact_range(Some(format!("{:0>4}", 0).as_bytes()), None::<Vec<u8>>);
+        // trigger all sst files in L1/2 instead of L0
+        opts.set_max_bytes_for_level_base(64 << 10); // 64KB
 
-        // get -> trigger caching
-        let _ = db.get(b"1");
-        assert!(cache.get_usage() > 0);
+        {
+            // set block based table and cache
+            let cache = Cache::new_lru_cache(512 << 10);
+            assert_eq!(cache.get_usage(), 0);
+            let mut block_based_opts = BlockBasedOptions::default();
+            block_based_opts.set_block_cache(&cache);
+            block_based_opts.set_cache_index_and_filter_blocks(true);
+            opts.set_block_based_table_factory(&block_based_opts);
+
+            // open db
+            let db = DB::open(&opts, &path).unwrap();
+
+            // write a lot
+            for i in 0..10_000 {
+                batch.put(format!("{i:0>4}").as_bytes(), b"v");
+            }
+            assert!(db.write(batch).is_ok());
 
-        // get approximated memory usage
-        let mem_usage = get_memory_usage_stats(Some(&[&db]), None).unwrap();
-        assert!(mem_usage.mem_table_total > 0);
+            // flush memory table to sst and manual compaction
+            assert!(db.flush().is_ok());
+            db.compact_range(Some(format!("{:0>4}", 0).as_bytes()), None::<Vec<u8>>);
 
-        // get approximated cache usage
-        let mem_usage = get_memory_usage_stats(None, Some(&[&cache])).unwrap();
-        assert!(mem_usage.cache_total > 0);
-    }
+            // get -> trigger caching
+            let _ = db.get(b"1");
+            assert!(cache.get_usage() > 0);
 
-    // bulk loading
-    {
-        // open db
-        let db = DB::open(&opts, &path).unwrap();
+            // get approximated memory usage
+            let mem_usage = get_memory_usage_stats(Some(&[&db]), None).unwrap();
+            assert!(mem_usage.mem_table_total > 0);
 
-        // try to get key
-        let iter = db.iterator(IteratorMode::Start);
-        for (expected, (k, _)) in iter.map(Result::unwrap).enumerate() {
-            assert_eq!(k.as_ref(), format!("{expected:0>4}").as_bytes());
+            // get approximated cache usage
+            let mem_usage = get_memory_usage_stats(None, Some(&[&cache])).unwrap();
+            assert!(mem_usage.cache_total > 0);
         }
 
-        // check live files (sst files meta)
-        let livefiles = db.live_files().unwrap();
-        assert_eq!(livefiles.len(), 1);
-        livefiles.iter().for_each(|f| {
-            assert_eq!(f.level, 2);
-            assert_eq!(f.column_family_name, "default");
-            assert!(!f.name.is_empty());
-            assert_eq!(
-                f.start_key.as_ref().unwrap().as_slice(),
-                format!("{:0>4}", 0).as_bytes()
-            );
-            assert_eq!(
-                f.end_key.as_ref().unwrap().as_slice(),
-                format!("{:0>4}", 9999).as_bytes()
-            );
-            assert_eq!(f.num_entries, 10000);
-            assert_eq!(f.num_deletions, 0);
-        });
+        // bulk loading
+        {
+            // open db
+            let db = DB::open(&opts, &path).unwrap();
 
-        // delete sst file in range (except L0)
-        assert!(db
-            .delete_file_in_range(
-                format!("{:0>4}", 0).as_bytes(),
-                format!("{:0>4}", 9999).as_bytes()
-            )
-            .is_ok());
-        let livefiles = db.live_files().unwrap();
-        assert_eq!(livefiles.len(), 0);
+            // try to get key
+            let iter = db.iterator(IteratorMode::Start);
+            for (expected, (k, _)) in iter.map(Result::unwrap).enumerate() {
+                assert_eq!(k.as_ref(), format!("{expected:0>4}").as_bytes());
+            }
 
-        // try to get a deleted key
-        assert!(db.get(format!("{:0>4}", 123).as_bytes()).unwrap().is_none());
-    }
+            // check live files (sst files meta)
+            let livefiles = db.live_files().unwrap();
+            assert_eq!(livefiles.len(), 1);
+            livefiles.iter().for_each(|f| {
+                assert_eq!(f.level, 2);
+                assert_eq!(f.column_family_name, "default");
+                assert!(!f.name.is_empty());
+                assert_eq!(
+                    f.start_key.as_ref().unwrap().as_slice(),
+                    format!("{:0>4}", 0).as_bytes()
+                );
+                assert_eq!(
+                    f.end_key.as_ref().unwrap().as_slice(),
+                    format!("{:0>4}", 9999).as_bytes()
+                );
+                assert_eq!(f.num_entries, 10000);
+                assert_eq!(f.num_deletions, 0);
+            });
+
+            // delete sst file in range (except L0)
+            assert!(db
+                .delete_file_in_range(
+                    format!("{:0>4}", 0).as_bytes(),
+                    format!("{:0>4}", 9999).as_bytes()
+                )
+                .is_ok());
+            let livefiles = db.live_files().unwrap();
+            assert_eq!(livefiles.len(), 0);
+
+            // try to get a deleted key
+            assert!(db.get(format!("{:0>4}", 123).as_bytes()).unwrap().is_none());
+        }
 
-    // raise error when db exists
-    {
-        opts.set_error_if_exists(true);
-        assert!(DB::open(&opts, &path).is_err());
-    }
+        // raise error when db exists
+        {
+            opts.set_error_if_exists(true);
+            assert!(DB::open(&opts, &path).is_err());
+        }
 
-    // disable all threads
-    {
-        // create new options
-        let mut opts = Options::default();
-        opts.set_max_background_jobs(0);
-        opts.set_stats_dump_period_sec(0);
-        opts.set_stats_persist_period_sec(0);
+        // disable all threads
+        {
+            // create new options
+            let mut opts = Options::default();
+            opts.set_max_background_jobs(0);
+            opts.set_stats_dump_period_sec(0);
+            opts.set_stats_persist_period_sec(0);
 
-        // test Env::Default()->SetBackgroundThreads(0, Env::Priority::BOTTOM);
-        let mut env = Env::new().unwrap();
-        env.set_bottom_priority_background_threads(0);
-        opts.set_env(&env);
+            // test Env::Default()->SetBackgroundThreads(0, Env::Priority::BOTTOM);
+            let mut env = Env::new().unwrap();
+            env.set_bottom_priority_background_threads(0);
+            opts.set_env(&env);
 
-        // open db
-        let db = DB::open(&opts, &path).unwrap();
+            // open db
+            let db = DB::open(&opts, &path).unwrap();
 
-        // try to get key
-        let iter = db.iterator(IteratorMode::Start);
-        for (expected, (k, _)) in iter.map(Result::unwrap).enumerate() {
-            assert_eq!(k.as_ref(), format!("{expected:0>4}").as_bytes());
+            // try to get key
+            let iter = db.iterator(IteratorMode::Start);
+            for (expected, (k, _)) in iter.map(Result::unwrap).enumerate() {
+                assert_eq!(k.as_ref(), format!("{expected:0>4}").as_bytes());
+            }
         }
     }
 }