Skip to content

[worker][refactor] change rpc handlers to return their next steps instead of taking actions inline #5043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MohamedBassem wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[worker][refactor] change rpc handlers to return their next steps instead of taking actions inline #5043

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion crates/worker/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
67 changes: 57 additions & 10 deletions crates/worker/src/partition/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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(
Expand Down
204 changes: 80 additions & 124 deletions crates/worker/src/partition/rpc/append_invocation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,21 +21,15 @@ pub(super) struct Request {
pub(super) append_invocation_reply_on: AppendInvocationReplyOn,
}

impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler<Request>
for RpcContext<'a, TActuator, TSchemas, TStorage>
{
type Output = PartitionProcessorRpcResponse;
type Error = ();

impl<'a, TSchemas, TStorage> RpcHandler<Request> for RpcContext<'a, TSchemas, TStorage> {
async fn handle(
self,
Request {
request_id,
invocation_request,
append_invocation_reply_on,
}: Request,
replier: Replier<Self::Output>,
) -> Result<(), Self::Error> {
) -> Decision {
let mut service_invocation = ServiceInvocation::from_request(
Arc::unwrap_or_clone(invocation_request),
invocation::Source::ingress(request_id),
Expand All @@ -58,148 +52,110 @@ impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler<Request>
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 }
}
},
})
}
}

#[cfg(test)]
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::<PartitionProcessorRpcResponse>()
.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::<PartitionProcessorRpcResponse>()
.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::<PartitionProcessorRpcResponse>()
.never();
proposer
.expect_handle_rpc_proposal_command::<PartitionProcessorRpcResponse>()
.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::<PartitionProcessorRpcResponse>()
.never();
proposer
.expect_handle_rpc_proposal_command::<PartitionProcessorRpcResponse>()
.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()),
))
);
}
}
27 changes: 9 additions & 18 deletions crates/worker/src/partition/rpc/append_invocation_response.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,28 +18,19 @@ pub(super) struct Request {
pub(super) invocation_response: InvocationResponse,
}

impl<'a, TActuator: Actuator, TSchemas, TStorage> RpcHandler<Request>
for RpcContext<'a, TActuator, TSchemas, TStorage>
{
type Output = PartitionProcessorRpcResponse;
type Error = ();

impl<'a, TSchemas, TStorage> RpcHandler<Request> for RpcContext<'a, TSchemas, TStorage> {
async fn handle(
self,
Request {
invocation_response,
}: Request,
replier: Replier<Self::Output>,
) -> 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,
},
})
}
}
Loading
Loading