[wip] shuffler with ingress client

Author: muhamadazmy

crates/worker/src/lib.rs

@@ -20,6 +20,7 @@ mod subscription_controller;
 mod subscription_integration;
 
 use std::sync::Arc;
+use std::time::Duration;
 
 use codederror::CodedError;
 use tracing::info;
@@ -35,6 +36,7 @@ use restate_core::worker_api::ProcessorsManagerHandle;
 use restate_core::{Metadata, TaskKind};
 use restate_core::{MetadataWriter, TaskCenter};
 use restate_ingress_client::IngressClient;
+use restate_ingress_client::SessionOptions;
 use restate_ingress_kafka::Service as IngressKafkaService;
 use restate_invoker_impl::InvokerHandle as InvokerChannelServiceHandle;
 use restate_partition_store::snapshots::SnapshotRepository;
@@ -97,7 +99,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>
@@ -140,6 +142,20 @@ where
             )));
         }
 
+        // This instance of ingress client is dedicated for the PPM.
+        // so it has separate budget AND batch_timeout (set to zero)
+        // to optimize for latency.
+        let ppm_ingress = IngressClient::new(
+            networking.clone(),
+            Metadata::with_current(|m| m.updateable_partition_table()),
+            partition_routing.clone(),
+            1000,
+            Some(SessionOptions {
+                batch_timeout: Duration::ZERO,
+                ..Default::default()
+            }),
+        );
+
         let partition_processor_manager = PartitionProcessorManager::new(
             health_status,
             Configuration::live(),
@@ -154,6 +170,7 @@ where
             )
             .await
             .map_err(BuildError::SnapshotRepository)?,
+            ppm_ingress,
         );
 
         let remote_scanner_manager = RemoteScannerManager::new(

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

@@ -26,9 +26,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_ingress_client::IngressClient;
 use restate_invoker_api::InvokeInputJournal;
 use restate_partition_store::PartitionStore;
 use restate_storage_api::deduplication_table::EpochSequenceNumber;
@@ -151,25 +152,28 @@ 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,
+    ingress: IngressClient<T>,
     bifrost: Bifrost,
     #[allow(unused)]
     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,
+        ingress: IngressClient<T>,
         bifrost: Bifrost,
         last_seen_leader_epoch: Option<LeaderEpoch>,
         trim_queue: TrimQueue,
@@ -178,6 +182,7 @@ where
             state: State::Follower,
             partition,
             invoker_tx,
+            ingress,
             bifrost,
             last_seen_leader_epoch,
             trim_queue,
@@ -374,7 +379,7 @@ where
                 OutboxReader::from(partition_store.clone()),
                 shuffle_tx,
                 config.worker.internal_queue_length(),
-                self.bifrost.clone(),
+                self.ingress.clone(),
             );
 
             let shuffle_hint_tx = shuffle.create_hint_sender();
@@ -562,7 +567,7 @@ where
     }
 }
 
-impl<I> LeadershipState<I> {
+impl<T, I> LeadershipState<T, I> {
     pub async fn handle_rpc_proposal_command(
         &mut self,
         request_id: PartitionProcessorRpcRequestId,
@@ -682,7 +687,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_ingress_client::IngressClient;
     use restate_invoker_api::test_util::MockInvokerHandle;
     use restate_partition_store::PartitionStoreManager;
     use restate_rocksdb::RocksDbManager;
@@ -714,10 +721,19 @@ mod tests {
 
         let partition_store_manager = PartitionStoreManager::create().await?;
 
+        let ingress = IngressClient::new(
+            env.networking.clone(),
+            env.metadata.updateable_partition_table(),
+            PartitionRouting::new(replica_set_states.clone(), TaskCenter::current()),
+            1000,
+            None,
+        );
+
         let invoker_tx = MockInvokerHandle::default();
         let mut state = LeadershipState::new(
             Arc::new(PARTITION),
             invoker_tx,
+            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_ingress_client::IngressClient;
 use restate_partition_store::{PartitionStore, PartitionStoreTransaction};
 use restate_storage_api::deduplication_table::{
     DedupInformation, DedupSequenceNumber, ProducerId, ReadDeduplicationTable,
@@ -107,12 +110,16 @@ where
         }
     }
 
-    pub async fn build(
+    pub async fn build<T>(
         self,
         bifrost: Bifrost,
+        ingress: IngressClient<T>,
         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,
@@ -155,6 +162,7 @@ where
         let leadership_state = LeadershipState::new(
             Arc::clone(partition_store.partition()),
             invoker_tx,
+            ingress,
             bifrost.clone(),
             last_seen_leader_epoch,
             trim_queue.clone(),
@@ -208,9 +216,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>,
@@ -274,8 +282,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,