Use ingestion-client in the Shuffler

Avoid direct writes to bifrost in shuffler by using a
dedicated ingestion-client instance.

Author: muhamadazmy

crates/ingestion-client/src/chunks_size.rs

@@ -67,6 +67,7 @@ where
                     let item_size = (me.size_fn)(&item);
 
                     if me.items.is_empty() || *me.size + item_size <= *me.cap {
+                        *me.size += item_size;
                         me.items.push(item);
                     } else {
                         // not empty and adding the item will go over the cap

crates/types/src/config/worker.rs

@@ -20,6 +20,7 @@ use restate_serde_util::NonZeroByteCount;
 use restate_time_util::{FriendlyDuration, NonZeroFriendlyDuration};
 
 use super::{CommonOptions, ObjectStoreOptions, RocksDbOptions, RocksDbOptionsBuilder};
+use crate::config::IngestionOptions;
 use crate::identifiers::PartitionId;
 use crate::rate::Rate;
 use crate::retries::RetryPolicy;
@@ -96,6 +97,12 @@ pub struct WorkerOptions {
     #[cfg_attr(feature = "schemars", schemars(skip))]
     #[serde(default, skip_serializing_if = "FriendlyDuration::is_zero")]
     trim_delay_interval: FriendlyDuration,
+
+    /// # Worker Ingestion Options
+    ///
+    /// Settings for the shared ingestion client used by all workers to
+    /// manage record ingestion across partitions.
+    pub ingestion_options: IngestionOptions,
 }
 
 impl WorkerOptions {
@@ -132,6 +139,7 @@ impl Default for WorkerOptions {
             snapshots: SnapshotsOptions::default(),
             trim_delay_interval: FriendlyDuration::ZERO,
             durability_mode: None,
+            ingestion_options: Default::default(),
         }
     }
 }

crates/worker/src/lib.rs

@@ -22,6 +22,9 @@ mod subscription_integration;
 use std::sync::Arc;
 
 use codederror::CodedError;
+use restate_core::network::Swimlane;
+use restate_ingestion_client::SessionOptions;
+use restate_wal_protocol::Envelope;
 use tracing::info;
 
 use restate_bifrost::Bifrost;
@@ -52,7 +55,6 @@ use restate_types::health::HealthStatus;
 use restate_types::partitions::state::PartitionReplicaSetStates;
 use restate_types::protobuf::common::WorkerStatus;
 use restate_types::schema::subscriptions::SubscriptionResolver;
-use restate_wal_protocol::Envelope;
 
 use crate::partition::invoker_storage_reader::InvokerStorageReader;
 use crate::partition_processor_manager::PartitionProcessorManager;
@@ -98,7 +100,7 @@ pub struct Worker<T> {
     datafusion_remote_scanner: RemoteQueryScannerServer,
     ingress_kafka: IngressKafkaService<T>,
     subscription_controller_handle: SubscriptionControllerHandle,
-    partition_processor_manager: PartitionProcessorManager,
+    partition_processor_manager: PartitionProcessorManager<T>,
 }
 
 impl<T> Worker<T>
@@ -142,6 +144,32 @@ where
             )));
         }
 
+        // A dedicated ingestion client for PPM that uses
+        // BifrostData swimlane
+        let ppm_ingestion_client = IngestionClient::new(
+            networking.clone(),
+            Metadata::with_current(|m| m.updateable_partition_table()),
+            partition_routing.clone(),
+            config
+                .worker
+                .ingestion_options
+                .inflight_memory_budget
+                .as_non_zero_usize(),
+            Some(SessionOptions {
+                batch_size: config
+                    .worker
+                    .ingestion_options
+                    .request_batch_size
+                    .as_usize(),
+                connection_retry_policy: config
+                    .worker
+                    .ingestion_options
+                    .connection_retry_policy
+                    .clone(),
+                swimlane: Swimlane::BifrostData,
+            }),
+        );
+
         let partition_processor_manager = PartitionProcessorManager::new(
             health_status,
             Configuration::live(),
@@ -156,6 +184,7 @@ where
             )
             .await
             .map_err(BuildError::SnapshotRepository)?,
+            ppm_ingestion_client,
         );
 
         let remote_scanner_manager = RemoteScannerManager::new(

crates/worker/src/partition/leadership/mod.rs

@@ -25,9 +25,10 @@ use tokio_stream::wrappers::ReceiverStream;
 use tracing::{debug, instrument, warn};
 
 use restate_bifrost::Bifrost;
-use restate_core::network::{Oneshot, Reciprocal};
+use restate_core::network::{Oneshot, Reciprocal, TransportConnect};
 use restate_core::{ShutdownError, TaskCenter, TaskKind, my_node_id};
 use restate_errors::NotRunningError;
+use restate_ingestion_client::IngestionClient;
 use restate_invoker_api::InvokeInputJournal;
 use restate_invoker_api::capacity::InvokerCapacity;
 use restate_partition_store::PartitionStore;
@@ -59,9 +60,9 @@ use restate_types::retries::with_jitter;
 use restate_types::schema::Schema;
 use restate_types::storage::{StorageDecodeError, StorageEncodeError};
 use restate_vqueues::{SchedulerService, VQueuesMeta, VQueuesMetaMut};
-use restate_wal_protocol::Command;
 use restate_wal_protocol::control::{AnnounceLeader, PartitionDurability};
 use restate_wal_protocol::timer::TimerKeyValue;
+use restate_wal_protocol::{Command, Envelope};
 
 use crate::partition::cleaner::{self, Cleaner};
 use crate::partition::invoker_storage_reader::InvokerStorageReader;
@@ -156,26 +157,29 @@ impl State {
     }
 }
 
-pub(crate) struct LeadershipState<I> {
+pub(crate) struct LeadershipState<T, I> {
     state: State,
     last_seen_leader_epoch: Option<LeaderEpoch>,
 
     partition: Arc<Partition>,
     invoker_tx: I,
+    ingestion_client: IngestionClient<T, Envelope>,
     invoker_capacity: InvokerCapacity,
     bifrost: Bifrost,
     trim_queue: TrimQueue,
 }
 
-impl<I> LeadershipState<I>
+impl<T, I> LeadershipState<T, I>
 where
     I: restate_invoker_api::InvokerHandle<InvokerStorageReader<PartitionStore>>,
+    T: TransportConnect,
 {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         partition: Arc<Partition>,
         invoker_tx: I,
         invoker_capacity: InvokerCapacity,
+        ingestion_client: IngestionClient<T, Envelope>,
         bifrost: Bifrost,
         last_seen_leader_epoch: Option<LeaderEpoch>,
         trim_queue: TrimQueue,
@@ -184,6 +188,7 @@ where
             state: State::Follower,
             partition,
             invoker_tx,
+            ingestion_client,
             invoker_capacity,
             bifrost,
             last_seen_leader_epoch,
@@ -411,7 +416,7 @@ where
                 OutboxReader::from(partition_store.clone()),
                 shuffle_tx,
                 config.worker.internal_queue_length(),
-                self.bifrost.clone(),
+                self.ingestion_client.clone(),
             );
 
             let shuffle_hint_tx = shuffle.create_hint_sender();
@@ -598,7 +603,7 @@ where
     }
 }
 
-impl<I> LeadershipState<I> {
+impl<T, I> LeadershipState<T, I> {
     pub async fn handle_rpc_proposal_command(
         &mut self,
         request_id: PartitionProcessorRpcRequestId,
@@ -718,7 +723,9 @@ mod tests {
     use crate::partition::leadership::{LeadershipState, State};
     use assert2::let_assert;
     use restate_bifrost::Bifrost;
+    use restate_core::partitions::PartitionRouting;
     use restate_core::{TaskCenter, TestCoreEnv};
+    use restate_ingestion_client::IngestionClient;
     use restate_invoker_api::capacity::InvokerCapacity;
     use restate_invoker_api::test_util::MockInvokerHandle;
     use restate_partition_store::PartitionStoreManager;
@@ -732,6 +739,7 @@ mod tests {
     use restate_vqueues::VQueuesMetaMut;
     use restate_wal_protocol::control::AnnounceLeader;
     use restate_wal_protocol::{Command, Envelope};
+    use std::num::NonZeroUsize;
     use std::ops::RangeInclusive;
     use std::sync::Arc;
     use test_log::test;
@@ -752,11 +760,20 @@ mod tests {
 
         let partition_store_manager = PartitionStoreManager::create().await?;
 
+        let ingress = IngestionClient::new(
+            env.networking.clone(),
+            env.metadata.updateable_partition_table(),
+            PartitionRouting::new(replica_set_states.clone(), TaskCenter::current()),
+            NonZeroUsize::new(10 * 1024 * 1024).unwrap(),
+            None,
+        );
+
         let invoker_tx = MockInvokerHandle::default();
         let mut state = LeadershipState::new(
             Arc::new(PARTITION),
             invoker_tx,
             InvokerCapacity::new_unlimited(),
+            ingress,
             bifrost.clone(),
             None,
             TrimQueue::default(),

crates/worker/src/partition/mod.rs

@@ -31,8 +31,11 @@ use tracing::{Span, debug, error, info, instrument, trace, warn};
 
 use restate_bifrost::loglet::FindTailOptions;
 use restate_bifrost::{Bifrost, LogEntry, MaybeRecord};
-use restate_core::network::{Incoming, Oneshot, Reciprocal, Rpc, ServiceMessage, Verdict};
+use restate_core::network::{
+    Incoming, Oneshot, Reciprocal, Rpc, ServiceMessage, TransportConnect, Verdict,
+};
 use restate_core::{Metadata, ShutdownError, cancellation_watcher, my_node_id};
+use restate_ingestion_client::IngestionClient;
 use restate_invoker_api::capacity::InvokerCapacity;
 use restate_partition_store::{PartitionStore, PartitionStoreTransaction};
 use restate_storage_api::deduplication_table::{
@@ -115,12 +118,16 @@ where
         }
     }
 
-    pub async fn build(
+    pub async fn build<T>(
         self,
         bifrost: Bifrost,
+        ingestion_client: IngestionClient<T, Envelope>,
         mut partition_store: PartitionStore,
         replica_set_states: PartitionReplicaSetStates,
-    ) -> Result<PartitionProcessor<InvokerInputSender>, state_machine::Error> {
+    ) -> Result<PartitionProcessor<T, InvokerInputSender>, state_machine::Error>
+    where
+        T: TransportConnect,
+    {
         let PartitionProcessorBuilder {
             invoker_tx,
             target_leader_state_rx,
@@ -165,6 +172,7 @@ where
             Arc::clone(partition_store.partition()),
             invoker_tx,
             invoker_capacity,
+            ingestion_client,
             bifrost.clone(),
             last_seen_leader_epoch,
             trim_queue.clone(),
@@ -218,9 +226,9 @@ where
     }
 }
 
-pub struct PartitionProcessor<InvokerSender> {
+pub struct PartitionProcessor<T, InvokerSender> {
     partition_id_str: SharedString,
-    leadership_state: LeadershipState<InvokerSender>,
+    leadership_state: LeadershipState<T, InvokerSender>,
     state_machine: StateMachine,
     bifrost: Bifrost,
     target_leader_state_rx: watch::Receiver<TargetLeaderState>,
@@ -284,8 +292,9 @@ struct LsnEnvelope {
     pub envelope: Arc<Envelope>,
 }
 
-impl<InvokerSender> PartitionProcessor<InvokerSender>
+impl<T, InvokerSender> PartitionProcessor<T, InvokerSender>
 where
+    T: TransportConnect,
     InvokerSender: restate_invoker_api::InvokerHandle<InvokerStorageReader<PartitionStore>> + Clone,
 {
     #[instrument(

crates/worker/src/partition/rpc/mod.rs

@@ -22,7 +22,7 @@ mod resume_invocation;
 
 use crate::partition;
 use crate::partition::leadership::LeadershipState;
-use restate_core::network::{Oneshot, Reciprocal};
+use restate_core::network::{Oneshot, Reciprocal, TransportConnect};
 use restate_invoker_api::InvokerHandle;
 use restate_storage_api::idempotency_table::ReadOnlyIdempotencyTable;
 use restate_storage_api::invocation_status_table::ReadInvocationStatusTable;
@@ -71,13 +71,14 @@ pub(super) trait Actuator {
     );
 }
 
-impl<
+impl<T, I> Actuator for LeadershipState<T, I>
+where
+    T: TransportConnect,
     I: InvokerHandle<
         partition::invoker_storage_reader::InvokerStorageReader<
             restate_partition_store::PartitionStore,
         >,
     >,
-> Actuator for LeadershipState<I>
 {
     async fn self_propose_and_respond_asynchronously<O: Into<PartitionProcessorRpcResponse>>(
         &mut self,