diff --git a/Cargo.lock b/Cargo.lock index 6558465035..baa10e55f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9094,7 +9094,6 @@ dependencies = [ "googletest", "itertools 0.14.0", "metrics", - "mockall", "opentelemetry", "parking_lot", "prost", diff --git a/crates/worker/Cargo.toml b/crates/worker/Cargo.toml index 32ecb5b5b6..add8bd11ce 100644 --- a/crates/worker/Cargo.toml +++ b/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 } diff --git a/crates/worker/src/partition/leadership/leader_state.rs b/crates/worker/src/partition/leadership/leader_state.rs index 0f06046f59..ec6158a05d 100644 --- a/crates/worker/src/partition/leadership/leader_state.rs +++ b/crates/worker/src/partition/leadership/leader_state.rs @@ -59,8 +59,11 @@ use restate_worker_api::{SchedulerStatusEntry, UserLimitCounterEntry}; use crate::metric_definitions::{PARTITION_HANDLE_LEADER_ACTIONS, USAGE_LEADER_ACTION_COUNT}; use crate::partition::cleaner::{CleanerEffect, CleanerHandle}; use crate::partition::leadership::self_proposer::SelfProposer; -use crate::partition::leadership::{ActionEffect, Error, InvokerStream, TimerService}; +use crate::partition::leadership::{ + AppendCallback, Error, InvokerStream, LeaderEvent, RpcReciprocal, TimerService, +}; use crate::partition::processor::{FsmAccess, Processor}; +use crate::partition::rpc::{ReplyOn, RpcProposal}; use crate::partition::shuffle; use crate::partition::shuffle::HintSender; use crate::partition::state_machine::Action; @@ -69,9 +72,7 @@ use crate::partition_processor_manager::LeaderQueryGuard; use super::durability_tracker::DurabilityTracker; const BATCH_READY_UP_TO: usize = 10; - -type RpcReciprocal = - Reciprocal>>; +const NETWORK_EVENTS_BUFFER_SIZE: usize = 1; pub struct LeaderState { pub(crate) partition_id: PartitionId, @@ -111,6 +112,8 @@ pub struct LeaderState { cleaner_handle: CleanerHandle, trimmer_task_id: TaskId, durability_tracker: DurabilityTracker, + network_events_tx: tokio::sync::mpsc::Sender, + network_events_stream: ReceiverStream, // Unregisters the leader-query registry entry on drop. Must live as long as // the partition processor's select! is willing to serve scheduler queries. _leader_query_guard: LeaderQueryGuard, @@ -141,6 +144,8 @@ impl LeaderState { leader_query_guard: LeaderQueryGuard, rule_book_rx: tokio::sync::watch::Receiver>, ) -> Self { + let (network_events_tx, network_events_rx) = + tokio::sync::mpsc::channel(NETWORK_EVENTS_BUFFER_SIZE); LeaderState { partition_id, leader_epoch, @@ -165,6 +170,8 @@ impl LeaderState { invoker_stream: invoker_rx, shuffle_stream: ReceiverStream::new(shuffle_rx), durability_tracker, + network_events_tx, + network_events_stream: ReceiverStream::new(network_events_rx), _leader_query_guard: leader_query_guard, } } @@ -189,20 +196,28 @@ impl LeaderState { self.scheduler.scan_user_limit_counters(keys.start()) } - /// Runs the leader specific task which is the awaiting of action effects and the monitoring - /// of unmanaged tasks. + pub(super) fn try_reserve_network_event_permit( + &self, + ) -> Result< + tokio::sync::mpsc::OwnedPermit, + tokio::sync::mpsc::error::TrySendError>, + > { + self.network_events_tx.clone().try_reserve_owned() + } + + /// Runs the leader-specific task by awaiting events and monitoring unmanaged tasks. /// /// Important: The future needs to be cancellation safe since it is polled as a tokio::select /// arm! pub async fn run( &mut self, ctx: impl Processor + HasVQueues, - ) -> Result, Error> { + ) -> Result, Error> { let timer_stream = std::pin::pin!(stream::unfold( &mut self.timer_service, |timer_service| async { let timer_value = timer_service.as_mut().next_timer().await; - Some((ActionEffect::Timer(timer_value), timer_service)) + Some((LeaderEvent::Timer(timer_value), timer_service)) } )); let vqueue_metas = ctx.vqueues(); @@ -212,7 +227,7 @@ impl LeaderState { let scheduler_stream = std::pin::pin!(stream::unfold(&mut self.scheduler, |scheduler| async { match scheduler.schedule_next(vqueue_metas).await { - Ok(decisions) => Some((ActionEffect::Scheduler(decisions), scheduler)), + Ok(decisions) => Some((LeaderEvent::Scheduler(decisions), scheduler)), Err(e) => { error!("Fatal error when polling scheduler: {e}"); None @@ -227,7 +242,7 @@ impl LeaderState { std::future::ready( Some(Metadata::with_current(|m| m.schema())) .filter(|schema| schema.version() > current_version) - .map(|schema| ActionEffect::UpsertSchema(schema.clone())), + .map(|schema| LeaderEvent::UpsertSchema(schema.clone())), ) }); @@ -236,19 +251,21 @@ impl LeaderState { // newer than the partition's current in-memory book. let current_version = ctx.fsm().rule_book().version(); std::future::ready( - (book.version() > current_version).then(|| ActionEffect::UpsertRuleBook(book)), + (book.version() > current_version).then(|| LeaderEvent::UpsertRuleBook(book)), ) }); - let invoker_stream = (&mut self.invoker_stream).map(ActionEffect::Invoker); - let shuffle_stream = (&mut self.shuffle_stream).map(ActionEffect::Shuffle); - let cleaner_stream = self.cleaner_handle.effects().map(ActionEffect::Cleaner); + let invoker_stream = (&mut self.invoker_stream).map(LeaderEvent::Invoker); + let shuffle_stream = (&mut self.shuffle_stream).map(LeaderEvent::Shuffle); + let cleaner_stream = self.cleaner_handle.effects().map(LeaderEvent::Cleaner); let dur_tracker_stream = - (&mut self.durability_tracker).map(ActionEffect::PartitionMaintenance); + (&mut self.durability_tracker).map(LeaderEvent::PartitionMaintenance); let awaiting_rpc_self_propose_stream = - (&mut self.awaiting_rpc_self_propose).map(|_| ActionEffect::AwaitingRpcSelfProposeDone); + (&mut self.awaiting_rpc_self_propose).map(|_| LeaderEvent::AwaitingRpcSelfProposeDone); + + let network_event_stream = &mut self.network_events_stream; let all_streams = futures::stream_select!( scheduler_stream, @@ -259,7 +276,8 @@ impl LeaderState { awaiting_rpc_self_propose_stream, dur_tracker_stream, schema_stream, - rule_book_stream + rule_book_stream, + network_event_stream ); let mut all_streams = all_streams.ready_chunks(BATCH_READY_UP_TO); @@ -351,16 +369,30 @@ impl LeaderState { for fut in self.awaiting_rpc_self_propose.iter_mut() { fut.fail_with_lost_leadership(self.partition_id); } + // No network processing permit can outlive the partition processor's select iteration, + // so after closing the receiver all buffered events are immediately available. + self.network_events_stream.close(); + while let Some(Some(event)) = self.network_events_stream.next().now_or_never() { + match event { + LeaderEvent::RpcProposal { reciprocal, .. } => reciprocal.send(Err( + PartitionProcessorRpcError::LostLeadership(self.partition_id), + )), + LeaderEvent::IngestRecords { callback, .. } => callback(Err( + PartitionProcessorRpcError::LostLeadership(self.partition_id), + )), + _ => unreachable!("only network events are buffered"), + } + } } - pub async fn handle_action_effects( + pub async fn handle_events( &mut self, - action_effects: impl IntoIterator, + events: impl IntoIterator, ) -> Result<(), Error> { let mut arena = BytesMut::with_capacity(128 * 1024); - for effect in action_effects { - match effect { - ActionEffect::Scheduler(decisions) => { + for event in events { + match event { + LeaderEvent::Scheduler(decisions) => { let Decisions { qids, num_run, @@ -398,7 +430,7 @@ impl LeaderState { .self_propose_many(commands.into_iter()) .await?; } - ActionEffect::PartitionMaintenance(partition_durability) => { + LeaderEvent::PartitionMaintenance(partition_durability) => { // based on configuration, whether to consider partition-local durability in // the replica-set as a sufficient source of durability, or only snapshots. self.self_proposer @@ -408,7 +440,7 @@ impl LeaderState { ) .await?; } - ActionEffect::Invoker(fenced) => { + LeaderEvent::Invoker(fenced) => { let invocation_id = fenced.effect.invocation_id; // Fence stale effects at write time: only self-propose if the effect carries // the token of the invocation's *current* attempt. A mismatch (or a missing @@ -434,7 +466,7 @@ impl LeaderState { ); } } - ActionEffect::Shuffle(outbox_truncation) => { + LeaderEvent::Shuffle(outbox_truncation) => { // todo: Until we support partition splits we need to get rid of outboxes or introduce partition // specific destination messages that are identified by a partition_id self.self_proposer @@ -444,12 +476,12 @@ impl LeaderState { ) .await?; } - ActionEffect::Timer(timer) => { + LeaderEvent::Timer(timer) => { self.self_proposer .self_propose(timer.invocation_id().partition_key(), Command::Timer(timer)) .await?; } - ActionEffect::Cleaner(effect) => { + LeaderEvent::Cleaner(effect) => { let (invocation_id, cmd) = match effect { CleanerEffect::PurgeJournal(invocation_id) => ( invocation_id, @@ -471,7 +503,7 @@ impl LeaderState { .self_propose(invocation_id.partition_key(), cmd) .await?; } - ActionEffect::UpsertSchema(schema) => { + LeaderEvent::UpsertSchema(schema) => { if SemanticRestateVersion::current() .is_equal_or_newer_than(&RESTATE_VERSION_1_7_0) { @@ -488,7 +520,7 @@ impl LeaderState { .await?; } } - ActionEffect::UpsertRuleBook(rule_book) => { + LeaderEvent::UpsertRuleBook(rule_book) => { let cmd = restate_wal_protocol::control::UpsertRuleBookCommand { rule_book }; arena.reserve(cmd.encoded_len()); // safe to unwrap because we reserved enough space @@ -504,16 +536,52 @@ impl LeaderState { ) .await?; } - ActionEffect::AwaitingRpcSelfProposeDone => { + LeaderEvent::AwaitingRpcSelfProposeDone => { // Nothing to do here } + LeaderEvent::RpcProposal { + proposal, + reciprocal, + } => { + self.handle_rpc_proposal(proposal, reciprocal).await; + } + LeaderEvent::IngestRecords { records, callback } => { + self.forward_many_with_callback(records.into_iter(), callback) + .await; + } } } Ok(()) } - pub async fn handle_rpc_proposal_command( + async fn handle_rpc_proposal(&mut self, proposal: RpcProposal, reciprocal: RpcReciprocal) { + let RpcProposal { + partition_key, + cmd, + reply_on, + } = proposal; + + match reply_on { + ReplyOn::Apply { request_id } => { + self.handle_rpc_proposal_command(request_id, reciprocal, partition_key, cmd) + .await; + } + ReplyOn::Commit { response } => { + self.append_and_respond_asynchronously(partition_key, cmd, reciprocal, response) + .await; + } + ReplyOn::ApplyAndFence { + request_id, + invocation_id, + } => { + self.propose_pause_and_fence(request_id, reciprocal, invocation_id, cmd) + .await; + } + } + } + + async fn handle_rpc_proposal_command( &mut self, request_id: PartitionProcessorRpcRequestId, reciprocal: Reciprocal< @@ -554,7 +622,7 @@ impl LeaderState { /// dropped at write time with no pause in the committed log. After a successful append the /// pause's LSN is fixed, and because the leader self-proposes from a single task, every later /// invoker-effect self-propose observes the cleared map and is fenced. - pub async fn propose_pause_and_fence( + async fn propose_pause_and_fence( &mut self, request_id: PartitionProcessorRpcRequestId, reciprocal: Reciprocal< @@ -595,7 +663,7 @@ impl LeaderState { /// Records appended this way are never filtered by the dedup mechanism during leadership /// transitions, making this safe for fire-and-forget ingress commands (signals, invocation /// responses). - pub async fn append_and_respond_asynchronously( + async fn append_and_respond_asynchronously( &mut self, partition_key: PartitionKey, cmd: Command, @@ -610,22 +678,20 @@ impl LeaderState { Ok(commit_token) => { self.awaiting_rpc_self_propose.push(SelfAppendFuture::new( commit_token, - |result: Result<(), PartitionProcessorRpcError>| { + Box::new(|result: Result<(), PartitionProcessorRpcError>| { reciprocal.send(result.map(|_| success_response)); - }, + }), )); } Err(e) => reciprocal.send(Err(PartitionProcessorRpcError::Internal(e.to_string()))), } } - pub async fn forward_many_with_callback( + async fn forward_many_with_callback( &mut self, records: impl ExactSizeIterator, - callback: F, - ) where - F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static, - { + callback: AppendCallback, + ) { match self .self_proposer .forward_many_with_notification(records) @@ -894,56 +960,22 @@ impl LeaderState { } } -trait CallbackInner: Send + Sync + 'static { - fn call(self: Box, result: Result<(), PartitionProcessorRpcError>); -} - -impl CallbackInner for F -where - F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static, -{ - fn call(self: Box, result: Result<(), PartitionProcessorRpcError>) { - self(result) - } -} - -struct Callback { - inner: Box, -} - -impl Callback { - fn call(self, result: Result<(), PartitionProcessorRpcError>) { - self.inner.call(result); - } -} - -impl From for Callback -where - I: CallbackInner, -{ - fn from(value: I) -> Self { - Self { - inner: Box::new(value), - } - } -} - struct SelfAppendFuture { commit_token: CommitToken, - callback: Option, + callback: Option, } impl SelfAppendFuture { - fn new(commit_token: CommitToken, callback: impl Into) -> Self { + fn new(commit_token: CommitToken, callback: AppendCallback) -> Self { Self { commit_token, - callback: Some(callback.into()), + callback: Some(callback), } } fn fail_with_internal(&mut self) { if let Some(callback) = self.callback.take() { - callback.call(Err(PartitionProcessorRpcError::Internal( + callback(Err(PartitionProcessorRpcError::Internal( "error when proposing to bifrost".to_string(), ))); } @@ -951,7 +983,7 @@ impl SelfAppendFuture { fn fail_with_lost_leadership(&mut self, this_partition_id: PartitionId) { if let Some(callback) = self.callback.take() { - callback.call(Err(PartitionProcessorRpcError::LostLeadership( + callback(Err(PartitionProcessorRpcError::LostLeadership( this_partition_id, ))); } @@ -959,7 +991,7 @@ impl SelfAppendFuture { fn succeed_with_appended(&mut self) { if let Some(callback) = self.callback.take() { - callback.call(Ok(())) + callback(Ok(())) } } } diff --git a/crates/worker/src/partition/leadership/mod.rs b/crates/worker/src/partition/leadership/mod.rs index 0e3868ce13..17727c5294 100644 --- a/crates/worker/src/partition/leadership/mod.rs +++ b/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; type InvokerStream = ReceiverStream; +type RpcReciprocal = + Reciprocal>>; +pub(super) type AppendCallback = + Box) + 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), AwaitingRpcSelfProposeDone, + RpcProposal { + proposal: rpc::RpcProposal, + #[debug(skip)] + reciprocal: RpcReciprocal, + }, + IngestRecords { + records: Vec, + #[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 { + 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, Error> { + ) -> Result, Error> { match &mut self.state { State::Follower => Ok(futures::future::pending::>().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, + events: impl IntoIterator, ) -> 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 LeadershipState { } } } - - pub async fn handle_rpc_proposal_command( - &mut self, - request_id: PartitionProcessorRpcRequestId, - reciprocal: Reciprocal< - Oneshot>, - >, - 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>, - >, - 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>, - >, - 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( - &mut self, - records: impl ExactSizeIterator, - 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), + 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, 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; diff --git a/crates/worker/src/partition/mod.rs b/crates/worker/src/partition/mod.rs index 9655847c0b..2443a0aeea 100644 --- a/crates/worker/src/partition/mod.rs +++ b/crates/worker/src/partition/mod.rs @@ -29,6 +29,7 @@ pub mod state_machine; mod state_machine; pub mod types; +use self::leadership::RpcProcessingPermit; pub use self::node::NodeContext; // Re-exported so external drivers (e.g. pp-bench) can build a context to drive `StateMachine::apply`. #[cfg(feature = "expose-internals")] @@ -86,6 +87,7 @@ use restate_vqueues::context::HasVQueues; use restate_wal_protocol::control::{CurrentReplicaSetConfiguration, NextReplicaSetConfiguration}; use restate_wal_protocol::v2::CommandScope; use restate_wal_protocol::{Envelope, v2}; +use restate_worker_api::invoker::InvokerHandle; use restate_worker_api::{LeaderQueryCommand, LeaderQueryReceiver}; use self::processor::commands::{ @@ -97,7 +99,7 @@ use self::state_machine::StateMachine; use crate::metric_definitions::{ LEADER_LABEL, LEADER_LABEL_LEADER, PARTITION_RECORD_COMMITTED_TO_READ_LATENCY_SECONDS, }; -use crate::partition::leadership::LeadershipState; +use crate::partition::leadership::{AppendCallback, LeadershipState}; use crate::partition::processor::leadership::LeadershipContext; use crate::partition::state_machine::ActionCollector; @@ -287,7 +289,7 @@ pub enum ProcessorError { )] MigrationBarrier { reason: String }, #[error(transparent)] - ActionEffect(#[from] leadership::Error), + Leadership(#[from] leadership::Error), #[error(transparent)] ShutdownError(#[from] ShutdownError), #[error("log read stream has terminated")] @@ -526,6 +528,8 @@ where let config = self.node_ctx.config.live_load(); let max_batching_size = config.worker.max_command_batch_size(); let bytes_limit = config.worker.max_command_batch_bytes.as_usize(); + let network_processing_permit = + self.leadership_state.try_reserve_rpc_processing_permit(); tokio::select! { _ = self.target_leader_state_rx.changed() => { @@ -539,9 +543,9 @@ where self.leadership_state.maybe_step_down(&mut self.ctx, new_state.current_leader_epoch, new_state.current_leader).await; self.refresh_status(&mut durable_lsn_watch)?; } - Some(msg) = self.network_leader_svc_rx.next(), if self.leadership_state.should_process_rpc() => { + Some(msg) = self.network_leader_svc_rx.next(), if network_processing_permit.is_some() => { // todo: replace the live schema with the leader's consistent schema - self.on_rpc(msg, live_schemas.live_load(), &last_applied_lsn_watch).await; + self.on_rpc(msg, live_schemas.live_load(), &last_applied_lsn_watch, network_processing_permit.unwrap()).await; } _ = status_update_timer.tick() => { // Update the status to ensure changes are observable @@ -624,11 +628,11 @@ where self.leadership_state.handle_actions(&mut self.ctx, action_collector.drain(..))?; }, result = self.leadership_state.run(&mut self.ctx) => { - let action_effects = result?; - // We process the action_effects not directly in the run future because it + let events = result?; + // We process the events not directly in the run future because it // requires the run future to be cancellation safe. In the future this could be // implemented. - self.leadership_state.handle_action_effects(action_effects).await?; + self.leadership_state.handle_events(events).await?; } Some(leader_query_cmd) = self.leader_query_rx.recv() => { self.on_leader_query(leader_query_cmd); @@ -674,17 +678,36 @@ where >, body: PartitionProcessorRpcRequest, schemas: &Schema, + permit: RpcProcessingPermit, ) { - let _ = rpc::RpcHandler::handle( - rpc::RpcContext::new( - &mut self.leadership_state, - schemas, - &mut self.partition_store, - ), - body, - rpc::Replier::new(response_tx), - ) - .await; + let context = rpc::RpcContext::new( + self.leadership_state.is_leader(), + self.leadership_state.partition_id(), + schemas, + &mut self.partition_store, + ); + let decision = rpc::RpcHandler::handle(context, body).await; + + match decision { + rpc::Decision::Propose(proposal) => permit.buffer_rpc_proposal(proposal, response_tx), + rpc::Decision::Reply(reply) => response_tx.send(reply), + rpc::Decision::NotifyInvokerAndReply { + notification, + reply, + } => { + if let Some(invoker_handle) = self.leadership_state.invoker_handle() { + match notification { + rpc::InvokerNotification::RetryNow(invocation_id) => { + let _ = invoker_handle.retry_invocation_now(invocation_id); + } + rpc::InvokerNotification::Pause(invocation_id) => { + let _ = invoker_handle.pause_invocation(invocation_id); + } + } + } + response_tx.send(Ok(reply)); + } + } } fn refresh_status( @@ -724,16 +747,18 @@ where msg: ServiceMessage, schemas: &Schema, last_applied_lsn_watch: &watch::Receiver, + permit: RpcProcessingPermit, ) { match msg { ServiceMessage::Rpc(msg) if msg.msg_type() == PartitionProcessorRpcRequest::TYPE => { let msg = msg.into_typed::(); // note: split() decodes the payload let (response_tx, body) = msg.split(); - self.on_pp_rpc_request(response_tx, body, schemas).await; + self.on_pp_rpc_request(response_tx, body, schemas, permit) + .await; } ServiceMessage::Rpc(msg) if msg.msg_type() == ReceivedIngestRequest::TYPE => { - self.on_pp_ingest_request(msg.into_typed()).await; + self.on_pp_ingest_request(msg.into_typed(), permit); } ServiceMessage::Rpc(msg) if msg.msg_type() == DedupSequenceNrQueryRequest::TYPE => { self.wait_for_tail_then( @@ -832,26 +857,24 @@ where } } - async fn on_pp_ingest_request(&mut self, msg: Incoming>) { + fn on_pp_ingest_request( + &mut self, + msg: Incoming>, + permit: RpcProcessingPermit, + ) { let (reciprocal, request) = msg.split(); + let callback: AppendCallback = Box::new(move |result| match result { + Ok(()) => reciprocal.send(ResponseStatus::Ack.into()), + Err( + PartitionProcessorRpcError::NotLeader(id) + | PartitionProcessorRpcError::LostLeadership(id), + ) => reciprocal.send(ResponseStatus::NotLeader { of: id }.into()), + Err(PartitionProcessorRpcError::Internal(msg)) => { + reciprocal.send(ResponseStatus::Internal { msg }.into()) + } + }); - self.leadership_state - .forward_many_with_callback( - request.records.into_iter(), - move |result: Result<(), PartitionProcessorRpcError>| match result { - Ok(_) => reciprocal.send(ResponseStatus::Ack.into()), - Err(err) => match err { - PartitionProcessorRpcError::NotLeader(id) - | PartitionProcessorRpcError::LostLeadership(id) => { - reciprocal.send(ResponseStatus::NotLeader { of: id }.into()) - } - PartitionProcessorRpcError::Internal(msg) => { - reciprocal.send(ResponseStatus::Internal { msg }.into()) - } - }, - }, - ) - .await; + permit.buffer_forwarded_records(request.records, callback); } // --- Apply new commands/records diff --git a/crates/worker/src/partition/rpc/append_invocation.rs b/crates/worker/src/partition/rpc/append_invocation.rs index 50e8fa0360..ed4bbbab9a 100644 --- a/crates/worker/src/partition/rpc/append_invocation.rs +++ b/crates/worker/src/partition/rpc/append_invocation.rs @@ -21,12 +21,7 @@ pub(super) struct Request { pub(super) append_invocation_reply_on: AppendInvocationReplyOn, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = PartitionProcessorRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { @@ -34,8 +29,7 @@ impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler invocation_request, append_invocation_reply_on, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { let mut service_invocation = ServiceInvocation::from_request( Arc::unwrap_or_clone(invocation_request), invocation::Source::ingress(request_id), @@ -58,25 +52,18 @@ impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler let partition_key = service_invocation.partition_key(); let cmd = Command::Invoke(Box::new(service_invocation)); - match append_invocation_reply_on { - AppendInvocationReplyOn::Appended => { - self.proposer - .append_and_respond_asynchronously( - partition_key, - cmd, - replier, - PartitionProcessorRpcResponse::Appended, - ) - .await; - } - AppendInvocationReplyOn::Submitted | AppendInvocationReplyOn::Output => { - self.proposer - .handle_rpc_proposal_command(partition_key, cmd, request_id, replier) - .await; - } - } - - Ok(()) + Decision::Propose(RpcProposal { + partition_key, + cmd, + reply_on: match append_invocation_reply_on { + AppendInvocationReplyOn::Appended => ReplyOn::Commit { + response: PartitionProcessorRpcResponse::Appended, + }, + AppendInvocationReplyOn::Submitted | AppendInvocationReplyOn::Output => { + ReplyOn::Apply { request_id } + } + }, + }) } } @@ -84,122 +71,91 @@ impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler mod tests { use super::*; - use crate::partition::rpc::MockActuator; - use futures::FutureExt; use googletest::prelude::*; use restate_test_util::let_assert; - use std::future::ready; use test_log::test; - #[test(restate_core::test)] - async fn reply_on_appended() { - let mut proposer = MockActuator::new(); - proposer - .expect_append_and_respond_asynchronously::() - .return_once_st(|_, cmd, _, _| { - let_assert!(Command::Invoke(service_invocation) = cmd); - assert_that!( - service_invocation, - points_to(all!( - field!(ServiceInvocation.response_sink, none()), - field!(ServiceInvocation.submit_notification_sink, none()), - )) - ); - ready(()).boxed() - }); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - - let (tx, _rx) = Reciprocal::mock(); + async fn handle( + request_id: PartitionProcessorRpcRequestId, + append_invocation_reply_on: AppendInvocationReplyOn, + ) -> Decision { RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut ()), + RpcContext::new(true, PartitionId::MIN, &(), &mut ()), Request { - request_id: Default::default(), + request_id, invocation_request: Arc::new(InvocationRequest::mock()), - append_invocation_reply_on: AppendInvocationReplyOn::Appended, + append_invocation_reply_on, }, - Replier::new(tx), ) .await - .unwrap(); + } + + #[test(restate_core::test)] + async fn reply_on_appended() { + let_assert!( + Decision::Propose(RpcProposal { + cmd: Command::Invoke(service_invocation), + reply_on: ReplyOn::Commit { response }, + .. + }) = handle(Default::default(), AppendInvocationReplyOn::Appended).await + ); + assert_eq!(response, PartitionProcessorRpcResponse::Appended); + assert_that!( + service_invocation, + points_to(all!( + field!(ServiceInvocation.response_sink, none()), + field!(ServiceInvocation.submit_notification_sink, none()), + )) + ); } #[test(restate_core::test)] async fn reply_on_submitted() { let request_id = PartitionProcessorRpcRequestId::new(); - let mut proposer = MockActuator::new(); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(|_, cmd, req_id, _| { - let_assert!(Command::Invoke(service_invocation) = cmd); - assert_that!( - service_invocation, - points_to(all!( - field!(ServiceInvocation.response_sink, none()), - field!( - ServiceInvocation.submit_notification_sink, - some(eq(SubmitNotificationSink::Ingress { request_id: req_id })) - ), - )) - ); - ready(()).boxed() - }); - - let (tx, _rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut ()), - Request { - request_id, - invocation_request: Arc::new(InvocationRequest::mock()), - append_invocation_reply_on: AppendInvocationReplyOn::Submitted, - }, - Replier::new(tx), - ) - .await - .unwrap(); + let_assert!( + Decision::Propose(RpcProposal { + cmd: Command::Invoke(service_invocation), + reply_on: ReplyOn::Apply { + request_id: actual_request_id, + }, + .. + }) = handle(request_id, AppendInvocationReplyOn::Submitted).await + ); + assert_eq!(actual_request_id, request_id); + assert_that!( + service_invocation, + points_to(all!( + field!(ServiceInvocation.response_sink, none()), + field!( + ServiceInvocation.submit_notification_sink, + some(eq(SubmitNotificationSink::Ingress { request_id })) + ), + )) + ); } #[test(restate_core::test)] async fn reply_on_output() { let request_id = PartitionProcessorRpcRequestId::new(); - let mut proposer = MockActuator::new(); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(|_, cmd, req_id, _| { - let_assert!(Command::Invoke(service_invocation) = cmd); - assert_that!( - service_invocation, - points_to(all!( - field!( - ServiceInvocation.response_sink, - some(eq(ServiceInvocationResponseSink::Ingress { - request_id: req_id - })) - ), - field!(ServiceInvocation.submit_notification_sink, none()), - )) - ); - ready(()).boxed() - }); - - let (tx, _rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut ()), - Request { - request_id, - invocation_request: Arc::new(InvocationRequest::mock()), - append_invocation_reply_on: AppendInvocationReplyOn::Output, - }, - Replier::new(tx), - ) - .await - .unwrap(); + let_assert!( + Decision::Propose(RpcProposal { + cmd: Command::Invoke(service_invocation), + reply_on: ReplyOn::Apply { + request_id: actual_request_id, + }, + .. + }) = handle(request_id, AppendInvocationReplyOn::Output).await + ); + assert_eq!(actual_request_id, request_id); + assert_that!( + service_invocation, + points_to(all!( + field!( + ServiceInvocation.response_sink, + some(eq(ServiceInvocationResponseSink::Ingress { request_id })) + ), + field!(ServiceInvocation.submit_notification_sink, none()), + )) + ); } } diff --git a/crates/worker/src/partition/rpc/append_invocation_response.rs b/crates/worker/src/partition/rpc/append_invocation_response.rs index 7f6842a644..897216ccc7 100644 --- a/crates/worker/src/partition/rpc/append_invocation_response.rs +++ b/crates/worker/src/partition/rpc/append_invocation_response.rs @@ -18,28 +18,19 @@ pub(super) struct Request { pub(super) invocation_response: InvocationResponse, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = PartitionProcessorRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { invocation_response, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .append_and_respond_asynchronously( - invocation_response.partition_key(), - Command::InvocationResponse(invocation_response), - replier, - PartitionProcessorRpcResponse::Appended, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_response.partition_key(), + cmd: Command::InvocationResponse(invocation_response), + reply_on: ReplyOn::Commit { + response: PartitionProcessorRpcResponse::Appended, + }, + }) } } diff --git a/crates/worker/src/partition/rpc/append_signal.rs b/crates/worker/src/partition/rpc/append_signal.rs index d8f08d0df6..a932458ea0 100644 --- a/crates/worker/src/partition/rpc/append_signal.rs +++ b/crates/worker/src/partition/rpc/append_signal.rs @@ -20,32 +20,23 @@ pub(super) struct Request { pub(super) signal: Signal, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = PartitionProcessorRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { invocation_id, signal, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .append_and_respond_asynchronously( - invocation_id.partition_key(), - Command::NotifySignal(NotifySignalRequest { - invocation_id, - signal, - }), - replier, - PartitionProcessorRpcResponse::Appended, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::NotifySignal(NotifySignalRequest { + invocation_id, + signal, + }), + reply_on: ReplyOn::Commit { + response: PartitionProcessorRpcResponse::Appended, + }, + }) } } diff --git a/crates/worker/src/partition/rpc/cancel_invocation.rs b/crates/worker/src/partition/rpc/cancel_invocation.rs index d7829db624..0a06fcc8cb 100644 --- a/crates/worker/src/partition/rpc/cancel_invocation.rs +++ b/crates/worker/src/partition/rpc/cancel_invocation.rs @@ -14,7 +14,6 @@ use restate_types::invocation::{ IngressInvocationResponseSink, InvocationMutationResponseSink, InvocationTermination, TerminationFlavor, }; -use restate_types::net::partition_processor::CancelInvocationRpcResponse; use restate_wal_protocol::Command; pub(super) struct Request { @@ -22,35 +21,24 @@ pub(super) struct Request { pub(super) invocation_id: InvocationId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = CancelInvocationRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { request_id, invocation_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::TerminateInvocation(InvocationTermination { - invocation_id, - flavor: TerminationFlavor::Cancel, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::TerminateInvocation(InvocationTermination { + invocation_id, + flavor: TerminationFlavor::Cancel, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } } diff --git a/crates/worker/src/partition/rpc/get_invocation_output.rs b/crates/worker/src/partition/rpc/get_invocation_output.rs index f864239a6a..1d17124a23 100644 --- a/crates/worker/src/partition/rpc/get_invocation_output.rs +++ b/crates/worker/src/partition/rpc/get_invocation_output.rs @@ -28,9 +28,8 @@ pub(super) struct Request { pub(super) response_mode: GetInvocationOutputResponseMode, } -impl<'a, TActuator, TSchemas, TStorage> RpcContext<'a, TActuator, TSchemas, TStorage> +impl<'a, TSchemas, TStorage> RpcContext<'a, TSchemas, TStorage> where - TActuator: Actuator, TStorage: ReadInvocationStatusTable, { async fn get_invocation_output( @@ -65,14 +64,10 @@ where } } -impl<'a, Proposer: Actuator, TSchemas, Storage> RpcHandler - for RpcContext<'a, Proposer, TSchemas, Storage> +impl<'a, TSchemas, Storage> RpcHandler for RpcContext<'a, TSchemas, Storage> where Storage: ReadInvocationStatusTable, { - type Output = PartitionProcessorRpcResponse; - type Error = (); - async fn handle( mut self, Request { @@ -80,8 +75,7 @@ where invocation_query, response_mode, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { match response_mode { GetInvocationOutputResponseMode::BlockWhenNotReady => { // Try to get invocation output now, if it's ready reply immediately with it @@ -89,22 +83,18 @@ where .get_invocation_output(request_id, invocation_query.clone()) .await { - replier.send(ready_result); - return Ok(()); + return Decision::Reply(Ok(ready_result)); } - self.proposer - .handle_rpc_proposal_command( - invocation_query.partition_key(), - Command::AttachInvocation(AttachInvocationRequest { - invocation_query, - block_on_inflight: true, - response_sink: ServiceInvocationResponseSink::Ingress { request_id }, - }), - request_id, - replier, - ) - .await; + Decision::Propose(RpcProposal { + partition_key: invocation_query.partition_key(), + cmd: Command::AttachInvocation(AttachInvocationRequest { + invocation_query, + block_on_inflight: true, + response_sink: ServiceInvocationResponseSink::Ingress { request_id }, + }), + reply_on: ReplyOn::Apply { request_id }, + }) } GetInvocationOutputResponseMode::ReplyIfNotReady => { // Reading invocation output from a non-leader partition processor can return @@ -121,20 +111,17 @@ where // An open question is how to handle partitioned followers that can no longer apply // the latest records (e.g. due to network partitions) — they would need to time // out and return an error rather than blocking indefinitely. - if !self.proposer.is_leader() { - replier.send_result(Err(PartitionProcessorRpcError::NotLeader( - self.proposer.partition_id(), + if !self.is_leader { + return Decision::Reply(Err(PartitionProcessorRpcError::NotLeader( + self.partition_id, ))); - return Ok(()); } - replier.send_result( + Decision::Reply( self.get_invocation_output(request_id, invocation_query) .await .map_err(|err| PartitionProcessorRpcError::Internal(err.to_string())), - ); + ) } } - - Ok(()) } } diff --git a/crates/worker/src/partition/rpc/kill_invocation.rs b/crates/worker/src/partition/rpc/kill_invocation.rs index 4401199918..c279e2ea76 100644 --- a/crates/worker/src/partition/rpc/kill_invocation.rs +++ b/crates/worker/src/partition/rpc/kill_invocation.rs @@ -14,7 +14,6 @@ use restate_types::invocation::{ IngressInvocationResponseSink, InvocationMutationResponseSink, InvocationTermination, TerminationFlavor, }; -use restate_types::net::partition_processor::KillInvocationRpcResponse; use restate_wal_protocol::Command; pub(super) struct Request { @@ -22,35 +21,24 @@ pub(super) struct Request { pub(super) invocation_id: InvocationId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = KillInvocationRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { request_id, invocation_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::TerminateInvocation(InvocationTermination { - invocation_id, - flavor: TerminationFlavor::Kill, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::TerminateInvocation(InvocationTermination { + invocation_id, + flavor: TerminationFlavor::Kill, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } } diff --git a/crates/worker/src/partition/rpc/mod.rs b/crates/worker/src/partition/rpc/mod.rs index 4ca6e5f565..1beb557efa 100644 --- a/crates/worker/src/partition/rpc/mod.rs +++ b/crates/worker/src/partition/rpc/mod.rs @@ -20,10 +20,8 @@ mod purge_journal; mod restart_as_new_invocation; mod resume_invocation; -use std::marker::PhantomData; use std::sync::Arc; -use restate_core::network::{Oneshot, Reciprocal, TransportConnect}; use restate_storage_api::invocation_status_table::ReadInvocationStatusTable; use restate_storage_api::journal_table as journal_table_v1; use restate_storage_api::journal_table_v2::ReadJournalTable; @@ -37,195 +35,84 @@ use restate_types::net::partition_processor::{ }; use restate_types::schema::deployment::DeploymentResolver; use restate_wal_protocol::Command; -use restate_worker_api::invoker::InvokerHandle; -use crate::partition::leadership::LeadershipState; - -#[cfg_attr(test, mockall::automock)] -pub(super) trait Actuator { - fn handle_rpc_proposal_command( - &mut self, - partition_key: PartitionKey, - cmd: Command, - request_id: PartitionProcessorRpcRequestId, - replier: Replier, - ) -> impl Future; - - /// Like [`Self::handle_rpc_proposal_command`], but for the PauseInvocation command: once the - /// command is appended, it clears `invocation_id`'s in-memory fencing token so a straggler - /// effect from the attempt being paused is dropped at write time. - fn propose_pause_and_fence( - &mut self, - invocation_id: InvocationId, - cmd: Command, - request_id: PartitionProcessorRpcRequestId, - replier: Replier, - ) -> impl Future; - - /// Appends a command to Bifrost **without** dedup information, responding on Bifrost commit. - /// - /// Records appended this way are never filtered by the dedup mechanism during leadership - /// transitions, making this safe for fire-and-forget ingress commands (signals, invocation - /// responses). - fn append_and_respond_asynchronously< - O: 'static + Into + Send + Sync, - >( - &mut self, - partition_key: PartitionKey, - cmd: Command, - replier: Replier, - success_response: O, - ) -> impl Future; - - fn notify_invoker_to_retry_now(&mut self, invocation_id: InvocationId); - - fn notify_invoker_to_pause(&mut self, invocation_id: InvocationId); - - fn is_leader(&self) -> bool; - - fn partition_id(&self) -> PartitionId; +#[derive(Debug, Clone)] +pub(crate) struct RpcProposal { + pub(crate) partition_key: PartitionKey, + pub(crate) cmd: Command, + pub(crate) reply_on: ReplyOn, } -impl Actuator for LeadershipState -where - T: TransportConnect, -{ - async fn append_and_respond_asynchronously>( - &mut self, - partition_key: PartitionKey, - cmd: Command, - replier: Replier, - on_proposed_response: O, - ) { - LeadershipState::append_and_respond_asynchronously( - self, - partition_key, - cmd, - replier.0, - on_proposed_response.into(), - ) - .await - } +#[derive(Debug)] +#[allow(clippy::large_enum_variant)] +pub(crate) enum Decision { + Propose(RpcProposal), + /// Reply immediately; nothing is proposed. + Reply(Result), + /// Legacy invoker-owned paths only: poke the invoker, then reply immediately. + NotifyInvokerAndReply { + notification: InvokerNotification, + reply: PartitionProcessorRpcResponse, + }, +} - async fn handle_rpc_proposal_command( - &mut self, - partition_key: PartitionKey, - cmd: Command, +#[derive(Debug, Clone)] +pub(crate) enum ReplyOn { + /// Reponds to the request; the state machine's Action replies later. + Apply { request_id: PartitionProcessorRpcRequestId, - replier: Replier, - ) { - LeadershipState::handle_rpc_proposal_command( - self, - request_id, - replier.0, - partition_key, - cmd, - ) - .await - } - - async fn propose_pause_and_fence( - &mut self, - invocation_id: InvocationId, - cmd: Command, + }, + /// Append WITHOUT dedup ESN; reply `response` on Bifrost commit. + Commit { + response: PartitionProcessorRpcResponse, + }, + /// Like Apply, but clear the invocation's fencing token strictly AFTER the + /// append succeeds. + ApplyAndFence { request_id: PartitionProcessorRpcRequestId, - replier: Replier, - ) { - LeadershipState::propose_pause_and_fence(self, request_id, replier.0, invocation_id, cmd) - .await - } - - fn notify_invoker_to_retry_now(&mut self, invocation_id: InvocationId) { - let Some(invoker_handle) = LeadershipState::invoker_handle(self) else { - return; - }; - let _ = invoker_handle.retry_invocation_now(invocation_id); - } - - fn notify_invoker_to_pause(&mut self, invocation_id: InvocationId) { - let Some(invoker_handle) = LeadershipState::invoker_handle(self) else { - return; - }; - let _ = invoker_handle.pause_invocation(invocation_id); - } - - fn is_leader(&self) -> bool { - LeadershipState::is_leader(self) - } + invocation_id: InvocationId, + }, +} - fn partition_id(&self) -> PartitionId { - LeadershipState::partition_id(self) - } +#[derive(Debug)] +pub(crate) enum InvokerNotification { + RetryNow(InvocationId), // resume legacy path, resume_invocation.rs:92 + Pause(InvocationId), // pause legacy path, pause_invocation.rs:93 } -pub(super) struct RpcContext<'a, Actuator, Schemas, Storage> { - proposer: &'a mut Actuator, +pub(super) struct RpcContext<'a, Schemas, Storage> { + is_leader: bool, + partition_id: PartitionId, schemas: &'a Schemas, storage: &'a mut Storage, } -impl<'a, Actuator, Schemas, Storage> RpcContext<'a, Actuator, Schemas, Storage> { +impl<'a, Schemas, Storage> RpcContext<'a, Schemas, Storage> { pub(super) fn new( - proposer: &'a mut Actuator, + is_leader: bool, + partition_id: PartitionId, schemas: &'a Schemas, storage: &'a mut Storage, ) -> Self { Self { - proposer, + is_leader, + partition_id, schemas, storage, } } } -pub(super) struct Replier( - Reciprocal>>, - PhantomData, -); - -impl> Replier { - pub fn new( - reciprocal: Reciprocal< - Oneshot>, - >, - ) -> Self { - Self(reciprocal, Default::default()) - } - - pub fn send(self, msg: O) { - self.0.send(Ok(msg.into())); - } - - pub fn send_result(self, res: Result) { - self.0.send(res.map(|msg| msg.into())); - } - - pub fn map>(self) -> Replier { - Replier::new(self.0) - } -} - pub(super) trait RpcHandler { - type Output: Into; - type Error; - - fn handle( - self, - input: Input, - response_tx: Replier, - ) -> impl Future>; + fn handle(self, input: Input) -> impl Future; } -impl<'a, TActuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> +impl<'a, TSchemas, TStorage> RpcHandler + for RpcContext<'a, TSchemas, TStorage> where - TActuator: Actuator, TSchemas: DeploymentResolver, TStorage: ReadInvocationStatusTable + ReadJournalTable + journal_table_v1::ReadJournalTable, { - type Output = PartitionProcessorRpcResponse; - type Error = (); - async fn handle( self, PartitionProcessorRpcRequest { @@ -233,94 +120,69 @@ where partition_id: _, inner, }: PartitionProcessorRpcRequest, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { match inner { PartitionProcessorRpcRequestInner::AppendInvocation( invocation_request, append_invocation_reply_on, ) => { - self.handle( - append_invocation::Request { - request_id, - invocation_request, - append_invocation_reply_on, - }, - replier, - ) + self.handle(append_invocation::Request { + request_id, + invocation_request, + append_invocation_reply_on, + }) .await } PartitionProcessorRpcRequestInner::GetInvocationOutput( invocation_query, response_mode, ) => { - self.handle( - get_invocation_output::Request { - request_id, - invocation_query, - response_mode, - }, - replier, - ) + self.handle(get_invocation_output::Request { + request_id, + invocation_query, + response_mode, + }) .await } PartitionProcessorRpcRequestInner::AppendInvocationResponse(invocation_response) => { - self.handle( - append_invocation_response::Request { - invocation_response, - }, - replier, - ) + self.handle(append_invocation_response::Request { + invocation_response, + }) .await } PartitionProcessorRpcRequestInner::AppendSignal(invocation_id, signal) => { - self.handle( - append_signal::Request { - invocation_id, - signal, - }, - replier, - ) + self.handle(append_signal::Request { + invocation_id, + signal, + }) .await } PartitionProcessorRpcRequestInner::CancelInvocation { invocation_id } => { - self.handle( - cancel_invocation::Request { - request_id, - invocation_id, - }, - replier.map(), - ) + self.handle(cancel_invocation::Request { + request_id, + invocation_id, + }) .await } PartitionProcessorRpcRequestInner::KillInvocation { invocation_id } => { - self.handle( - kill_invocation::Request { - request_id, - invocation_id, - }, - replier.map(), - ) + self.handle(kill_invocation::Request { + request_id, + invocation_id, + }) .await } PartitionProcessorRpcRequestInner::PurgeInvocation { invocation_id } => { - self.handle( - purge_invocation::Request { - request_id, - invocation_id, - }, - replier.map(), - ) + self.handle(purge_invocation::Request { + request_id, + invocation_id, + }) .await } PartitionProcessorRpcRequestInner::PurgeJournal { invocation_id } => { - self.handle( - purge_journal::Request { - request_id, - invocation_id, - }, - replier.map(), - ) + self.handle(purge_journal::Request { + request_id, + invocation_id, + }) .await } PartitionProcessorRpcRequestInner::RestartAsNewInvocation { @@ -328,39 +190,30 @@ where copy_prefix_up_to_index_included, patch_deployment_id, } => { - self.handle( - restart_as_new_invocation::Request { - request_id, - invocation_id, - copy_prefix_up_to_index_included, - patch_deployment_id, - }, - replier.map(), - ) + self.handle(restart_as_new_invocation::Request { + request_id, + invocation_id, + copy_prefix_up_to_index_included, + patch_deployment_id, + }) .await } PartitionProcessorRpcRequestInner::ResumeInvocation { invocation_id, deployment_id, } => { - self.handle( - resume_invocation::Request { - request_id, - invocation_id, - update_deployment_id: deployment_id, - }, - replier.map(), - ) + self.handle(resume_invocation::Request { + request_id, + invocation_id, + update_deployment_id: deployment_id, + }) .await } PartitionProcessorRpcRequestInner::PauseInvocation { invocation_id } => { - self.handle( - pause_invocation::PauseRequest { - request_id, - invocation_id, - }, - replier.map(), - ) + self.handle(pause_invocation::PauseRequest { + request_id, + invocation_id, + }) .await } } diff --git a/crates/worker/src/partition/rpc/pause_invocation.rs b/crates/worker/src/partition/rpc/pause_invocation.rs index b058b45918..5e4f1ad940 100644 --- a/crates/worker/src/partition/rpc/pause_invocation.rs +++ b/crates/worker/src/partition/rpc/pause_invocation.rs @@ -10,7 +10,7 @@ use super::*; use restate_storage_api::invocation_status_table::InvocationStatus; -use restate_types::identifiers::InvocationId; +use restate_types::identifiers::{InvocationId, WithPartitionKey}; use restate_types::net::partition_processor::PauseInvocationRpcResponse; use restate_wal_protocol::invocation::PauseInvocationCommand; @@ -19,39 +19,32 @@ pub(super) struct PauseRequest { pub(super) invocation_id: InvocationId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> where TStorage: ReadInvocationStatusTable, { - type Output = PauseInvocationRpcResponse; - type Error = (); - async fn handle( self, PauseRequest { request_id, invocation_id, }: PauseRequest, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { // Reading from a non-leader partition processor can return stale results // (e.g. NotFound for an invocation that exists on the leader) because the // follower's local store may not have replayed all log entries yet. - if !self.proposer.is_leader() { - replier.send_result(Err(PartitionProcessorRpcError::NotLeader( - self.proposer.partition_id(), + if !self.is_leader { + return Decision::Reply(Err(PartitionProcessorRpcError::NotLeader( + self.partition_id, ))); - return Ok(()); } let status = match self.storage.get_invocation_status(&invocation_id).await { Ok(status) => status, Err(storage_error) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal( + return Decision::Reply(Err(PartitionProcessorRpcError::Internal( storage_error.to_string(), ))); - return Ok(()); } }; @@ -70,58 +63,53 @@ where // replies via Action::ForwardPauseInvocationResponse. propose_pause_and_fence clears // the leader's in-memory fencing token (after appending the command) so that any // straggler effect from the attempt we are pausing is dropped at write time. - self.proposer - .propose_pause_and_fence( - invocation_id, - Command::PauseInvocation( - PauseInvocationCommand { - invocation_id, - request_id: Some(request_id), - } - .bilrost_encode_to_bytes(), - ), + return Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::PauseInvocation( + PauseInvocationCommand { + invocation_id, + request_id: Some(request_id), + } + .bilrost_encode_to_bytes(), + ), + reply_on: ReplyOn::ApplyAndFence { request_id, - replier, - ) - .await; - return Ok(()); + invocation_id, + }, + }); } // -- Legacy path for invoker-owned (non-VQueue) invocations: best-effort invoker poke. match status { - InvocationStatus::Invoked(_) => { - self.proposer.notify_invoker_to_pause(invocation_id); - replier.send(PauseInvocationRpcResponse::Accepted); - } + InvocationStatus::Invoked(_) => Decision::NotifyInvokerAndReply { + notification: InvokerNotification::Pause(invocation_id), + reply: PauseInvocationRpcResponse::Accepted.into(), + }, InvocationStatus::Completed(_) | InvocationStatus::Scheduled(_) | InvocationStatus::Inboxed(_) => { - replier.send(PauseInvocationRpcResponse::NotRunning); + Decision::Reply(Ok(PauseInvocationRpcResponse::NotRunning.into())) } InvocationStatus::Paused(_) | InvocationStatus::Suspended { .. } => { - replier.send(PauseInvocationRpcResponse::AlreadyPaused); + Decision::Reply(Ok(PauseInvocationRpcResponse::AlreadyPaused.into())) } InvocationStatus::Free => { - replier.send(PauseInvocationRpcResponse::NotFound); + Decision::Reply(Ok(PauseInvocationRpcResponse::NotFound.into())) } - }; - - Ok(()) + } } } #[cfg(test)] mod tests { - use super::*; + use std::assert_matches; + use std::future::ready; - use crate::partition::rpc::MockActuator; - use futures::FutureExt; - use restate_core::network::Reciprocal; use restate_storage_api::invocation_status_table::InFlightInvocationMetadata; - use restate_types::identifiers::WithPartitionKey; - use std::future::ready; use test_log::test; + use super::*; + struct MockStorage { status: InvocationStatus, } @@ -146,12 +134,6 @@ mod tests { async fn reply_not_leader_when_not_leader() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(false); - proposer - .expect_partition_id() - .return_const(PartitionId::from(0)); - struct NoopStorage; impl ReadInvocationStatusTable for NoopStorage { #[allow(unreachable_code)] @@ -176,22 +158,19 @@ mod tests { let mut storage = NoopStorage; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = RpcHandler::handle( + RpcContext::new(false, PartitionId::from(0), &(), &mut storage), PauseRequest { request_id: Default::default(), invocation_id, }, - Replier::new(tx), ) - .await - .unwrap(); + .await; - assert!(matches!( - rx.recv().await, - Err(PartitionProcessorRpcError::NotLeader(_)) - )); + assert_matches!( + decision, + Decision::Reply(Err(PartitionProcessorRpcError::NotLeader(_))) + ); } /// A non-VQueue (invoker-owned) invocation uses the legacy invoker-poke path and is not @@ -200,34 +179,28 @@ mod tests { async fn non_vqueue_invocation_uses_legacy_invoker_poke() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_notify_invoker_to_pause() - .return_once_st(move |got_invocation_id| { - assert_eq!(got_invocation_id, invocation_id); - }); - let mut storage = MockStorage { // mock() has no vqueue_id. status: InvocationStatus::Invoked(InFlightInvocationMetadata::mock()), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = RpcHandler::handle( + RpcContext::new(true, PartitionId::MIN, &(), &mut storage), PauseRequest { request_id: Default::default(), invocation_id, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::PauseInvocation(PauseInvocationRpcResponse::Accepted) + .await; + + assert_matches!( + decision, + Decision::NotifyInvokerAndReply { + notification: InvokerNotification::Pause(actual_invocation_id), + reply: PartitionProcessorRpcResponse::PauseInvocation( + PauseInvocationRpcResponse::Accepted + ), + } if actual_invocation_id == invocation_id ); } @@ -236,43 +209,39 @@ mod tests { async fn vqueue_invocation_proposes_pause_command() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_propose_pause_and_fence::() - .return_once_st(move |got_invocation_id, cmd, request_id, replier| { - assert_eq!(got_invocation_id, invocation_id); - let Command::PauseInvocation(bytes) = cmd else { - panic!("expected a PauseInvocation command"); - }; - let pause = PauseInvocationCommand::bilrost_decode(bytes).unwrap(); - assert_eq!(pause.invocation_id, invocation_id); - assert_eq!(pause.request_id, Some(request_id)); - replier.send(PauseInvocationRpcResponse::Accepted); - ready(()).boxed() - }); - let mut storage = MockStorage { status: InvocationStatus::Invoked(InFlightInvocationMetadata::mock_with_vqueue( invocation_id.partition_key(), )), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let request_id = PartitionProcessorRpcRequestId::new(); + let decision = RpcHandler::handle( + RpcContext::new(true, PartitionId::MIN, &(), &mut storage), PauseRequest { - request_id: Default::default(), + request_id, invocation_id, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::PauseInvocation(PauseInvocationRpcResponse::Accepted) - ); + .await; + + let Decision::Propose(RpcProposal { + partition_key, + cmd: Command::PauseInvocation(bytes), + reply_on: + ReplyOn::ApplyAndFence { + request_id: actual_request_id, + invocation_id: actual_invocation_id, + }, + }) = decision + else { + panic!("expected a fenced pause proposal"); + }; + assert_eq!(partition_key, invocation_id.partition_key()); + assert_eq!(actual_request_id, request_id); + assert_eq!(actual_invocation_id, invocation_id); + let pause = PauseInvocationCommand::bilrost_decode(bytes).unwrap(); + assert_eq!(pause.invocation_id, invocation_id); + assert_eq!(pause.request_id, Some(request_id)); } } diff --git a/crates/worker/src/partition/rpc/purge_invocation.rs b/crates/worker/src/partition/rpc/purge_invocation.rs index 291c268c3d..38d0f07c5d 100644 --- a/crates/worker/src/partition/rpc/purge_invocation.rs +++ b/crates/worker/src/partition/rpc/purge_invocation.rs @@ -13,7 +13,6 @@ use restate_types::identifiers::{InvocationId, WithPartitionKey}; use restate_types::invocation::{ IngressInvocationResponseSink, InvocationMutationResponseSink, PurgeInvocationRequest, }; -use restate_types::net::partition_processor::PurgeInvocationRpcResponse; use restate_wal_protocol::Command; pub(super) struct Request { @@ -21,34 +20,23 @@ pub(super) struct Request { pub(super) invocation_id: InvocationId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = PurgeInvocationRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { request_id, invocation_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::PurgeInvocation(PurgeInvocationRequest { - invocation_id, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::PurgeInvocation(PurgeInvocationRequest { + invocation_id, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } } diff --git a/crates/worker/src/partition/rpc/purge_journal.rs b/crates/worker/src/partition/rpc/purge_journal.rs index 7630c28038..5a85a07b6d 100644 --- a/crates/worker/src/partition/rpc/purge_journal.rs +++ b/crates/worker/src/partition/rpc/purge_journal.rs @@ -13,7 +13,6 @@ use restate_types::identifiers::{InvocationId, WithPartitionKey}; use restate_types::invocation::{ IngressInvocationResponseSink, InvocationMutationResponseSink, PurgeInvocationRequest, }; -use restate_types::net::partition_processor::PurgeInvocationRpcResponse; use restate_wal_protocol::Command; pub(super) struct Request { @@ -21,34 +20,23 @@ pub(super) struct Request { pub(super) invocation_id: InvocationId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> -{ - type Output = PurgeInvocationRpcResponse; - type Error = (); - +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> { async fn handle( self, Request { request_id, invocation_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::PurgeJournal(PurgeInvocationRequest { - invocation_id, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; - - Ok(()) + ) -> Decision { + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::PurgeJournal(PurgeInvocationRequest { + invocation_id, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } } diff --git a/crates/worker/src/partition/rpc/restart_as_new_invocation.rs b/crates/worker/src/partition/rpc/restart_as_new_invocation.rs index 79fb672343..59943b3440 100644 --- a/crates/worker/src/partition/rpc/restart_as_new_invocation.rs +++ b/crates/worker/src/partition/rpc/restart_as_new_invocation.rs @@ -9,8 +9,10 @@ // by the Apache License, Version 2.0. use super::*; + use assert2::let_assert; use opentelemetry::trace::Span; + use restate_service_protocol::codec::ProtobufRawEntryCodec as OldProtocolEntryCodec; use restate_service_protocol_v4::entry_codec::ServiceProtocolV4Codec; use restate_storage_api::invocation_status_table::{InvocationStatus, ReadInvocationStatusTable}; @@ -37,26 +39,20 @@ pub(super) struct Request { } macro_rules! bail { - ($replier:expr, $err:expr) => { + ($err:expr) => { use RestartAsNewInvocationRpcResponse::*; - $replier.send($err); - return Ok(()); + return Decision::Reply(Ok($err.into())); }; } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> where - TActuator: Actuator, TSchemas: DeploymentResolver, TStorage: ReadInvocationStatusTable + journal_table_v2::ReadJournalTable + journal_table_v1::ReadJournalTable, { - type Output = RestartAsNewInvocationRpcResponse; - type Error = (); - async fn handle( self, Request { @@ -65,16 +61,14 @@ where copy_prefix_up_to_index_included, patch_deployment_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { // Reading from a non-leader partition processor can return stale results // (e.g. NotFound for an invocation that exists on the leader) because the // follower's local store may not have replayed all log entries yet. - if !self.proposer.is_leader() { - replier.send_result(Err(PartitionProcessorRpcError::NotLeader( - self.proposer.partition_id(), + if !self.is_leader { + return Decision::Reply(Err(PartitionProcessorRpcError::NotLeader( + self.partition_id, ))); - return Ok(()); } // -- Resolve completed invocation status and input command @@ -83,30 +77,29 @@ where let completed_invocation = match self.storage.get_invocation_status(&invocation_id).await { Ok(InvocationStatus::Completed(completed_invocation)) => completed_invocation, Ok(InvocationStatus::Free) => { - bail!(replier, NotFound); + bail!(NotFound); } Ok(InvocationStatus::Scheduled(_) | InvocationStatus::Inboxed(_)) => { - bail!(replier, NotStarted); + bail!(NotStarted); } Ok(_) => { - bail!(replier, StillRunning); + bail!(StillRunning); } Err(storage_error) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal( + return Decision::Reply(Err(PartitionProcessorRpcError::Internal( storage_error.to_string(), ))); - return Ok(()); } }; // Check if there's any journal stored if completed_invocation.journal_metadata.length == 0 { - bail!(replier, MissingInput); + bail!(MissingInput); } // Check that is not a workflow if completed_invocation.invocation_target.service_ty() == ServiceType::Workflow { - bail!(replier, Unsupported); + bail!(Unsupported); } // If the invocation is using the old protocol version or no version is set at all. @@ -125,10 +118,9 @@ where { Ok(opt_entry) => opt_entry.is_none(), Err(storage_error) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal( + return Decision::Reply(Err(PartitionProcessorRpcError::Internal( storage_error.to_string(), ))); - return Ok(()); } } } @@ -147,7 +139,7 @@ where if use_old_journal_workaround { // Only copying the input entry works with this workaround! if copy_prefix_up_to_index_included > 0 { - bail!(replier, Unsupported); + bail!(Unsupported); } // Restarting an invocation that is using the journal v1 works by creating a new @@ -158,7 +150,7 @@ where patch_deployment_id, PatchDeploymentId::KeepPinned | PatchDeploymentId::PinTo { .. } ) { - bail!(replier, CannotPatchDeploymentId); + bail!(CannotPatchDeploymentId); } // Now retrieve the input command @@ -192,12 +184,11 @@ where Ok(Some(ic)) => ic, // No matching entry. Ok(None) => { - bail!(replier, MissingInput); + bail!(MissingInput); } // Failure: a storage or decoding error occurred. Err(err) => { - replier.send_result(Err(err)); - return Ok(()); + return Decision::Reply(Err(err)); } }; @@ -244,17 +235,13 @@ where // Propose and done // This path should be no longer needed once we switch to the journal v2 by default. - self.proposer - .append_and_respond_asynchronously( - invocation_id.partition_key(), - cmd, - replier, - RestartAsNewInvocationRpcResponse::Ok { new_invocation_id }, - ) - .await; - - // All good - return Ok(()); + return Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd, + reply_on: ReplyOn::Commit { + response: RestartAsNewInvocationRpcResponse::Ok { new_invocation_id }.into(), + }, + }); } // For Restart from prefix, the PP will actually execute the operation, @@ -269,7 +256,7 @@ where if copy_prefix_up_to_index_included > 0 && pinned_service_protocol_version.is_none_or(|sp| sp < ServiceProtocolVersion::V6) { - bail!(replier, Unsupported); + bail!(Unsupported); } // Figure out the deployment id, validate the protocol version constraints. @@ -291,7 +278,7 @@ where unreachable!() } }) else { - bail!(replier, DeploymentNotFound); + bail!(DeploymentNotFound); }; // Check the protocol constraints are respected. @@ -300,14 +287,14 @@ where .supported_protocol_versions .contains(&(pinned_service_protocol as i32)) { - replier.send( + return Decision::Reply(Ok( RestartAsNewInvocationRpcResponse::IncompatibleDeploymentId { pinned_protocol_version: pinned_service_protocol as i32, deployment_id: deployment.id, supported_protocol_versions: deployment.supported_protocol_versions, - }, - ); - return Ok(()); + } + .into(), + )); } Some(deployment.id) } @@ -339,10 +326,9 @@ where let cmd = match entry.decode::() { Ok(cmd) => cmd, Err(err) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal( + return Decision::Reply(Err(PartitionProcessorRpcError::Internal( err.to_string(), ))); - return Ok(()); } }; for completion_id in cmd.related_completion_ids() { @@ -352,7 +338,7 @@ where .await .is_ok_and(|b| b) { - bail!(replier, JournalCopyRangeInvalid); + bail!(JournalCopyRangeInvalid); } } } @@ -361,12 +347,13 @@ where } Ok(None) => { // Not sure what else to do here... - bail!(replier, JournalCopyRangeInvalid); + bail!(JournalCopyRangeInvalid); } Err(err) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal(err.to_string()))); - return Ok(()); + return Decision::Reply(Err(PartitionProcessorRpcError::Internal( + err.to_string(), + ))); } }; } @@ -381,31 +368,33 @@ where IngressInvocationResponseSink { request_id }, )), }); - self.proposer - .handle_rpc_proposal_command(invocation_id.partition_key(), cmd, request_id, replier) - .await; - - Ok(()) + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd, + reply_on: ReplyOn::Apply { request_id }, + }) } } #[cfg(test)] mod tests { + use std::assert_matches; use std::collections::HashMap; use std::future::ready; - use assert2::let_assert; use bytes::Bytes; - use futures::{FutureExt, Stream, StreamExt, stream}; + use futures::{Stream, StreamExt, stream}; use googletest::prelude::*; - use restate_storage_api::journal_table_v2::NotificationEntryIndex; use rstest::rstest; use test_log::test; + use journal_v2::{InputCommand, SleepCommand}; + use restate_storage_api::BudgetedReadError; use restate_storage_api::invocation_status_table::{ CompletedInvocation, InFlightInvocationMetadata, InboxedInvocation, JournalMetadata, PreFlightInvocationMetadata, ScheduledInvocation, }; + use restate_storage_api::journal_table_v2::NotificationEntryIndex; use restate_test_util::rand; use restate_test_util::rand::bytestring; use restate_types::deployment::PinnedDeployment; @@ -421,10 +410,6 @@ mod tests { use super::*; - use crate::partition::rpc::MockActuator; - use journal_v2::{InputCommand, SleepCommand}; - use restate_storage_api::BudgetedReadError; - struct MockStorage { expected_invocation_id: InvocationId, status: InvocationStatus, @@ -722,6 +707,19 @@ mod tests { } } + async fn handle( + is_leader: bool, + schemas: &R, + storage: &mut MockStorage, + request: Request, + ) -> Decision { + RpcHandler::handle( + RpcContext::new(is_leader, PartitionId::MIN, schemas, storage), + request, + ) + .await + } + #[test(restate_core::test)] async fn copy_prefix_zero_with_no_version_uses_service_invocation_command() { let old_invocation_id = InvocationId::mock_random(); @@ -729,40 +727,9 @@ mod tests { let headers = vec![Header::new("key", "value")]; let payload = rand::bytes(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); let invocation_target_clone = invocation_target.clone(); let headers_clone = vec![Header::new("key", "value")]; let payload_clone = payload.clone(); - proposer - .expect_append_and_respond_asynchronously::() - .return_once_st(move |_, cmd, _, response| { - let_assert!(Command::Invoke(service_invocation) = cmd); - assert_that!( - service_invocation, - points_to(all!( - field!(ServiceInvocation.invocation_id, not(eq(old_invocation_id))), - field!(ServiceInvocation.argument, eq(payload_clone)), - field!(ServiceInvocation.headers, eq(headers_clone)), - field!( - ServiceInvocation.invocation_target, - eq(invocation_target_clone) - ), - field!(ServiceInvocation.response_sink, none()), - field!(ServiceInvocation.submit_notification_sink, none()), - )) - ); - assert_that!( - response, - pat!(RestartAsNewInvocationRpcResponse::Ok { - new_invocation_id: eq(service_invocation.invocation_id) - }) - ); - ready(()).boxed() - }); - proposer - .expect_handle_rpc_proposal_command::() - .never(); let mut storage = MockStorage::new_with_input_v1( old_invocation_id, @@ -780,19 +747,52 @@ mod tests { headers.clone(), ); - let (tx, _rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id: old_invocation_id, copy_prefix_up_to_index_included: 0, - patch_deployment_id: Default::default(), + patch_deployment_id: PatchDeploymentId::PinToLatest, }, - Replier::new(tx), ) - .await - .unwrap(); + .await; + let (service_invocation, response) = match decision { + Decision::Propose(RpcProposal { + cmd: Command::Invoke(service_invocation), + reply_on: ReplyOn::Commit { response }, + .. + }) => (service_invocation, response), + Decision::Propose(RpcProposal { cmd, .. }) => panic!("unexpected proposal: {cmd:?}"), + Decision::Reply(reply) => panic!("unexpected reply: {reply:?}"), + Decision::NotifyInvokerAndReply { .. } => { + panic!("unexpected invoker notification") + } + }; + assert_that!( + service_invocation, + points_to(all!( + field!(ServiceInvocation.invocation_id, not(eq(old_invocation_id))), + field!(ServiceInvocation.argument, eq(payload_clone)), + field!(ServiceInvocation.headers, eq(headers_clone)), + field!( + ServiceInvocation.invocation_target, + eq(invocation_target_clone) + ), + field!(ServiceInvocation.response_sink, none()), + field!(ServiceInvocation.submit_notification_sink, none()), + )) + ); + assert_that!( + response, + pat!(PartitionProcessorRpcResponse::RestartAsNewInvocation(pat!( + RestartAsNewInvocationRpcResponse::Ok { + new_invocation_id: eq(service_invocation.invocation_id) + } + ))) + ); } #[test(restate_core::test)] @@ -800,15 +800,6 @@ mod tests { let invocation_id = InvocationId::mock_random(); let invocation_target = InvocationTarget::mock_virtual_object(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - // Completed with no pinned deployment triggers v1 workaround let status = InvocationStatus::Completed(CompletedInvocation { journal_metadata: JournalMetadata { @@ -826,25 +817,24 @@ mod tests { vec![Header::new("k", "v")], ); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::Unsupported - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::Unsupported + ) ); } @@ -855,15 +845,6 @@ mod tests { let payload = rand::bytes(); let headers = vec![Header::new("k", "v")]; - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let status = InvocationStatus::Completed(CompletedInvocation { journal_metadata: JournalMetadata { length: 1, @@ -875,25 +856,24 @@ mod tests { }); let mut storage = MockStorage::new_with_input_v1(invocation_id, status, payload, headers); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::CannotPatchDeploymentId - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::CannotPatchDeploymentId + ) ); } @@ -902,15 +882,6 @@ mod tests { let invocation_id = InvocationId::mock_random(); let invocation_target = InvocationTarget::mock_virtual_object(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let status = InvocationStatus::Completed(CompletedInvocation { journal_metadata: JournalMetadata { length: 1, @@ -927,9 +898,10 @@ mod tests { vec![Header::new("k", "v")], ); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, @@ -938,16 +910,14 @@ mod tests { id: DeploymentId::new(), }, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::CannotPatchDeploymentId - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::CannotPatchDeploymentId + ) ); } @@ -956,15 +926,6 @@ mod tests { let invocation_id = InvocationId::mock_random(); let invocation_target = InvocationTarget::mock_virtual_object(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - // Completed with journal length>0 but no v1 input present in storage let status = InvocationStatus::Completed(CompletedInvocation { journal_metadata: JournalMetadata { @@ -977,25 +938,24 @@ mod tests { }); let mut storage = MockStorage::new_without_journal(invocation_id, status); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: PatchDeploymentId::PinToLatest, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::MissingInput - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::MissingInput + ) ); } @@ -1012,45 +972,28 @@ mod tests { true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, _, _| { - assert_that!( - cmd, - pat!(Command::RestartAsNewInvocation(pat!( - RestartAsNewInvocationRequest { - copy_prefix_up_to_index_included: eq(0), - patch_deployment_id: none() - } - ))) - ); - ready(()).boxed() - }); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new( - &mut proposer, - &MockDeploymentMetadataRegistry::default(), - &mut storage, - ), + let decision = handle( + true, + &MockDeploymentMetadataRegistry::default(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - rx.assert_not_received(); + .await; + let Decision::Propose(RpcProposal { + cmd: Command::RestartAsNewInvocation(request), + reply_on: ReplyOn::Apply { .. }, + .. + }) = decision + else { + panic!("expected an on-apply restart-as-new proposal"); + }; + assert_eq!(request.copy_prefix_up_to_index_included, 0); + assert_eq!(request.patch_deployment_id, None); } #[test(restate_core::test)] @@ -1074,81 +1017,54 @@ mod tests { true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, _, _| { - assert_that!( - cmd, - pat!(Command::RestartAsNewInvocation(pat!( - RestartAsNewInvocationRequest { - copy_prefix_up_to_index_included: eq(0), - patch_deployment_id: none() - } - ))) - ); - ready(()).boxed() - }); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new( - &mut proposer, - &MockDeploymentMetadataRegistry::default(), - &mut storage, - ), + let decision = handle( + true, + &MockDeploymentMetadataRegistry::default(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - rx.assert_not_received(); + .await; + let Decision::Propose(RpcProposal { + cmd: Command::RestartAsNewInvocation(request), + reply_on: ReplyOn::Apply { .. }, + .. + }) = decision + else { + panic!("expected an on-apply restart-as-new proposal"); + }; + assert_eq!(request.copy_prefix_up_to_index_included, 0); + assert_eq!(request.patch_deployment_id, None); } #[test(restate_core::test)] async fn reply_not_found_for_unknown_invocation() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage::new_without_journal(invocation_id, Default::default()); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::NotFound - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::NotFound + ) ); } @@ -1156,15 +1072,6 @@ mod tests { async fn reply_unsupported() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage::new_with_input_v1( invocation_id, InvocationStatus::Completed(CompletedInvocation { @@ -1179,25 +1086,24 @@ mod tests { vec![Header::new("key", "value")], ); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::Unsupported - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::Unsupported + ) ); } @@ -1215,36 +1121,26 @@ mod tests { ) { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage::new_without_journal(invocation_id, status); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::StillRunning - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::StillRunning + ) ); } @@ -1264,36 +1160,26 @@ mod tests { ) { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage::new_without_journal(invocation_id, status); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + true, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::NotStarted - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::NotStarted + ) ); } @@ -1348,45 +1234,28 @@ mod tests { /* has completion */ true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, _, _| { - assert_that!( - cmd, - pat!(Command::RestartAsNewInvocation(pat!( - RestartAsNewInvocationRequest { - copy_prefix_up_to_index_included: eq(1), - patch_deployment_id: none() - } - ))) - ); - ready(()).boxed() - }); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new( - &mut proposer, - &MockDeploymentMetadataRegistry::default(), - &mut storage, - ), + let decision = handle( + true, + &MockDeploymentMetadataRegistry::default(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - rx.assert_not_received(); + .await; + let Decision::Propose(RpcProposal { + cmd: Command::RestartAsNewInvocation(request), + reply_on: ReplyOn::Apply { .. }, + .. + }) = decision + else { + panic!("expected an on-apply restart-as-new proposal"); + }; + assert_eq!(request.copy_prefix_up_to_index_included, 1); + assert_eq!(request.patch_deployment_id, None); } #[test(restate_core::test)] @@ -1412,41 +1281,28 @@ mod tests { /* has completion */ true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, _, _| { - assert_that!( - cmd, - pat!(Command::RestartAsNewInvocation(pat!( - RestartAsNewInvocationRequest { - copy_prefix_up_to_index_included: eq(1), - patch_deployment_id: some(eq(latest_id)) - } - ))) - ); - ready(()).boxed() - }); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, ®istry, &mut storage), + let decision = handle( + true, + ®istry, + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::PinToLatest, }, - Replier::new(tx), ) - .await - .unwrap(); - - rx.assert_not_received(); + .await; + let Decision::Propose(RpcProposal { + cmd: Command::RestartAsNewInvocation(request), + reply_on: ReplyOn::Apply { .. }, + .. + }) = decision + else { + panic!("expected an on-apply restart-as-new proposal"); + }; + assert_eq!(request.copy_prefix_up_to_index_included, 1); + assert_eq!(request.patch_deployment_id, Some(latest_id)); } #[test(restate_core::test)] @@ -1470,38 +1326,28 @@ mod tests { true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, ®istry, &mut storage), + let decision = handle( + true, + ®istry, + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::PinTo { id }, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::IncompatibleDeploymentId { - pinned_protocol_version: pinned.service_protocol_version as i32, - deployment_id: id, - supported_protocol_versions: 1..=2, - } - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::IncompatibleDeploymentId { + pinned_protocol_version: pinned.service_protocol_version as i32, + deployment_id: id, + supported_protocol_versions: 1..=2, + } + ) ); } @@ -1521,31 +1367,24 @@ mod tests { true, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, ®istry, &mut storage), + let decision = handle( + true, + ®istry, + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::PinTo { id: some_id }, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::DeploymentNotFound - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::DeploymentNotFound + ) ); } @@ -1562,35 +1401,24 @@ mod tests { /* has completion */ false, ); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new( - &mut proposer, - &MockDeploymentMetadataRegistry::default(), - &mut storage, - ), + let decision = handle( + true, + &MockDeploymentMetadataRegistry::default(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::JournalCopyRangeInvalid - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::JournalCopyRangeInvalid + ) ); } @@ -1603,35 +1431,24 @@ mod tests { let mut storage = MockStorage::new_without_journal(invocation_id, InvocationStatus::Completed(completed)); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new( - &mut proposer, - &MockDeploymentMetadataRegistry::default(), - &mut storage, - ), + let decision = handle( + true, + &MockDeploymentMetadataRegistry::default(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 1, patch_deployment_id: PatchDeploymentId::KeepPinned, }, - Replier::new(tx), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::RestartAsNewInvocation( - RestartAsNewInvocationRpcResponse::Unsupported - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + RestartAsNewInvocationRpcResponse::Unsupported + ) ); } @@ -1639,37 +1456,24 @@ mod tests { async fn reply_not_leader_when_not_leader() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(false); - proposer - .expect_partition_id() - .return_const(PartitionId::from(0)); - proposer - .expect_append_and_respond_asynchronously::() - .never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage::new_without_journal(invocation_id, Default::default()); - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), + let decision = handle( + false, + &(), + &mut storage, Request { request_id: Default::default(), invocation_id, copy_prefix_up_to_index_included: 0, patch_deployment_id: Default::default(), }, - Replier::new(tx), ) - .await - .unwrap(); + .await; - assert!(matches!( - rx.recv().await, - Err(PartitionProcessorRpcError::NotLeader(_)) - )); + assert_matches!( + decision, + Decision::Reply(Err(PartitionProcessorRpcError::NotLeader(_))) + ); } } diff --git a/crates/worker/src/partition/rpc/resume_invocation.rs b/crates/worker/src/partition/rpc/resume_invocation.rs index 5d04e9a417..161682000d 100644 --- a/crates/worker/src/partition/rpc/resume_invocation.rs +++ b/crates/worker/src/partition/rpc/resume_invocation.rs @@ -25,17 +25,12 @@ pub(super) struct Request { pub(super) update_deployment_id: PatchDeploymentId, } -impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler - for RpcContext<'a, TActuator, TSchemas, TStorage> +impl<'a, TSchemas, TStorage> RpcHandler for RpcContext<'a, TSchemas, TStorage> where - TActuator: Actuator, // Needed for the non-VQueue path, which resolves the deployment here (see `handle`). TSchemas: DeploymentResolver, TStorage: ReadInvocationStatusTable, { - type Output = ResumeInvocationRpcResponse; - type Error = (); - async fn handle( self, Request { @@ -43,16 +38,14 @@ where invocation_id, update_deployment_id, }: Request, - replier: Replier, - ) -> Result<(), Self::Error> { + ) -> Decision { // Reading from a non-leader partition processor can return stale results // (e.g. NotFound for an invocation that exists on the leader) because the // follower's local store may not have replayed all log entries yet. - if !self.proposer.is_leader() { - replier.send_result(Err(PartitionProcessorRpcError::NotLeader( - self.proposer.partition_id(), + if !self.is_leader { + return Decision::Reply(Err(PartitionProcessorRpcError::NotLeader( + self.partition_id, ))); - return Ok(()); } // -- Figure out the invocation status @@ -66,31 +59,31 @@ where // is no running attempt to fence, so use the plain proposal path -- unlike // pause's `propose_pause_and_fence`. Writing the new `update_deployment_id` field // is safe here: vqueues being enabled implies a cluster min version >= 1.7.0. - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::ResumeInvocation(ResumeInvocationRequest { - invocation_id, - update_deployment_id: Some(update_deployment_id), - update_pinned_deployment_id: None, - run_at: None, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::ResumeInvocation(ResumeInvocationRequest { + invocation_id, + update_deployment_id: Some(update_deployment_id), + update_pinned_deployment_id: None, + run_at: None, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } else { // Legacy invoker-owned path: there is a live attempt, so the pinned deployment // cannot be patched. Poke the invoker to retry now, if possible. if !matches!(update_deployment_id, PatchDeploymentId::KeepPinned) { - replier.send(ResumeInvocationRpcResponse::CannotPatchDeploymentId); - return Ok(()); + return Decision::Reply(Ok( + ResumeInvocationRpcResponse::CannotPatchDeploymentId.into(), + )); + } + Decision::NotifyInvokerAndReply { + notification: InvokerNotification::RetryNow(invocation_id), + reply: ResumeInvocationRpcResponse::Ok.into(), } - self.proposer.notify_invoker_to_retry_now(invocation_id); - replier.send(ResumeInvocationRpcResponse::Ok); } } Ok(InvocationStatus::Suspended { metadata, .. }) @@ -99,23 +92,19 @@ where // VQueue path: forward the unresolved patch; the apply path resolves it against // the status as of the command's log position. Safe to write the new // `update_deployment_id` field -- vqueues imply a cluster min version >= 1.7.0. - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::ResumeInvocation(ResumeInvocationRequest { - invocation_id, - update_deployment_id: Some(update_deployment_id), - update_pinned_deployment_id: None, - run_at: None, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; - return Ok(()); + return Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::ResumeInvocation(ResumeInvocationRequest { + invocation_id, + update_deployment_id: Some(update_deployment_id), + update_pinned_deployment_id: None, + run_at: None, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }); } // Non-VQueue path: a pre-1.7.0 node may still apply this command and does not know @@ -130,55 +119,48 @@ where ) { Ok(resolved) => resolved, Err(response) => { - replier.send(response.into()); - return Ok(()); + return Decision::Reply(Ok( + ResumeInvocationRpcResponse::from(response).into() + )); } }; - self.proposer - .handle_rpc_proposal_command( - invocation_id.partition_key(), - Command::ResumeInvocation(ResumeInvocationRequest { - invocation_id, - update_deployment_id: None, - update_pinned_deployment_id, - run_at: None, - response_sink: Some(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id }, - )), - }), - request_id, - replier, - ) - .await; + Decision::Propose(RpcProposal { + partition_key: invocation_id.partition_key(), + cmd: Command::ResumeInvocation(ResumeInvocationRequest { + invocation_id, + update_deployment_id: None, + update_pinned_deployment_id, + run_at: None, + response_sink: Some(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id }, + )), + }), + reply_on: ReplyOn::Apply { request_id }, + }) } Ok(InvocationStatus::Scheduled(_)) | Ok(InvocationStatus::Inboxed(_)) => { - replier.send(ResumeInvocationRpcResponse::NotStarted); + Decision::Reply(Ok(ResumeInvocationRpcResponse::NotStarted.into())) } Ok(InvocationStatus::Completed(_)) => { - replier.send(ResumeInvocationRpcResponse::Completed); + Decision::Reply(Ok(ResumeInvocationRpcResponse::Completed.into())) } Ok(InvocationStatus::Free) => { - replier.send(ResumeInvocationRpcResponse::NotFound); + Decision::Reply(Ok(ResumeInvocationRpcResponse::NotFound.into())) } - Err(storage_error) => { - replier.send_result(Err(PartitionProcessorRpcError::Internal( - storage_error.to_string(), - ))); - } - }; - - Ok(()) + Err(storage_error) => Decision::Reply(Err(PartitionProcessorRpcError::Internal( + storage_error.to_string(), + ))), + } } } #[cfg(test)] mod tests { - use super::*; + use std::assert_matches; + use std::future::ready; - use crate::partition::rpc::MockActuator; use assert2::let_assert; - use futures::FutureExt; use googletest::prelude::*; use restate_storage_api::invocation_status_table::{ CompletedInvocation, InFlightInvocationMetadata, InboxedInvocation, @@ -192,9 +174,10 @@ mod tests { use restate_types::schema::deployment::test_util::MockDeploymentMetadataRegistry; use restate_types::service_protocol::ServiceProtocolVersion; use rstest::rstest; - use std::future::ready; use test_log::test; + use super::*; + struct MockStorage { expected_invocation_id: InvocationId, status: InvocationStatus, @@ -217,39 +200,51 @@ mod tests { } } + async fn handle( + is_leader: bool, + schemas: &R, + storage: &mut MockStorage, + request_id: PartitionProcessorRpcRequestId, + invocation_id: InvocationId, + update_deployment_id: PatchDeploymentId, + ) -> Decision { + RpcHandler::handle( + RpcContext::new(is_leader, PartitionId::MIN, schemas, storage), + Request { + request_id, + invocation_id, + update_deployment_id, + }, + ) + .await + } + #[test(restate_core::test)] async fn reply_ok_when_invoked() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_notify_invoker_to_retry_now() - .return_once_st(move |got_invocation_id| { - assert_eq!(got_invocation_id, invocation_id); - }); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: InvocationStatus::Invoked(InFlightInvocationMetadata::mock()), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation(ResumeInvocationRpcResponse::Ok) + .await; + assert_matches!( + decision, + Decision::NotifyInvokerAndReply { + notification: InvokerNotification::RetryNow(actual_invocation_id), + reply: PartitionProcessorRpcResponse::ResumeInvocation( + ResumeInvocationRpcResponse::Ok + ), + } if actual_invocation_id == invocation_id ); } @@ -259,22 +254,6 @@ mod tests { async fn vqueue_invoked_proposes_resume_command() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - // The invoker must NOT be poked for a VQueue invocation. - proposer.expect_notify_invoker_to_retry_now().never(); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, _request_id, replier| { - // The command shape (response_sink, deployment patching) is covered by - // `propose_resume_command_on_paused_and_suspended`; here we only care that the - // Invoked+VQueue path proposes ResumeInvocation rather than poking the invoker. - let_assert!(Command::ResumeInvocation(request) = cmd); - assert_eq!(request.invocation_id, invocation_id); - replier.send(ResumeInvocationRpcResponse::Ok); - ready(()).boxed() - }); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: InvocationStatus::Invoked(InFlightInvocationMetadata::mock_with_vqueue( @@ -282,23 +261,23 @@ mod tests { )), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation(ResumeInvocationRpcResponse::Ok) + .await; + let_assert!( + Decision::Propose(RpcProposal { + cmd: Command::ResumeInvocation(request), + reply_on: ReplyOn::Apply { .. }, + .. + }) = decision ); + assert_eq!(request.invocation_id, invocation_id); } #[rstest] @@ -315,49 +294,42 @@ mod tests { ) { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .return_once_st(move |_, cmd, request_id, replier| { - let_assert!(Command::ResumeInvocation(resume_invocation_request) = cmd); - assert_that!( - resume_invocation_request, - all!( - field!(ResumeInvocationRequest.invocation_id, eq(invocation_id)), - field!( - ResumeInvocationRequest.response_sink, - some(eq(InvocationMutationResponseSink::Ingress( - IngressInvocationResponseSink { request_id } - ))) - ), - ) - ); - replier.send(ResumeInvocationRpcResponse::Ok); - ready(()).boxed() - }); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status, }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let request_id = PartitionProcessorRpcRequestId::new(); + let decision = handle( + true, + &(), + &mut storage, + request_id, + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation(ResumeInvocationRpcResponse::Ok) + .await; + let_assert!( + Decision::Propose(RpcProposal { + cmd: Command::ResumeInvocation(resume_invocation_request), + reply_on: ReplyOn::Apply { + request_id: actual_request_id, + }, + .. + }) = decision + ); + assert_eq!(actual_request_id, request_id); + assert_that!( + resume_invocation_request, + all!( + field!(ResumeInvocationRequest.invocation_id, eq(invocation_id)), + field!( + ResumeInvocationRequest.response_sink, + some(eq(InvocationMutationResponseSink::Ingress( + IngressInvocationResponseSink { request_id } + ))) + ), + ) ); } @@ -365,33 +337,26 @@ mod tests { async fn reply_not_found_for_unknown_invocation() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: Default::default(), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation(ResumeInvocationRpcResponse::NotFound) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + ResumeInvocationRpcResponse::NotFound + ) ); } @@ -399,33 +364,26 @@ mod tests { async fn reply_completed() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: InvocationStatus::Completed(CompletedInvocation::mock_neo()), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation(ResumeInvocationRpcResponse::Completed) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + ResumeInvocationRpcResponse::Completed + ) ); } @@ -445,35 +403,26 @@ mod tests { ) { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status, }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation( - ResumeInvocationRpcResponse::NotStarted - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + ResumeInvocationRpcResponse::NotStarted + ) ); } @@ -498,12 +447,6 @@ mod tests { schemas.mock_latest_service(invocation_target.service_name(), dep.id); } - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let metadata = InFlightInvocationMetadata { invocation_target, pinned_deployment: Some(PinnedDeployment::new(DeploymentId::new(), pinned_version)), @@ -521,7 +464,6 @@ mod tests { }, }; - let (tx, rx) = Reciprocal::mock(); let update_deployment_id = if pin_to_specific { PatchDeploymentId::PinTo { id: expected_deployment_id, @@ -529,22 +471,23 @@ mod tests { } else { PatchDeploymentId::PinToLatest }; - RpcHandler::handle( - RpcContext::new(&mut proposer, &schemas, &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id, - }, - Replier::new(tx), + let decision = handle( + true, + &schemas, + &mut storage, + Default::default(), + invocation_id, + update_deployment_id, ) - .await - .unwrap(); + .await; // Non-VQueue path resolves the deployment in the RPC and rejects the incompatible pin // synchronously (no command proposed). assert_eq!( - rx.recv().await.unwrap(), + match decision { + Decision::Reply(Ok(response)) => response, + _ => panic!("expected an immediate successful reply"), + }, PartitionProcessorRpcResponse::ResumeInvocation( ResumeInvocationRpcResponse::IncompatibleDeploymentId { pinned_protocol_version: i32::from(pinned_version), @@ -561,36 +504,26 @@ mod tests { async fn legacy_invoked_rejects_deployment_patch() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(true); - proposer.expect_notify_invoker_to_retry_now().never(); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: InvocationStatus::Invoked(InFlightInvocationMetadata::mock()), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: PatchDeploymentId::PinToLatest, - }, - Replier::new(tx), + let decision = handle( + true, + &(), + &mut storage, + Default::default(), + invocation_id, + PatchDeploymentId::PinToLatest, ) - .await - .unwrap(); - - assert_eq!( - rx.recv().await.unwrap(), - PartitionProcessorRpcResponse::ResumeInvocation( - ResumeInvocationRpcResponse::CannotPatchDeploymentId - ) + .await; + assert_matches!( + decision, + Decision::Reply(Ok(response)) + if response == PartitionProcessorRpcResponse::from( + ResumeInvocationRpcResponse::CannotPatchDeploymentId + ) ); } @@ -598,36 +531,24 @@ mod tests { async fn reply_not_leader_when_not_leader() { let invocation_id = InvocationId::mock_random(); - let mut proposer = MockActuator::new(); - proposer.expect_is_leader().return_const(false); - proposer - .expect_partition_id() - .return_const(PartitionId::from(0)); - proposer - .expect_handle_rpc_proposal_command::() - .never(); - let mut storage = MockStorage { expected_invocation_id: invocation_id, status: Default::default(), }; - let (tx, rx) = Reciprocal::mock(); - RpcHandler::handle( - RpcContext::new(&mut proposer, &(), &mut storage), - Request { - request_id: Default::default(), - invocation_id, - update_deployment_id: Default::default(), - }, - Replier::new(tx), + let decision = handle( + false, + &(), + &mut storage, + Default::default(), + invocation_id, + Default::default(), ) - .await - .unwrap(); + .await; - assert!(matches!( - rx.recv().await, - Err(PartitionProcessorRpcError::NotLeader(_)) - )); + assert_matches!( + decision, + Decision::Reply(Err(PartitionProcessorRpcError::NotLeader(_))) + ); } }