[worker] Funnel all command proposals through LeaderState::run

crates/worker/Cargo.toml

@@ -88,7 +88,6 @@ restate-vqueues = { workspace = true, features = ["test-util"] }
 restate-wal-protocol = { workspace = true, features = ["test-util"] }
 
 googletest = { workspace = true }
-mockall = { workspace = true }
 prost = { workspace = true }
 rstest = { workspace = true }
 test-log = { workspace = true }

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

@@ -49,7 +49,6 @@ use restate_types::cluster::cluster_state::RunMode;
 use restate_types::config::Configuration;
 use restate_types::errors::GenericError;
 use restate_types::identifiers::{InvocationId, LeaderEpoch, PartitionId};
-use restate_types::identifiers::{PartitionKey, PartitionProcessorRpcRequestId};
 use restate_types::invocation::FencingToken;
 use restate_types::live::LiveLoadExt;
 use restate_types::logs::Keys;
@@ -91,10 +90,14 @@ use crate::partition::state_machine::Action;
 use crate::partition::types::InvokerEffect;
 
 use super::node::NodeContext;
-use super::processor::*;
+use super::{processor::*, rpc};
 
 type TimerService = restate_timer::TimerService<TimerKeyValue, TokioClock, TimerReader>;
 type InvokerStream = ReceiverStream<InvokerEffect>;
+type RpcReciprocal =
+    Reciprocal<Oneshot<Result<PartitionProcessorRpcResponse, PartitionProcessorRpcError>>>;
+pub(super) type AppendCallback =
+    Box<dyn FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static>;
 
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum Error {
@@ -145,8 +148,11 @@ pub(crate) enum TaskTermination {
     Failure(GenericError),
 }
 
-#[derive(Debug)]
-pub(crate) enum ActionEffect {
+#[derive(derive_more::Debug)]
+// Keep proposals inline to avoid allocating on the network request path. At most one network
+// event can be buffered at a time.
+#[allow(clippy::large_enum_variant)]
+pub(crate) enum LeaderEvent {
     Scheduler(scheduler::Decisions),
     Invoker(InvokerEffect),
     Shuffle(shuffle::OutboxTruncation),
@@ -156,6 +162,16 @@ pub(crate) enum ActionEffect {
     UpsertSchema(Schema),
     UpsertRuleBook(Arc<restate_limiter::RuleBook>),
     AwaitingRpcSelfProposeDone,
+    RpcProposal {
+        proposal: rpc::RpcProposal,
+        #[debug(skip)]
+        reciprocal: RpcReciprocal,
+    },
+    IngestRecords {
+        records: Vec<IngestRecord>,
+        #[debug(skip)]
+        callback: AppendCallback,
+    },
 }
 
 enum State {
@@ -230,13 +246,19 @@ where
         }
     }
 
-    pub(super) fn should_process_rpc(&self) -> bool {
-        // In case of BecomingLeader we prefer to park RPC requests
-        // until we transition out of the current state.
-        matches!(
-            self.state,
-            State::Leader(_) | State::Follower | State::Candidate { .. }
-        )
+    pub(super) fn try_reserve_rpc_processing_permit(&self) -> Option<RpcProcessingPermit> {
+        match &self.state {
+            State::Leader(leader_state) => leader_state
+                .try_reserve_network_event_permit()
+                .ok()
+                .map(RpcProcessingPermit::Leader),
+            State::Follower | State::Candidate { .. } => Some(RpcProcessingPermit::NonLeader {
+                partition_id: self.partition_id(),
+            }),
+            // In case of BecomingLeader we prefer to park RPC requests
+            // until we transition out of the current state.
+            State::BecomingLeader { .. } => None,
+        }
     }
 
     #[instrument(level = "debug", skip_all, fields(leader_epoch = %leadership_info.leader_epoch))]
@@ -824,11 +846,11 @@ where
     ///
     /// * Follower: Nothing to do
     /// * Candidate: Monitor appender task
-    /// * Leader: Await action effects and monitor appender task
+    /// * Leader: Await events and monitor appender task
     pub async fn run(
         &mut self,
         ctx: impl Processor + HasVQueues,
-    ) -> Result<Vec<ActionEffect>, Error> {
+    ) -> Result<Vec<LeaderEvent>, Error> {
         match &mut self.state {
             State::Follower => Ok(futures::future::pending::<Vec<_>>().await),
             State::Candidate { self_proposer, .. }
@@ -842,17 +864,15 @@ where
         }
     }
 
-    pub async fn handle_action_effects(
+    pub async fn handle_events(
         &mut self,
-        action_effects: impl IntoIterator<Item = ActionEffect>,
+        events: impl IntoIterator<Item = LeaderEvent>,
     ) -> Result<(), Error> {
         match &mut self.state {
             State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => {
                 // nothing to do :-)
             }
-            State::Leader(leader_state) => {
-                leader_state.handle_action_effects(action_effects).await?
-            }
+            State::Leader(leader_state) => leader_state.handle_events(events).await?,
         }
 
         Ok(())
@@ -891,102 +911,6 @@ impl<T> LeadershipState<T> {
             }
         }
     }
-
-    pub async fn handle_rpc_proposal_command(
-        &mut self,
-        request_id: PartitionProcessorRpcRequestId,
-        reciprocal: Reciprocal<
-            Oneshot<Result<PartitionProcessorRpcResponse, PartitionProcessorRpcError>>,
-        >,
-        partition_key: PartitionKey,
-        cmd: Command,
-    ) {
-        match &mut self.state {
-            State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => {
-                // Just fail the rpc
-                reciprocal.send(Err(PartitionProcessorRpcError::NotLeader(
-                    self.partition_id,
-                )))
-            }
-            State::Leader(leader_state) => {
-                leader_state
-                    .handle_rpc_proposal_command(request_id, reciprocal, partition_key, cmd)
-                    .await;
-            }
-        }
-    }
-
-    pub async fn propose_pause_and_fence(
-        &mut self,
-        request_id: PartitionProcessorRpcRequestId,
-        reciprocal: Reciprocal<
-            Oneshot<Result<PartitionProcessorRpcResponse, PartitionProcessorRpcError>>,
-        >,
-        invocation_id: InvocationId,
-        cmd: Command,
-    ) {
-        match &mut self.state {
-            State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => {
-                // Just fail the rpc
-                reciprocal.send(Err(PartitionProcessorRpcError::NotLeader(
-                    self.partition_id,
-                )))
-            }
-            State::Leader(leader_state) => {
-                leader_state
-                    .propose_pause_and_fence(request_id, reciprocal, invocation_id, cmd)
-                    .await;
-            }
-        }
-    }
-
-    /// Append a command to Bifrost without dedup information, responding on Bifrost commit.
-    pub async fn append_and_respond_asynchronously(
-        &mut self,
-        partition_key: PartitionKey,
-        cmd: Command,
-        reciprocal: Reciprocal<
-            Oneshot<Result<PartitionProcessorRpcResponse, PartitionProcessorRpcError>>,
-        >,
-        success_response: PartitionProcessorRpcResponse,
-    ) {
-        match &mut self.state {
-            State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => reciprocal
-                .send(Err(PartitionProcessorRpcError::NotLeader(
-                    self.partition_id,
-                ))),
-            State::Leader(leader_state) => {
-                leader_state
-                    .append_and_respond_asynchronously(
-                        partition_key,
-                        cmd,
-                        reciprocal,
-                        success_response,
-                    )
-                    .await;
-            }
-        }
-    }
-
-    /// Forward externally-created records to this partition.
-    pub async fn forward_many_with_callback<F>(
-        &mut self,
-        records: impl ExactSizeIterator<Item = IngestRecord>,
-        callback: F,
-    ) where
-        F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static,
-    {
-        match &mut self.state {
-            State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => callback(
-                Err(PartitionProcessorRpcError::NotLeader(self.partition_id)),
-            ),
-            State::Leader(leader_state) => {
-                leader_state
-                    .forward_many_with_callback(records, callback)
-                    .await;
-            }
-        }
-    }
 }
 
 #[derive(Debug, derive_more::From)]
@@ -1031,6 +955,38 @@ impl shuffle::OutboxReader for OutboxReader {
     }
 }
 
+pub(super) enum RpcProcessingPermit {
+    Leader(tokio::sync::mpsc::OwnedPermit<LeaderEvent>),
+    NonLeader { partition_id: PartitionId },
+}
+
+impl RpcProcessingPermit {
+    pub fn buffer_rpc_proposal(self, proposal: rpc::RpcProposal, reciprocal: RpcReciprocal) {
+        match self {
+            RpcProcessingPermit::NonLeader { partition_id } => {
+                reciprocal.send(Err(PartitionProcessorRpcError::NotLeader(partition_id)))
+            }
+            RpcProcessingPermit::Leader(permit) => {
+                permit.send(LeaderEvent::RpcProposal {
+                    proposal,
+                    reciprocal,
+                });
+            }
+        }
+    }
+
+    pub fn buffer_forwarded_records(self, records: Vec<IngestRecord>, callback: AppendCallback) {
+        match self {
+            RpcProcessingPermit::NonLeader { partition_id } => {
+                callback(Err(PartitionProcessorRpcError::NotLeader(partition_id)))
+            }
+            RpcProcessingPermit::Leader(permit) => {
+                permit.send(LeaderEvent::IngestRecords { records, callback });
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::num::NonZeroUsize;