[RocksDB] Allow for WriteBatch buffer reuse after commit

Author: AhmedSoliman

crates/metadata-server/src/raft/storage/rocksdb.rs

@@ -327,6 +327,7 @@ impl RocksDbStorage {
                 write_batch,
             )
             .await
+            .map(|_| ())
             .map_err(Into::into)
     }
 

crates/partition-store/src/partition_store.rs

@@ -971,6 +971,7 @@ impl Transaction for PartitionStoreTransaction<'_> {
                 self.write_batch_with_index,
             )
             .await
+            .map(|_| ())
             .map_err(|error| StorageError::Generic(error.into()))
     }
 }

crates/rocksdb/src/lib.rs

@@ -155,6 +155,11 @@ impl RocksDb {
         self.db.cfs()
     }
 
+    /// Write a [`rocksdb::WriteBatch`] to the database.
+    ///
+    /// The batch is consumed and returned on success, allowing callers to reuse
+    /// it (e.g. by calling [`rocksdb::WriteBatch::clear()`]) without
+    /// re-allocating. On error the batch is lost.
     #[tracing::instrument(skip_all, fields(db = %self.name()))]
     pub async fn write_batch(
         self: &Arc<Self>,
@@ -163,17 +168,22 @@ impl RocksDb {
         io_mode: IoMode,
         write_options: rocksdb::WriteOptions,
         write_batch: rocksdb::WriteBatch,
-    ) -> Result<(), RocksError> {
+    ) -> Result<rocksdb::WriteBatch, RocksError> {
         self.write_batch_internal(
             name,
             priority,
             io_mode,
             write_options,
-            move |db, write_options| db.write_batch(&write_batch, write_options),
+            write_batch,
+            |db, write_options, batch| db.write_batch(batch, write_options),
         )
         .await
     }
 
+    /// Write a [`rocksdb::WriteBatchWithIndex`] to the database.
+    ///
+    /// The batch is consumed and returned on success, allowing callers to reuse
+    /// it without re-allocating. On error the batch is lost.
     #[tracing::instrument(skip_all, fields(db = %self.name()))]
     pub async fn write_batch_with_index(
         self: &Arc<Self>,
@@ -182,27 +192,32 @@ impl RocksDb {
         io_mode: IoMode,
         write_options: rocksdb::WriteOptions,
         write_batch: rocksdb::WriteBatchWithIndex,
-    ) -> Result<(), RocksError> {
+    ) -> Result<rocksdb::WriteBatchWithIndex, RocksError> {
         self.write_batch_internal(
             name,
             priority,
             io_mode,
             write_options,
-            move |db, write_options| db.write_batch_with_index(&write_batch, write_options),
+            write_batch,
+            |db, write_options, batch| db.write_batch_with_index(batch, write_options),
         )
         .await
     }
 
-    async fn write_batch_internal<OP>(
+    async fn write_batch_internal<B, OP>(
         self: &Arc<Self>,
         name: &'static str,
         priority: Priority,
         io_mode: IoMode,
         mut write_options: rocksdb::WriteOptions,
+        batch: B,
         write_op: OP,
-    ) -> Result<(), RocksError>
+    ) -> Result<B, RocksError>
     where
-        OP: Fn(&RocksAccess, &rocksdb::WriteOptions) -> Result<(), rocksdb::Error> + Send + 'static,
+        B: Send + 'static,
+        OP: Fn(&RocksAccess, &rocksdb::WriteOptions, &B) -> Result<(), rocksdb::Error>
+            + Send
+            + 'static,
     {
         //  depending on the IoMode, we decide how to do the write.
         match io_mode {
@@ -212,14 +227,14 @@ impl RocksDb {
                     "Blocking IO is allowed for write_batch, stall detection will not be used in this operation!"
                 );
                 write_options.set_no_slowdown(false);
-                write_op(&self.db, &write_options)?;
+                write_op(&self.db, &write_options, &batch)?;
                 counter!(STORAGE_IO_OP,
                     DISPOSITION => DISPOSITION_MAYBE_BLOCKING,
                     OP_TYPE => StorageTaskKind::WriteBatch.as_static_str(),
                     PRIORITY => priority.as_static_str(),
                 )
                 .increment(1);
-                return Ok(());
+                return Ok(batch);
             }
             IoMode::AlwaysBackground => {
                 // Operation will block, dispatch to background.
@@ -229,9 +244,10 @@ impl RocksDb {
                 let task = StorageTask::default()
                     .priority(priority)
                     .kind(StorageTaskKind::WriteBatch)
-                    .op(move || {
+                    .op(move || -> Result<B, rocksdb::Error> {
                         let _x = RocksDbPerfGuard::new(name);
-                        write_op(&db.db, &write_options)
+                        write_op(&db.db, &write_options, &batch)?;
+                        Ok(batch)
                     })
                     .build()
                     .unwrap();
@@ -248,14 +264,14 @@ impl RocksDb {
             IoMode::OnlyIfNonBlocking => {
                 let _x = RocksDbPerfGuard::new(name);
                 write_options.set_no_slowdown(true);
-                write_op(&self.db, &write_options)?;
+                write_op(&self.db, &write_options, &batch)?;
                 counter!(STORAGE_IO_OP,
                     DISPOSITION => DISPOSITION_NON_BLOCKING,
                     OP_TYPE => StorageTaskKind::WriteBatch.as_static_str(),
                     PRIORITY => priority.as_static_str(),
                 )
                 .increment(1);
-                return Ok(());
+                return Ok(batch);
             }
             _ => {}
         }
@@ -265,7 +281,7 @@ impl RocksDb {
         write_options.set_no_slowdown(true);
 
         let perf_guard = RocksDbPerfGuard::new(name);
-        let result = write_op(&self.db, &write_options);
+        let result = write_op(&self.db, &write_options, &batch);
         match result {
             Ok(_) => {
                 counter!(STORAGE_IO_OP,
@@ -274,7 +290,7 @@ impl RocksDb {
                     PRIORITY => priority.as_static_str(),
                 )
                 .increment(1);
-                Ok(())
+                Ok(batch)
             }
             Err(e) if is_retryable_error(e.kind()) => {
                 counter!(STORAGE_IO_OP,
@@ -294,9 +310,10 @@ impl RocksDb {
                 let task = StorageTask::default()
                     .priority(priority)
                     .kind(StorageTaskKind::WriteBatch)
-                    .op(move || {
+                    .op(move || -> Result<B, rocksdb::Error> {
                         let _x = RocksDbPerfGuard::new(name);
-                        write_op(&db.db, &write_options)
+                        write_op(&db.db, &write_options, &batch)?;
+                        Ok(batch)
                     })
                     .build()
                     .unwrap();