[RocksDB] Allow for WriteBatch buffer reuse after commit

crates/core/src/network/io/reactor.rs

@@ -34,7 +34,9 @@ use restate_types::{Version, Versioned};
 
 use crate::network::incoming::{RawRpc, RawUnary, RpcReplyPort};
 use crate::network::io::EgressMessage;
-use crate::network::metric_definitions::NETWORK_MESSAGE_RECEIVED_BYTES;
+use crate::network::metric_definitions::{
+    NETWORK_MESSAGE_RECEIVED_BYTES, NETWORK_MESSAGE_RECEIVED_DROPPED_BYTES,
+};
 use crate::network::protobuf::network::message::{Body, Signal};
 use crate::network::protobuf::network::{Datagram, RpcReply, datagram, rpc_reply};
 use crate::network::protobuf::network::{Header, Message};
@@ -332,7 +334,7 @@ impl ConnectionReactor {
                 };
                 let target_service = rpc_call.service();
 
-                let encoded_len = rpc_call.payload.len();
+                let encoded_len = rpc_call.payload.len() as u64;
                 let (reply_port, reply_rx) = RpcReplyPort::new();
                 let raw_rpc = RawRpc {
                     reply_port,
@@ -352,22 +354,21 @@ impl ConnectionReactor {
                     "Received RPC call: {target_service}::{}",
                     incoming.msg_type()
                 );
+
                 // ship to the service router, dropping the reply port will close the responder
                 // task.
                 match tokio::task::unconstrained(self.router.call_rpc(target_service, incoming))
                     .await
                 {
-                    Ok(()) => { /* spawn reply task */ }
+                    Ok(()) => {
+                        counter!(NETWORK_MESSAGE_RECEIVED_BYTES, "target" => target_service.as_str_name()).increment(encoded_len);
+                        spawn_rpc_responder(tx.clone(), rpc_call.id, reply_rx, target_service);
+                    }
                     Err(err) => {
                         send_rpc_error(tx, err, rpc_call.id);
+                        counter!(NETWORK_MESSAGE_RECEIVED_DROPPED_BYTES, "target" => target_service.as_str_name()).increment(encoded_len);
                     }
                 }
-
-                counter!(NETWORK_MESSAGE_RECEIVED_BYTES, "target" => target_service.as_str_name())
-                    .increment(encoded_len as u64);
-
-                spawn_rpc_responder(tx.clone(), rpc_call.id, reply_rx, target_service);
-
                 Decision::Continue
             }
             // UNARY MESSAGE
@@ -376,7 +377,7 @@ impl ConnectionReactor {
             }) => {
                 let metadata_versions = PeerMetadataVersion::from(header);
                 let target = unary.service();
-                let encoded_len = unary.payload.len();
+                let encoded_len = unary.payload.len() as u64;
                 let incoming = Incoming::new(
                     self.connection.protocol_version,
                     RawUnary {
@@ -392,7 +393,7 @@ impl ConnectionReactor {
                 let _ = tokio::task::unconstrained(self.router.call_unary(target, incoming)).await;
 
                 counter!(NETWORK_MESSAGE_RECEIVED_BYTES, "target" => target.as_str_name())
-                    .increment(encoded_len as u64);
+                    .increment(encoded_len);
                 Decision::Continue
             }
             // RPC REPLY

crates/core/src/network/metric_definitions.rs

@@ -13,6 +13,8 @@ use metrics::{Unit, describe_counter, describe_histogram};
 pub const NETWORK_CONNECTION_CREATED: &str = "restate.network.connection_created.total";
 pub const NETWORK_CONNECTION_DROPPED: &str = "restate.network.connection_dropped.total";
 pub const NETWORK_MESSAGE_RECEIVED_BYTES: &str = "restate.network.message_received_bytes.total";
+pub const NETWORK_MESSAGE_RECEIVED_DROPPED_BYTES: &str =
+    "restate.network.message_received_dropped_bytes.total";
 
 pub const NETWORK_MESSAGE_PROCESSING_DURATION: &str =
     "restate.network.message_processing_duration.seconds";
@@ -31,7 +33,13 @@ pub fn describe_metrics() {
     describe_counter!(
         NETWORK_MESSAGE_RECEIVED_BYTES,
         Unit::Bytes,
-        "Number of bytes received by message name"
+        "Number of bytes received by service name"
+    );
+
+    describe_counter!(
+        NETWORK_MESSAGE_RECEIVED_DROPPED_BYTES,
+        Unit::Bytes,
+        "Number of bytes received and dropped/rejected by service name"
     );
 
     describe_histogram!(

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();