diff --git a/crates/wal-protocol/src/invocation.rs b/crates/wal-protocol/src/invocation.rs new file mode 100644 index 0000000000..afdd45cecc --- /dev/null +++ b/crates/wal-protocol/src/invocation.rs @@ -0,0 +1,43 @@ +// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH. +// All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use bytes::{Buf, BufMut, Bytes}; + +use restate_types::bilrost_storage_encode_decode; +use restate_types::identifiers::{InvocationId, PartitionProcessorRpcRequestId}; + +/// Pause an invocation rpc triggered from external sources (http ingress) +#[derive(Debug, Clone, bilrost::Message)] +pub struct PauseInvocationRpcRequest { + #[bilrost(tag(1))] + pub invocation_id: InvocationId, + #[bilrost(tag(2))] + pub request_id: PartitionProcessorRpcRequestId, +} + +bilrost_storage_encode_decode!(PauseInvocationRpcRequest); + +impl PauseInvocationRpcRequest { + pub fn bilrost_encode(&self, b: &mut B) -> Result<(), bilrost::EncodeError> { + bilrost::Message::encode(self, b) + } + + pub fn encoded_len(&self) -> usize { + bilrost::Message::encoded_len(self) + } + + pub fn bilrost_encode_to_bytes(&self) -> Bytes { + bilrost::Message::encode_to_bytes(self) + } + + pub fn bilrost_decode(buf: B) -> Result { + bilrost::OwnedMessage::decode(buf) + } +} diff --git a/crates/wal-protocol/src/lib.rs b/crates/wal-protocol/src/lib.rs index 060c3c6ad6..7b27dd2e6a 100644 --- a/crates/wal-protocol/src/lib.rs +++ b/crates/wal-protocol/src/lib.rs @@ -9,6 +9,7 @@ // by the Apache License, Version 2.0. pub mod control; +pub mod invocation; pub mod timer; pub mod v1; pub mod v2; diff --git a/crates/wal-protocol/src/v1.rs b/crates/wal-protocol/src/v1.rs index d30f2e027e..329ee0e62e 100644 --- a/crates/wal-protocol/src/v1.rs +++ b/crates/wal-protocol/src/v1.rs @@ -165,6 +165,11 @@ pub enum Command { AttachInvocation(AttachInvocationRequest), /// Resume an invocation ResumeInvocation(ResumeInvocationRequest), + /// Pause an invocation + /// payload is bilrost encoded [`invocation::PauseInvocationRpcRequest`] + /// + /// Introduced in v1.7.0 to support pausing invocations with vqueues + PauseInvocationRpcRequest(#[debug(skip)] Bytes), /// Restart as new invocation from prefix RestartAsNewInvocation(RestartAsNewInvocationRequest), @@ -252,6 +257,7 @@ impl HasRecordKeys for Envelope { Command::ProxyThrough(_) => Keys::Single(self.partition_key()), Command::AttachInvocation(_) => Keys::Single(self.partition_key()), Command::ResumeInvocation(req) => Keys::Single(req.invocation_id.partition_key()), + Command::PauseInvocationRpcRequest(_) => Keys::Single(self.partition_key()), Command::RestartAsNewInvocation(req) => Keys::Single(req.invocation_id.partition_key()), // todo: Handle journal entries that request cross-partition invocations Command::InvokerEffect(effect) => Keys::Single(effect.invocation_id.partition_key()), diff --git a/crates/wal-protocol/src/v2.rs b/crates/wal-protocol/src/v2.rs index 039df5a211..8a9dea5633 100644 --- a/crates/wal-protocol/src/v2.rs +++ b/crates/wal-protocol/src/v2.rs @@ -309,6 +309,10 @@ pub enum CommandKind { /// payload is bilrost encoded [`vqueues::VQueuesResume`] /// *Since v1.7.0 VQueuesResume = 24, + + /// Pause an invocation + /// Introduced in v1.7.0 to support pausing invocations with vqueues + PauseInvocationRpcRequest = 25, } mod bilrost_encoding { diff --git a/crates/wal-protocol/src/v2/commands.rs b/crates/wal-protocol/src/v2/commands.rs index cbecfd0e18..b72b5cfeb1 100644 --- a/crates/wal-protocol/src/v2/commands.rs +++ b/crates/wal-protocol/src/v2/commands.rs @@ -23,8 +23,9 @@ use restate_types::{ use super::sealed::Sealed; use super::{Command, CommandKind}; pub use crate::control::UpsertRuleBookCommand; +pub use crate::invocation::PauseInvocationRpcRequest; use crate::timer; -// Re-epxort vqueues commands +// Re-export vqueues commands pub use crate::vqueues::{VQueuesPauseCommand, VQueuesResumeCommand}; pub use crate::control::{ @@ -467,3 +468,8 @@ command! { @kind=CommandKind::VQueuesResume, @command=VQueuesResumeCommand } + +command! { + @kind=CommandKind::PauseInvocationRpcRequest, + @command=PauseInvocationRpcRequest +} diff --git a/crates/wal-protocol/src/v2/compatibility.rs b/crates/wal-protocol/src/v2/compatibility.rs index 9f4e0d6715..71595fb967 100644 --- a/crates/wal-protocol/src/v2/compatibility.rs +++ b/crates/wal-protocol/src/v2/compatibility.rs @@ -155,6 +155,12 @@ impl TryFrom for v2::Envelope { dedup, payload, ), + v1::Command::PauseInvocationRpcRequest(payload) => Envelope::from_bytes_unchecked( + v2::CommandKind::PauseInvocationRpcRequest, + StorageCodecKind::Bilrost, + dedup, + payload, + ), }; Ok(envelope) diff --git a/crates/worker/src/partition/leadership/leader_state.rs b/crates/worker/src/partition/leadership/leader_state.rs index 3d22faf70f..1e5f6db946 100644 --- a/crates/worker/src/partition/leadership/leader_state.rs +++ b/crates/worker/src/partition/leadership/leader_state.rs @@ -697,6 +697,16 @@ impl LeaderState { ))); } } + Action::ForwardPauseInvocationResponse { + request_id, + response, + } => { + if let Some(response_tx) = self.awaiting_rpc_actions.remove(&request_id) { + response_tx.send(Ok(PartitionProcessorRpcResponse::PauseInvocation( + response.into(), + ))); + } + } Action::ForwardPurgeJournalResponse { request_id, response, diff --git a/crates/worker/src/partition/state_machine/actions.rs b/crates/worker/src/partition/state_machine/actions.rs index 1e625e51d6..8da7732ec0 100644 --- a/crates/worker/src/partition/state_machine/actions.rs +++ b/crates/worker/src/partition/state_machine/actions.rs @@ -16,7 +16,8 @@ use restate_types::identifiers::{EntryIndex, InvocationId, PartitionProcessorRpc use restate_types::invocation::InvocationTarget; use restate_types::invocation::client::{ CancelInvocationResponse, InvocationOutputResponse, KillInvocationResponse, - PurgeInvocationResponse, RestartAsNewInvocationResponse, ResumeInvocationResponse, + PauseInvocationResponse, PurgeInvocationResponse, RestartAsNewInvocationResponse, + ResumeInvocationResponse, }; use restate_types::journal_v2::{CommandIndex, NotificationId}; use restate_types::message::MessageIndex; @@ -105,6 +106,10 @@ pub enum Action { request_id: PartitionProcessorRpcRequestId, response: RestartAsNewInvocationResponse, }, + ForwardPauseInvocationResponse { + request_id: PartitionProcessorRpcRequestId, + response: PauseInvocationResponse, + }, /// Forward a batch of rule-book diff entries to the leader's /// `UserLimiter` via the resource-manager mpsc. Emitted by /// `Command::UpsertRuleBook` apply when the rule book version diff --git a/crates/worker/src/partition/state_machine/lifecycle/mod.rs b/crates/worker/src/partition/state_machine/lifecycle/mod.rs index 7cfa923198..b2dbd7b1d5 100644 --- a/crates/worker/src/partition/state_machine/lifecycle/mod.rs +++ b/crates/worker/src/partition/state_machine/lifecycle/mod.rs @@ -16,6 +16,7 @@ mod notify_get_invocation_output_response; mod notify_invocation_response; mod notify_signal; mod notify_sleep_completion; +mod pause_invocation_rpc; mod paused; mod pinned_deployment; mod purge; @@ -34,6 +35,7 @@ pub(super) use notify_get_invocation_output_response::OnNotifyGetInvocationOutpu pub(super) use notify_invocation_response::OnNotifyInvocationResponse; pub(super) use notify_signal::OnNotifySignalCommand; pub(super) use notify_sleep_completion::OnNotifySleepCompletionCommand; +pub(super) use pause_invocation_rpc::OnInvocationPauseRpcRequest; pub(super) use paused::OnPausedCommand; pub(super) use pinned_deployment::OnPinnedDeploymentCommand; pub(super) use purge::OnPurgeCommand; diff --git a/crates/worker/src/partition/state_machine/lifecycle/pause_invocation_rpc.rs b/crates/worker/src/partition/state_machine/lifecycle/pause_invocation_rpc.rs new file mode 100644 index 0000000000..f1793c4489 --- /dev/null +++ b/crates/worker/src/partition/state_machine/lifecycle/pause_invocation_rpc.rs @@ -0,0 +1,269 @@ +// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH. +// All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use tracing::debug; + +use restate_clock::UniqueTimestamp; +use restate_types::invocation::client::PauseInvocationResponse; +use restate_storage_api::invocation_status_table::{ + InvocationStatus, ReadInvocationStatusTable, WriteInvocationStatusTable, +}; +use restate_storage_api::journal_events::WriteJournalEventsTable; +use restate_storage_api::lock_table::WriteLockTable; +use restate_storage_api::vqueue_table::{self, EntryStatusHeader, ReadVQueueTable, WriteVQueueTable}; +use restate_types::identifiers::{ + InvocationId, PartitionProcessorRpcRequestId, WithPartitionKey as _, +}; +use restate_types::vqueues::EntryId; +use restate_vqueues::VQueue; + +use crate::debug_if_leader; +use crate::partition::state_machine::lifecycle::event::ApplyEventCommand; +use crate::partition::state_machine::{CommandHandler, Error, StateMachineApplyContext}; + +pub struct OnInvocationPauseRpcRequest<'a> { + pub invocation_id: &'a InvocationId, + pub request_id: PartitionProcessorRpcRequestId, +} + +impl<'ctx, 's: 'ctx, S> CommandHandler<&'ctx mut StateMachineApplyContext<'s, S>> + for OnInvocationPauseRpcRequest<'_> +where + S: ReadInvocationStatusTable + + WriteInvocationStatusTable + + WriteJournalEventsTable + + WriteVQueueTable + + WriteLockTable + + ReadVQueueTable, +{ + async fn apply(self, ctx: &'ctx mut StateMachineApplyContext<'s, S>) -> Result<(), Error> { + let Self { + invocation_id, + request_id, + } = self; + + let at = UniqueTimestamp::from_unix_millis_unchecked(ctx.record_created_at); + + let entry_id = EntryId::from(invocation_id); + let Some(header) = ctx + .storage + .get_vqueue_entry_status(invocation_id.partition_key(), &entry_id) + .await? + else { + // This is equivalent to InvocationStatus::Free. + debug!( + "Trying to pause invocation {invocation_id} which does not exist as a vqueue entry, will ignore." + ); + ctx.reply_to_pause_invocation(self.request_id,PauseInvocationResponse::NotFound); + return Ok(()); + }; + + match header.stage() { + vqueue_table::Stage::Paused => { + ctx.reply_to_pause_invocation(self.request_id,PauseInvocationResponse::AlreadyPaused); + return Ok(()) + } + vqueue_table::Stage::Finished => { + debug!( + "Trying to pause invocation {invocation_id} which already finished." + ); + ctx.reply_to_pause_invocation(self.request_id,PauseInvocationResponse::NotFound); + return Ok(()); + } + _ => {} + }; + + VQueue::get( + header.vqueue_id(), + ctx.storage, + ctx.vqueues_cache, + ctx.is_leader.then_some(ctx.action_collector), + ) + .await? + .expect("pausing in a non-existent vqueue") + .pause_entry(at, &header); + + + + let invoked_meta = match ctx.get_invocation_status(invocation_id).await? { + InvocationStatus::Suspended { metadata: meta,.. }, + InvocationStatus::Invoked(meta) => meta, + InvocationStatus::Suspended { .. } + | InvocationStatus::Paused(_) + | InvocationStatus::Scheduled(_) + | InvocationStatus::Inboxed(_) + | InvocationStatus::Completed(_) + | InvocationStatus::Free => { + // Nothing to do in these cases, pause gets processed only if the invocation was Invoked. + return Ok(()); + } + }; + + // Invoker paused the invocation, let's record the event, then set the status to paused + debug_if_leader!(ctx.is_leader, "Paused the invocation"); + + if invoked_meta.vqueue_id.is_some() { + // todo: use the new status + let entry_id = EntryId::from(invocation_id); + let Some(header) = ctx + .storage + .get_vqueue_entry_status(invocation_id.partition_key(), &entry_id) + .await? + else { + // This is equivalent to InvocationStatus::Free. + debug!( + "Trying to pause invocation {invocation_id} which does not exist as a vqueue entry, will ignore." + ); + return Ok(()); + }; + + let at = UniqueTimestamp::from_unix_millis_unchecked(ctx.record_created_at); + VQueue::get( + header.vqueue_id(), + ctx.storage, + ctx.vqueues_cache, + ctx.is_leader.then_some(ctx.action_collector), + ) + .await? + .expect("pausing in a non-existent vqueue") + .pause_entry(at, &header); + } + + let mut invocation_status = InvocationStatus::Paused(invoked_meta); + + ApplyEventCommand { + invocation_id, + invocation_status: &invocation_status, + event: paused_event, + } + .apply(ctx) + .await?; + + // Update timestamps + if let Some(timestamps) = invocation_status.get_timestamps_mut() { + timestamps.update(ctx.record_created_at); + } + + ctx.storage + .put_invocation_status(self.invocation_id, &invocation_status)?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use googletest::prelude::*; + + use restate_storage_api::invocation_status_table::{ + InFlightInvocationMetadata, InvocationStatusDiscriminants, ReadInvocationStatusTable, + }; + use restate_types::journal_events::{Event, PausedEvent, TransientErrorEvent}; + use restate_wal_protocol::v2::{Command, commands}; + + use crate::partition::state_machine::tests::{TestEnv, fixtures, matchers}; + use crate::partition::types::InvokerEffectKind; + + #[restate_core::test] + async fn paused_with_pinned_deployment() { + let mut test_env = TestEnv::create().await; + let invocation_id = fixtures::mock_start_invocation(&mut test_env).await; + fixtures::mock_pinned_deployment_v5(&mut test_env, invocation_id).await; + + let paused_event = Event::from(PausedEvent { + last_failure: Some(TransientErrorEvent { + error_code: 501u16.into(), + error_message: "my bad".to_string(), + error_stacktrace: Some("something something".to_string()), + restate_doc_error_code: Some("RT0001".to_string()), + related_command_index: None, + related_command_name: Some("my command".to_string()), + related_command_type: None, + }), + }); + + // Check we just pause + let _ = test_env + .apply(commands::InvokerEffectCommand::test_envelope( + restate_worker_api::invoker::Effect { + invocation_id, + kind: InvokerEffectKind::Paused { + paused_event: paused_event.clone().into(), + }, + }, + )) + .await; + assert_that!( + test_env + .storage + .get_invocation_status(&invocation_id) + .await + .unwrap(), + all!( + matchers::storage::is_variant(InvocationStatusDiscriminants::Paused), + matchers::storage::in_flight_metadata(field!( + InFlightInvocationMetadata.pinned_deployment, + some(anything()) + )) + ) + ); + assert_that!( + test_env.read_journal_events(&invocation_id).await, + elements_are![eq(paused_event)] + ); + + test_env.shutdown().await; + } + + #[restate_core::test] + async fn paused_when_deployment_version_not_set_yet() { + let mut test_env = TestEnv::create().await; + let invocation_id = fixtures::mock_start_invocation(&mut test_env).await; + + let paused_event = Event::from(PausedEvent { + last_failure: Some(TransientErrorEvent { + error_code: 501u16.into(), + error_message: "my bad".to_string(), + error_stacktrace: Some("something something".to_string()), + restate_doc_error_code: Some("RT0001".to_string()), + related_command_index: None, + related_command_name: Some("my command".to_string()), + related_command_type: None, + }), + }) + .into(); + + // Check we just pause + let _ = test_env + .apply(commands::InvokerEffectCommand::test_envelope( + restate_worker_api::invoker::Effect { + invocation_id, + kind: InvokerEffectKind::Paused { paused_event }, + }, + )) + .await; + assert_that!( + test_env + .storage + .get_invocation_status(&invocation_id) + .await + .unwrap(), + all!( + matchers::storage::is_variant(InvocationStatusDiscriminants::Paused), + matchers::storage::in_flight_metadata(field!( + InFlightInvocationMetadata.pinned_deployment, + none() + )) + ) + ); + + test_env.shutdown().await; + } +} diff --git a/crates/worker/src/partition/state_machine/mod.rs b/crates/worker/src/partition/state_machine/mod.rs index 8055374207..ccf2306dce 100644 --- a/crates/worker/src/partition/state_machine/mod.rs +++ b/crates/worker/src/partition/state_machine/mod.rs @@ -78,7 +78,7 @@ use restate_types::identifiers::{ }; use restate_types::invocation::client::{ CancelInvocationResponse, InvocationOutputResponse, KillInvocationResponse, - PurgeInvocationResponse, ResumeInvocationResponse, + PauseInvocationResponse, PurgeInvocationResponse, ResumeInvocationResponse, }; use restate_types::invocation::{ AttachInvocationRequest, IngressInvocationResponseSink, InvocationInput, @@ -782,6 +782,19 @@ impl StateMachineApplyContext<'_, S> { .await?; Ok(()) } + CommandKind::PauseInvocationRpcRequest => { + let request = envelope + .into_typed::() + .into_inner()?; + + lifecycle::OnInvocationPauseRpcRequest { + invocation_id: &request.invocation_id, + request_id: &request.request_id, + } + .apply(self) + .await?; + Ok(()) + } CommandKind::RestartAsNewInvocation => { let restart_as_new_invocation_request: RestartAsNewInvocationRequest = envelope .into_typed::() @@ -2281,6 +2294,7 @@ impl StateMachineApplyContext<'_, S> { self.end_invocation( invocation_id, metadata, + Some(TerminationFlavor::Kill), Some(ResponseResult::Failure(KILLED_INVOCATION_ERROR)), ) .await?; @@ -2318,6 +2332,7 @@ impl StateMachineApplyContext<'_, S> { self.end_invocation( invocation_id, metadata, + Some(TerminationFlavor::Kill), Some(ResponseResult::Failure(KILLED_INVOCATION_ERROR)), ) .await?; @@ -2855,6 +2870,7 @@ impl StateMachineApplyContext<'_, S> { .into_invocation_metadata() .expect("Must be present if status is invoked"), None, + None, ) .await?; } @@ -2864,6 +2880,7 @@ impl StateMachineApplyContext<'_, S> { invocation_status .into_invocation_metadata() .expect("Must be present if status is invoked"), + None, Some(ResponseResult::Failure(e)), ) .await?; @@ -2959,10 +2976,14 @@ impl StateMachineApplyContext<'_, S> { } /// TODO(slinkydeveloper) move this to lifecycle command + /// + /// `flavor` lets us know how the invocation ended. If `flavor` is `None`, + /// then we determine the termination status based on the `response_result_override`. async fn end_invocation( &mut self, invocation_id: InvocationId, invocation_metadata: InFlightInvocationMetadata, + flavor: Option, // If given, this will override any Output Entry available in the journal table response_result_override: Option, ) -> Result<(), Error> @@ -3020,8 +3041,18 @@ impl StateMachineApplyContext<'_, S> { return Ok(()); }; - if let ResponseResult::Failure(_) = &response_result { - end_status = vqueue_table::Status::Failed; + if let ResponseResult::Failure(e) = &response_result { + if e.code() == restate_types::errors::codes::ABORTED { + // special handling for cancel/kill. Definitely not ideal, but the current + // design leaves me with no other options. In practice, to distinguish between + // cancel and kill, the flavor will be used (in vqueues) to make the distinction. + // + // Kill is always passed in `flavor` but cancel must be deduced from the aborted + // code. + end_status = vqueue_table::Status::Cancelled; + } else { + end_status = vqueue_table::Status::Failed; + } } // Send responses out @@ -3097,6 +3128,13 @@ impl StateMachineApplyContext<'_, S> { let record_unique_ts = UniqueTimestamp::from_unix_millis_unchecked(self.record_created_at); + // Make sure we report cancel/killed correctly. + end_status = match flavor { + Some(TerminationFlavor::Cancel) => vqueue_table::Status::Cancelled, + Some(TerminationFlavor::Kill) => vqueue_table::Status::Killed, + None => end_status, + }; + VQueue::get( &vqueue_id, self.storage, @@ -4561,6 +4599,25 @@ impl StateMachineApplyContext<'_, S> { }); } + fn reply_to_pause_invocation( + &mut self, + request_id: PartitionProcessorRpcRequestId, + response: PauseInvocationResponse, + ) { + debug_if_leader!( + self.is_leader, + "Send pause invocation response to request id '{:?}': {:?}", + request_id, + response + ); + + self.action_collector + .push(Action::ForwardPauseInvocationResponse { + request_id, + response, + }); + } + fn reply_to_cancel( &mut self, response_sink: Option,