From 77f7f5007135b4e3773ab4f2660971beb44a5c7a Mon Sep 17 00:00:00 2001 From: MohamedBassem Date: Wed, 15 Jul 2026 17:36:32 +0100 Subject: [PATCH] [worker][refactor] change rpc handlers to return their next steps instead of taking actions inline As part of the work on centeralizes writes to the self proposer inside one scheduler, this PR refactors the RPC handlers to return their decisions to the main PP loop where they get executed there. This is in prep for shipping the proposals to a channel in the next PR. This PR is mostly mechanical and does the following: 1. Changes the handlers to no longer take the "Actuator" trait and instead return their `Decision` from the function. 2. The decision can be either: Propose a command (with its different modes), reply now, or the pre-vqueues invoker poking. Tests are much simpler as they no longer need to mock the leadership state. Only meaningful change was a bug in the restart_as_new test `copy_prefix_zero_with_no_version_uses_service_invocation_command` where it was using the default `patch_deployment_id` and it didn't actually hit the assertions inside the mock handlers as it was failing with `CannotPatchDeploymentId`. The tests where not setting expectations on num calls to the mock, so it didn't fail when the call to `append_and_respond_asynchronously` didn't happen. So this was fixed along the way by using `PatchDeploymentId::PinToLatest` to hit the correct assertions. This PR anyways removes `mockall` altogether from that worker crate. --- Cargo.lock | 1 - crates/worker/Cargo.toml | 1 - crates/worker/src/partition/mod.rs | 67 +- .../src/partition/rpc/append_invocation.rs | 204 ++--- .../rpc/append_invocation_response.rs | 27 +- .../worker/src/partition/rpc/append_signal.rs | 33 +- .../src/partition/rpc/cancel_invocation.rs | 38 +- .../partition/rpc/get_invocation_output.rs | 49 +- .../src/partition/rpc/kill_invocation.rs | 38 +- crates/worker/src/partition/rpc/mod.rs | 339 +++----- .../src/partition/rpc/pause_invocation.rs | 181 ++-- .../src/partition/rpc/purge_invocation.rs | 36 +- .../worker/src/partition/rpc/purge_journal.rs | 36 +- .../rpc/restart_as_new_invocation.rs | 810 +++++++----------- .../src/partition/rpc/resume_invocation.rs | 529 +++++------- 15 files changed, 929 insertions(+), 1460 deletions(-) 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/mod.rs b/crates/worker/src/partition/mod.rs index 9655847c0b..339a15f851 100644 --- a/crates/worker/src/partition/mod.rs +++ b/crates/worker/src/partition/mod.rs @@ -86,6 +86,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::{ @@ -675,16 +676,62 @@ where body: PartitionProcessorRpcRequest, schemas: &Schema, ) { - 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(rpc::RpcProposal { + partition_key, + cmd, + reply_on, + }) => match reply_on { + rpc::ReplyOn::Apply { request_id } => { + self.leadership_state + .handle_rpc_proposal_command(request_id, response_tx, partition_key, cmd) + .await; + } + rpc::ReplyOn::Commit { response } => { + self.leadership_state + .append_and_respond_asynchronously( + partition_key, + cmd, + response_tx, + response, + ) + .await; + } + rpc::ReplyOn::ApplyAndFence { + request_id, + invocation_id, + } => { + self.leadership_state + .propose_pause_and_fence(request_id, response_tx, invocation_id, cmd) + .await; + } + }, + 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( 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(_))) + ); } }