From 77f7f5007135b4e3773ab4f2660971beb44a5c7a Mon Sep 17 00:00:00 2001 From: MohamedBassem Date: Wed, 15 Jul 2026 17:36:32 +0100 Subject: [PATCH 1/2] [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(_))) + ); } } From b7b8cf39c233aea7aa5e67cc45c519cf3db4ffc8 Mon Sep 17 00:00:00 2001 From: MohamedBassem Date: Thu, 16 Jul 2026 12:33:15 +0100 Subject: [PATCH 2/2] [worker] Funnel all command proposals through LeaderState::run --- .../src/partition/leadership/leader_state.rs | 192 ++++++++++-------- crates/worker/src/partition/leadership/mod.rs | 184 +++++++---------- crates/worker/src/partition/mod.rs | 88 +++----- 3 files changed, 214 insertions(+), 250 deletions(-) diff --git a/crates/worker/src/partition/leadership/leader_state.rs b/crates/worker/src/partition/leadership/leader_state.rs index 0f06046f59..ec6158a05d 100644 --- a/crates/worker/src/partition/leadership/leader_state.rs +++ b/crates/worker/src/partition/leadership/leader_state.rs @@ -59,8 +59,11 @@ use restate_worker_api::{SchedulerStatusEntry, UserLimitCounterEntry}; use crate::metric_definitions::{PARTITION_HANDLE_LEADER_ACTIONS, USAGE_LEADER_ACTION_COUNT}; use crate::partition::cleaner::{CleanerEffect, CleanerHandle}; use crate::partition::leadership::self_proposer::SelfProposer; -use crate::partition::leadership::{ActionEffect, Error, InvokerStream, TimerService}; +use crate::partition::leadership::{ + AppendCallback, Error, InvokerStream, LeaderEvent, RpcReciprocal, TimerService, +}; use crate::partition::processor::{FsmAccess, Processor}; +use crate::partition::rpc::{ReplyOn, RpcProposal}; use crate::partition::shuffle; use crate::partition::shuffle::HintSender; use crate::partition::state_machine::Action; @@ -69,9 +72,7 @@ use crate::partition_processor_manager::LeaderQueryGuard; use super::durability_tracker::DurabilityTracker; const BATCH_READY_UP_TO: usize = 10; - -type RpcReciprocal = - Reciprocal>>; +const NETWORK_EVENTS_BUFFER_SIZE: usize = 1; pub struct LeaderState { pub(crate) partition_id: PartitionId, @@ -111,6 +112,8 @@ pub struct LeaderState { cleaner_handle: CleanerHandle, trimmer_task_id: TaskId, durability_tracker: DurabilityTracker, + network_events_tx: tokio::sync::mpsc::Sender, + network_events_stream: ReceiverStream, // Unregisters the leader-query registry entry on drop. Must live as long as // the partition processor's select! is willing to serve scheduler queries. _leader_query_guard: LeaderQueryGuard, @@ -141,6 +144,8 @@ impl LeaderState { leader_query_guard: LeaderQueryGuard, rule_book_rx: tokio::sync::watch::Receiver>, ) -> Self { + let (network_events_tx, network_events_rx) = + tokio::sync::mpsc::channel(NETWORK_EVENTS_BUFFER_SIZE); LeaderState { partition_id, leader_epoch, @@ -165,6 +170,8 @@ impl LeaderState { invoker_stream: invoker_rx, shuffle_stream: ReceiverStream::new(shuffle_rx), durability_tracker, + network_events_tx, + network_events_stream: ReceiverStream::new(network_events_rx), _leader_query_guard: leader_query_guard, } } @@ -189,20 +196,28 @@ impl LeaderState { self.scheduler.scan_user_limit_counters(keys.start()) } - /// Runs the leader specific task which is the awaiting of action effects and the monitoring - /// of unmanaged tasks. + pub(super) fn try_reserve_network_event_permit( + &self, + ) -> Result< + tokio::sync::mpsc::OwnedPermit, + tokio::sync::mpsc::error::TrySendError>, + > { + self.network_events_tx.clone().try_reserve_owned() + } + + /// Runs the leader-specific task by awaiting events and monitoring unmanaged tasks. /// /// Important: The future needs to be cancellation safe since it is polled as a tokio::select /// arm! pub async fn run( &mut self, ctx: impl Processor + HasVQueues, - ) -> Result, Error> { + ) -> Result, Error> { let timer_stream = std::pin::pin!(stream::unfold( &mut self.timer_service, |timer_service| async { let timer_value = timer_service.as_mut().next_timer().await; - Some((ActionEffect::Timer(timer_value), timer_service)) + Some((LeaderEvent::Timer(timer_value), timer_service)) } )); let vqueue_metas = ctx.vqueues(); @@ -212,7 +227,7 @@ impl LeaderState { let scheduler_stream = std::pin::pin!(stream::unfold(&mut self.scheduler, |scheduler| async { match scheduler.schedule_next(vqueue_metas).await { - Ok(decisions) => Some((ActionEffect::Scheduler(decisions), scheduler)), + Ok(decisions) => Some((LeaderEvent::Scheduler(decisions), scheduler)), Err(e) => { error!("Fatal error when polling scheduler: {e}"); None @@ -227,7 +242,7 @@ impl LeaderState { std::future::ready( Some(Metadata::with_current(|m| m.schema())) .filter(|schema| schema.version() > current_version) - .map(|schema| ActionEffect::UpsertSchema(schema.clone())), + .map(|schema| LeaderEvent::UpsertSchema(schema.clone())), ) }); @@ -236,19 +251,21 @@ impl LeaderState { // newer than the partition's current in-memory book. let current_version = ctx.fsm().rule_book().version(); std::future::ready( - (book.version() > current_version).then(|| ActionEffect::UpsertRuleBook(book)), + (book.version() > current_version).then(|| LeaderEvent::UpsertRuleBook(book)), ) }); - let invoker_stream = (&mut self.invoker_stream).map(ActionEffect::Invoker); - let shuffle_stream = (&mut self.shuffle_stream).map(ActionEffect::Shuffle); - let cleaner_stream = self.cleaner_handle.effects().map(ActionEffect::Cleaner); + let invoker_stream = (&mut self.invoker_stream).map(LeaderEvent::Invoker); + let shuffle_stream = (&mut self.shuffle_stream).map(LeaderEvent::Shuffle); + let cleaner_stream = self.cleaner_handle.effects().map(LeaderEvent::Cleaner); let dur_tracker_stream = - (&mut self.durability_tracker).map(ActionEffect::PartitionMaintenance); + (&mut self.durability_tracker).map(LeaderEvent::PartitionMaintenance); let awaiting_rpc_self_propose_stream = - (&mut self.awaiting_rpc_self_propose).map(|_| ActionEffect::AwaitingRpcSelfProposeDone); + (&mut self.awaiting_rpc_self_propose).map(|_| LeaderEvent::AwaitingRpcSelfProposeDone); + + let network_event_stream = &mut self.network_events_stream; let all_streams = futures::stream_select!( scheduler_stream, @@ -259,7 +276,8 @@ impl LeaderState { awaiting_rpc_self_propose_stream, dur_tracker_stream, schema_stream, - rule_book_stream + rule_book_stream, + network_event_stream ); let mut all_streams = all_streams.ready_chunks(BATCH_READY_UP_TO); @@ -351,16 +369,30 @@ impl LeaderState { for fut in self.awaiting_rpc_self_propose.iter_mut() { fut.fail_with_lost_leadership(self.partition_id); } + // No network processing permit can outlive the partition processor's select iteration, + // so after closing the receiver all buffered events are immediately available. + self.network_events_stream.close(); + while let Some(Some(event)) = self.network_events_stream.next().now_or_never() { + match event { + LeaderEvent::RpcProposal { reciprocal, .. } => reciprocal.send(Err( + PartitionProcessorRpcError::LostLeadership(self.partition_id), + )), + LeaderEvent::IngestRecords { callback, .. } => callback(Err( + PartitionProcessorRpcError::LostLeadership(self.partition_id), + )), + _ => unreachable!("only network events are buffered"), + } + } } - pub async fn handle_action_effects( + pub async fn handle_events( &mut self, - action_effects: impl IntoIterator, + events: impl IntoIterator, ) -> Result<(), Error> { let mut arena = BytesMut::with_capacity(128 * 1024); - for effect in action_effects { - match effect { - ActionEffect::Scheduler(decisions) => { + for event in events { + match event { + LeaderEvent::Scheduler(decisions) => { let Decisions { qids, num_run, @@ -398,7 +430,7 @@ impl LeaderState { .self_propose_many(commands.into_iter()) .await?; } - ActionEffect::PartitionMaintenance(partition_durability) => { + LeaderEvent::PartitionMaintenance(partition_durability) => { // based on configuration, whether to consider partition-local durability in // the replica-set as a sufficient source of durability, or only snapshots. self.self_proposer @@ -408,7 +440,7 @@ impl LeaderState { ) .await?; } - ActionEffect::Invoker(fenced) => { + LeaderEvent::Invoker(fenced) => { let invocation_id = fenced.effect.invocation_id; // Fence stale effects at write time: only self-propose if the effect carries // the token of the invocation's *current* attempt. A mismatch (or a missing @@ -434,7 +466,7 @@ impl LeaderState { ); } } - ActionEffect::Shuffle(outbox_truncation) => { + LeaderEvent::Shuffle(outbox_truncation) => { // todo: Until we support partition splits we need to get rid of outboxes or introduce partition // specific destination messages that are identified by a partition_id self.self_proposer @@ -444,12 +476,12 @@ impl LeaderState { ) .await?; } - ActionEffect::Timer(timer) => { + LeaderEvent::Timer(timer) => { self.self_proposer .self_propose(timer.invocation_id().partition_key(), Command::Timer(timer)) .await?; } - ActionEffect::Cleaner(effect) => { + LeaderEvent::Cleaner(effect) => { let (invocation_id, cmd) = match effect { CleanerEffect::PurgeJournal(invocation_id) => ( invocation_id, @@ -471,7 +503,7 @@ impl LeaderState { .self_propose(invocation_id.partition_key(), cmd) .await?; } - ActionEffect::UpsertSchema(schema) => { + LeaderEvent::UpsertSchema(schema) => { if SemanticRestateVersion::current() .is_equal_or_newer_than(&RESTATE_VERSION_1_7_0) { @@ -488,7 +520,7 @@ impl LeaderState { .await?; } } - ActionEffect::UpsertRuleBook(rule_book) => { + LeaderEvent::UpsertRuleBook(rule_book) => { let cmd = restate_wal_protocol::control::UpsertRuleBookCommand { rule_book }; arena.reserve(cmd.encoded_len()); // safe to unwrap because we reserved enough space @@ -504,16 +536,52 @@ impl LeaderState { ) .await?; } - ActionEffect::AwaitingRpcSelfProposeDone => { + LeaderEvent::AwaitingRpcSelfProposeDone => { // Nothing to do here } + LeaderEvent::RpcProposal { + proposal, + reciprocal, + } => { + self.handle_rpc_proposal(proposal, reciprocal).await; + } + LeaderEvent::IngestRecords { records, callback } => { + self.forward_many_with_callback(records.into_iter(), callback) + .await; + } } } Ok(()) } - pub async fn handle_rpc_proposal_command( + async fn handle_rpc_proposal(&mut self, proposal: RpcProposal, reciprocal: RpcReciprocal) { + let RpcProposal { + partition_key, + cmd, + reply_on, + } = proposal; + + match reply_on { + ReplyOn::Apply { request_id } => { + self.handle_rpc_proposal_command(request_id, reciprocal, partition_key, cmd) + .await; + } + ReplyOn::Commit { response } => { + self.append_and_respond_asynchronously(partition_key, cmd, reciprocal, response) + .await; + } + ReplyOn::ApplyAndFence { + request_id, + invocation_id, + } => { + self.propose_pause_and_fence(request_id, reciprocal, invocation_id, cmd) + .await; + } + } + } + + async fn handle_rpc_proposal_command( &mut self, request_id: PartitionProcessorRpcRequestId, reciprocal: Reciprocal< @@ -554,7 +622,7 @@ impl LeaderState { /// dropped at write time with no pause in the committed log. After a successful append the /// pause's LSN is fixed, and because the leader self-proposes from a single task, every later /// invoker-effect self-propose observes the cleared map and is fenced. - pub async fn propose_pause_and_fence( + async fn propose_pause_and_fence( &mut self, request_id: PartitionProcessorRpcRequestId, reciprocal: Reciprocal< @@ -595,7 +663,7 @@ impl LeaderState { /// Records appended this way are never filtered by the dedup mechanism during leadership /// transitions, making this safe for fire-and-forget ingress commands (signals, invocation /// responses). - pub async fn append_and_respond_asynchronously( + async fn append_and_respond_asynchronously( &mut self, partition_key: PartitionKey, cmd: Command, @@ -610,22 +678,20 @@ impl LeaderState { Ok(commit_token) => { self.awaiting_rpc_self_propose.push(SelfAppendFuture::new( commit_token, - |result: Result<(), PartitionProcessorRpcError>| { + Box::new(|result: Result<(), PartitionProcessorRpcError>| { reciprocal.send(result.map(|_| success_response)); - }, + }), )); } Err(e) => reciprocal.send(Err(PartitionProcessorRpcError::Internal(e.to_string()))), } } - pub async fn forward_many_with_callback( + async fn forward_many_with_callback( &mut self, records: impl ExactSizeIterator, - callback: F, - ) where - F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static, - { + callback: AppendCallback, + ) { match self .self_proposer .forward_many_with_notification(records) @@ -894,56 +960,22 @@ impl LeaderState { } } -trait CallbackInner: Send + Sync + 'static { - fn call(self: Box, result: Result<(), PartitionProcessorRpcError>); -} - -impl CallbackInner for F -where - F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static, -{ - fn call(self: Box, result: Result<(), PartitionProcessorRpcError>) { - self(result) - } -} - -struct Callback { - inner: Box, -} - -impl Callback { - fn call(self, result: Result<(), PartitionProcessorRpcError>) { - self.inner.call(result); - } -} - -impl From for Callback -where - I: CallbackInner, -{ - fn from(value: I) -> Self { - Self { - inner: Box::new(value), - } - } -} - struct SelfAppendFuture { commit_token: CommitToken, - callback: Option, + callback: Option, } impl SelfAppendFuture { - fn new(commit_token: CommitToken, callback: impl Into) -> Self { + fn new(commit_token: CommitToken, callback: AppendCallback) -> Self { Self { commit_token, - callback: Some(callback.into()), + callback: Some(callback), } } fn fail_with_internal(&mut self) { if let Some(callback) = self.callback.take() { - callback.call(Err(PartitionProcessorRpcError::Internal( + callback(Err(PartitionProcessorRpcError::Internal( "error when proposing to bifrost".to_string(), ))); } @@ -951,7 +983,7 @@ impl SelfAppendFuture { fn fail_with_lost_leadership(&mut self, this_partition_id: PartitionId) { if let Some(callback) = self.callback.take() { - callback.call(Err(PartitionProcessorRpcError::LostLeadership( + callback(Err(PartitionProcessorRpcError::LostLeadership( this_partition_id, ))); } @@ -959,7 +991,7 @@ impl SelfAppendFuture { fn succeed_with_appended(&mut self) { if let Some(callback) = self.callback.take() { - callback.call(Ok(())) + callback(Ok(())) } } } diff --git a/crates/worker/src/partition/leadership/mod.rs b/crates/worker/src/partition/leadership/mod.rs index 0e3868ce13..17727c5294 100644 --- a/crates/worker/src/partition/leadership/mod.rs +++ b/crates/worker/src/partition/leadership/mod.rs @@ -49,7 +49,6 @@ use restate_types::cluster::cluster_state::RunMode; use restate_types::config::Configuration; use restate_types::errors::GenericError; use restate_types::identifiers::{InvocationId, LeaderEpoch, PartitionId}; -use restate_types::identifiers::{PartitionKey, PartitionProcessorRpcRequestId}; use restate_types::invocation::FencingToken; use restate_types::live::LiveLoadExt; use restate_types::logs::Keys; @@ -91,10 +90,14 @@ use crate::partition::state_machine::Action; use crate::partition::types::InvokerEffect; use super::node::NodeContext; -use super::processor::*; +use super::{processor::*, rpc}; type TimerService = restate_timer::TimerService; type InvokerStream = ReceiverStream; +type RpcReciprocal = + Reciprocal>>; +pub(super) type AppendCallback = + Box) + Send + Sync + 'static>; #[derive(Debug, thiserror::Error)] pub(crate) enum Error { @@ -145,8 +148,11 @@ pub(crate) enum TaskTermination { Failure(GenericError), } -#[derive(Debug)] -pub(crate) enum ActionEffect { +#[derive(derive_more::Debug)] +// Keep proposals inline to avoid allocating on the network request path. At most one network +// event can be buffered at a time. +#[allow(clippy::large_enum_variant)] +pub(crate) enum LeaderEvent { Scheduler(scheduler::Decisions), Invoker(InvokerEffect), Shuffle(shuffle::OutboxTruncation), @@ -156,6 +162,16 @@ pub(crate) enum ActionEffect { UpsertSchema(Schema), UpsertRuleBook(Arc), AwaitingRpcSelfProposeDone, + RpcProposal { + proposal: rpc::RpcProposal, + #[debug(skip)] + reciprocal: RpcReciprocal, + }, + IngestRecords { + records: Vec, + #[debug(skip)] + callback: AppendCallback, + }, } enum State { @@ -230,13 +246,19 @@ where } } - pub(super) fn should_process_rpc(&self) -> bool { - // In case of BecomingLeader we prefer to park RPC requests - // until we transition out of the current state. - matches!( - self.state, - State::Leader(_) | State::Follower | State::Candidate { .. } - ) + pub(super) fn try_reserve_rpc_processing_permit(&self) -> Option { + match &self.state { + State::Leader(leader_state) => leader_state + .try_reserve_network_event_permit() + .ok() + .map(RpcProcessingPermit::Leader), + State::Follower | State::Candidate { .. } => Some(RpcProcessingPermit::NonLeader { + partition_id: self.partition_id(), + }), + // In case of BecomingLeader we prefer to park RPC requests + // until we transition out of the current state. + State::BecomingLeader { .. } => None, + } } #[instrument(level = "debug", skip_all, fields(leader_epoch = %leadership_info.leader_epoch))] @@ -824,11 +846,11 @@ where /// /// * Follower: Nothing to do /// * Candidate: Monitor appender task - /// * Leader: Await action effects and monitor appender task + /// * Leader: Await events and monitor appender task pub async fn run( &mut self, ctx: impl Processor + HasVQueues, - ) -> Result, Error> { + ) -> Result, Error> { match &mut self.state { State::Follower => Ok(futures::future::pending::>().await), State::Candidate { self_proposer, .. } @@ -842,17 +864,15 @@ where } } - pub async fn handle_action_effects( + pub async fn handle_events( &mut self, - action_effects: impl IntoIterator, + events: impl IntoIterator, ) -> Result<(), Error> { match &mut self.state { State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => { // nothing to do :-) } - State::Leader(leader_state) => { - leader_state.handle_action_effects(action_effects).await? - } + State::Leader(leader_state) => leader_state.handle_events(events).await?, } Ok(()) @@ -891,102 +911,6 @@ impl LeadershipState { } } } - - pub async fn handle_rpc_proposal_command( - &mut self, - request_id: PartitionProcessorRpcRequestId, - reciprocal: Reciprocal< - Oneshot>, - >, - partition_key: PartitionKey, - cmd: Command, - ) { - match &mut self.state { - State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => { - // Just fail the rpc - reciprocal.send(Err(PartitionProcessorRpcError::NotLeader( - self.partition_id, - ))) - } - State::Leader(leader_state) => { - leader_state - .handle_rpc_proposal_command(request_id, reciprocal, partition_key, cmd) - .await; - } - } - } - - pub async fn propose_pause_and_fence( - &mut self, - request_id: PartitionProcessorRpcRequestId, - reciprocal: Reciprocal< - Oneshot>, - >, - invocation_id: InvocationId, - cmd: Command, - ) { - match &mut self.state { - State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => { - // Just fail the rpc - reciprocal.send(Err(PartitionProcessorRpcError::NotLeader( - self.partition_id, - ))) - } - State::Leader(leader_state) => { - leader_state - .propose_pause_and_fence(request_id, reciprocal, invocation_id, cmd) - .await; - } - } - } - - /// Append a command to Bifrost without dedup information, responding on Bifrost commit. - pub async fn append_and_respond_asynchronously( - &mut self, - partition_key: PartitionKey, - cmd: Command, - reciprocal: Reciprocal< - Oneshot>, - >, - success_response: PartitionProcessorRpcResponse, - ) { - match &mut self.state { - State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => reciprocal - .send(Err(PartitionProcessorRpcError::NotLeader( - self.partition_id, - ))), - State::Leader(leader_state) => { - leader_state - .append_and_respond_asynchronously( - partition_key, - cmd, - reciprocal, - success_response, - ) - .await; - } - } - } - - /// Forward externally-created records to this partition. - pub async fn forward_many_with_callback( - &mut self, - records: impl ExactSizeIterator, - callback: F, - ) where - F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static, - { - match &mut self.state { - State::Follower | State::Candidate { .. } | State::BecomingLeader { .. } => callback( - Err(PartitionProcessorRpcError::NotLeader(self.partition_id)), - ), - State::Leader(leader_state) => { - leader_state - .forward_many_with_callback(records, callback) - .await; - } - } - } } #[derive(Debug, derive_more::From)] @@ -1031,6 +955,38 @@ impl shuffle::OutboxReader for OutboxReader { } } +pub(super) enum RpcProcessingPermit { + Leader(tokio::sync::mpsc::OwnedPermit), + NonLeader { partition_id: PartitionId }, +} + +impl RpcProcessingPermit { + pub fn buffer_rpc_proposal(self, proposal: rpc::RpcProposal, reciprocal: RpcReciprocal) { + match self { + RpcProcessingPermit::NonLeader { partition_id } => { + reciprocal.send(Err(PartitionProcessorRpcError::NotLeader(partition_id))) + } + RpcProcessingPermit::Leader(permit) => { + permit.send(LeaderEvent::RpcProposal { + proposal, + reciprocal, + }); + } + } + } + + pub fn buffer_forwarded_records(self, records: Vec, callback: AppendCallback) { + match self { + RpcProcessingPermit::NonLeader { partition_id } => { + callback(Err(PartitionProcessorRpcError::NotLeader(partition_id))) + } + RpcProcessingPermit::Leader(permit) => { + permit.send(LeaderEvent::IngestRecords { records, callback }); + } + } + } +} + #[cfg(test)] mod tests { use std::num::NonZeroUsize; diff --git a/crates/worker/src/partition/mod.rs b/crates/worker/src/partition/mod.rs index 339a15f851..2443a0aeea 100644 --- a/crates/worker/src/partition/mod.rs +++ b/crates/worker/src/partition/mod.rs @@ -29,6 +29,7 @@ pub mod state_machine; mod state_machine; pub mod types; +use self::leadership::RpcProcessingPermit; pub use self::node::NodeContext; // Re-exported so external drivers (e.g. pp-bench) can build a context to drive `StateMachine::apply`. #[cfg(feature = "expose-internals")] @@ -98,7 +99,7 @@ use self::state_machine::StateMachine; use crate::metric_definitions::{ LEADER_LABEL, LEADER_LABEL_LEADER, PARTITION_RECORD_COMMITTED_TO_READ_LATENCY_SECONDS, }; -use crate::partition::leadership::LeadershipState; +use crate::partition::leadership::{AppendCallback, LeadershipState}; use crate::partition::processor::leadership::LeadershipContext; use crate::partition::state_machine::ActionCollector; @@ -288,7 +289,7 @@ pub enum ProcessorError { )] MigrationBarrier { reason: String }, #[error(transparent)] - ActionEffect(#[from] leadership::Error), + Leadership(#[from] leadership::Error), #[error(transparent)] ShutdownError(#[from] ShutdownError), #[error("log read stream has terminated")] @@ -527,6 +528,8 @@ where let config = self.node_ctx.config.live_load(); let max_batching_size = config.worker.max_command_batch_size(); let bytes_limit = config.worker.max_command_batch_bytes.as_usize(); + let network_processing_permit = + self.leadership_state.try_reserve_rpc_processing_permit(); tokio::select! { _ = self.target_leader_state_rx.changed() => { @@ -540,9 +543,9 @@ where self.leadership_state.maybe_step_down(&mut self.ctx, new_state.current_leader_epoch, new_state.current_leader).await; self.refresh_status(&mut durable_lsn_watch)?; } - Some(msg) = self.network_leader_svc_rx.next(), if self.leadership_state.should_process_rpc() => { + Some(msg) = self.network_leader_svc_rx.next(), if network_processing_permit.is_some() => { // todo: replace the live schema with the leader's consistent schema - self.on_rpc(msg, live_schemas.live_load(), &last_applied_lsn_watch).await; + self.on_rpc(msg, live_schemas.live_load(), &last_applied_lsn_watch, network_processing_permit.unwrap()).await; } _ = status_update_timer.tick() => { // Update the status to ensure changes are observable @@ -625,11 +628,11 @@ where self.leadership_state.handle_actions(&mut self.ctx, action_collector.drain(..))?; }, result = self.leadership_state.run(&mut self.ctx) => { - let action_effects = result?; - // We process the action_effects not directly in the run future because it + let events = result?; + // We process the events not directly in the run future because it // requires the run future to be cancellation safe. In the future this could be // implemented. - self.leadership_state.handle_action_effects(action_effects).await?; + self.leadership_state.handle_events(events).await?; } Some(leader_query_cmd) = self.leader_query_rx.recv() => { self.on_leader_query(leader_query_cmd); @@ -675,6 +678,7 @@ where >, body: PartitionProcessorRpcRequest, schemas: &Schema, + permit: RpcProcessingPermit, ) { let context = rpc::RpcContext::new( self.leadership_state.is_leader(), @@ -685,35 +689,7 @@ where 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::Propose(proposal) => permit.buffer_rpc_proposal(proposal, response_tx), rpc::Decision::Reply(reply) => response_tx.send(reply), rpc::Decision::NotifyInvokerAndReply { notification, @@ -771,16 +747,18 @@ where msg: ServiceMessage, schemas: &Schema, last_applied_lsn_watch: &watch::Receiver, + permit: RpcProcessingPermit, ) { match msg { ServiceMessage::Rpc(msg) if msg.msg_type() == PartitionProcessorRpcRequest::TYPE => { let msg = msg.into_typed::(); // note: split() decodes the payload let (response_tx, body) = msg.split(); - self.on_pp_rpc_request(response_tx, body, schemas).await; + self.on_pp_rpc_request(response_tx, body, schemas, permit) + .await; } ServiceMessage::Rpc(msg) if msg.msg_type() == ReceivedIngestRequest::TYPE => { - self.on_pp_ingest_request(msg.into_typed()).await; + self.on_pp_ingest_request(msg.into_typed(), permit); } ServiceMessage::Rpc(msg) if msg.msg_type() == DedupSequenceNrQueryRequest::TYPE => { self.wait_for_tail_then( @@ -879,26 +857,24 @@ where } } - async fn on_pp_ingest_request(&mut self, msg: Incoming>) { + fn on_pp_ingest_request( + &mut self, + msg: Incoming>, + permit: RpcProcessingPermit, + ) { let (reciprocal, request) = msg.split(); + let callback: AppendCallback = Box::new(move |result| match result { + Ok(()) => reciprocal.send(ResponseStatus::Ack.into()), + Err( + PartitionProcessorRpcError::NotLeader(id) + | PartitionProcessorRpcError::LostLeadership(id), + ) => reciprocal.send(ResponseStatus::NotLeader { of: id }.into()), + Err(PartitionProcessorRpcError::Internal(msg)) => { + reciprocal.send(ResponseStatus::Internal { msg }.into()) + } + }); - self.leadership_state - .forward_many_with_callback( - request.records.into_iter(), - move |result: Result<(), PartitionProcessorRpcError>| match result { - Ok(_) => reciprocal.send(ResponseStatus::Ack.into()), - Err(err) => match err { - PartitionProcessorRpcError::NotLeader(id) - | PartitionProcessorRpcError::LostLeadership(id) => { - reciprocal.send(ResponseStatus::NotLeader { of: id }.into()) - } - PartitionProcessorRpcError::Internal(msg) => { - reciprocal.send(ResponseStatus::Internal { msg }.into()) - } - }, - }, - ) - .await; + permit.buffer_forwarded_records(request.records, callback); } // --- Apply new commands/records