Implement Reset invocation feature

Author: slinkydeveloper

crates/admin/src/rest_api/error.rs

@@ -195,6 +195,20 @@ impl_meta_api_error!(RestartInvocationMissingInputError: GONE "The invocation ca
 pub(crate) struct RestartInvocationNotStartedError(pub(crate) String);
 impl_meta_api_error!(RestartInvocationNotStartedError: TOO_EARLY "The invocation cannot be restarted because it's not running yet, meaning it might have been scheduled or inboxed.");
 
+#[derive(Debug, thiserror::Error)]
+#[error(
+    "Resetting the invocation '{0}' is not supported, because it was started using the old service protocol."
+)]
+pub(crate) struct ResetInvocationUnsupportedError(pub(crate) String);
+impl_meta_api_error!(ResetInvocationUnsupportedError: UNPROCESSABLE_ENTITY "Resetting the invocation is not supported, because it was started using the old service protocol.");
+
+#[derive(Debug, thiserror::Error)]
+#[error(
+    "The invocation '{0}' cannot be reset because it's not running. For completed invocations, use restart instead."
+)]
+pub(crate) struct ResetInvocationNotRunningError(pub(crate) String);
+impl_meta_api_error!(ResetInvocationNotRunningError: TOO_EARLY "The invocation cannot be reset because it's not running. For completed invocations, use restart instead.");
+
 // --- Old Meta API errors. Please don't use these anymore.
 
 /// This error is used by handlers to propagate API errors,

crates/admin/src/rest_api/invocations.rs

@@ -21,8 +21,10 @@ use restate_types::invocation::client::{
     self, CancelInvocationResponse, InvocationClient, KillInvocationResponse,
     PurgeInvocationResponse,
 };
+use restate_types::invocation::reset::TruncateFrom;
 use restate_types::invocation::{
-    InvocationEpoch, InvocationTermination, PurgeInvocationRequest, TerminationFlavor, restart,
+    InvocationEpoch, InvocationTermination, PurgeInvocationRequest, TerminationFlavor, reset,
+    restart,
 };
 use restate_wal_protocol::{Command, Envelope};
 use serde::{Deserialize, Serialize};
@@ -514,3 +516,171 @@ where
         }
     }
 }
+
+#[derive(Default, Debug, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ResetInvocationApplyToChildInvocations {
+    Nothing,
+    /// Kill all the child invocations that have been created after the truncation point
+    #[default]
+    Kill,
+    /// Cancel all the child invocations that have been created after the truncation point
+    Cancel,
+}
+
+impl From<ResetInvocationApplyToChildInvocations> for reset::ApplyToChildInvocations {
+    fn from(value: ResetInvocationApplyToChildInvocations) -> Self {
+        match value {
+            ResetInvocationApplyToChildInvocations::Kill => reset::ApplyToChildInvocations::Kill,
+            ResetInvocationApplyToChildInvocations::Nothing => {
+                reset::ApplyToChildInvocations::Nothing
+            }
+            ResetInvocationApplyToChildInvocations::Cancel => {
+                reset::ApplyToChildInvocations::Cancel
+            }
+        }
+    }
+}
+
+#[derive(Default, Debug, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ResetInvocationApplyToPinnedDeployment {
+    #[default]
+    Keep,
+    /// Clear the pinned deployment.
+    ///
+    /// NOTE: If the new picked up deployment doesn't support the current service protocol version, the invocation will remain stuck in a retry loop. Use with caution!
+    Clear,
+}
+
+impl From<ResetInvocationApplyToPinnedDeployment> for reset::ApplyToPinnedDeployment {
+    fn from(value: ResetInvocationApplyToPinnedDeployment) -> Self {
+        match value {
+            ResetInvocationApplyToPinnedDeployment::Keep => reset::ApplyToPinnedDeployment::Keep,
+            ResetInvocationApplyToPinnedDeployment::Clear => reset::ApplyToPinnedDeployment::Clear,
+        }
+    }
+}
+
+#[derive(Debug, Default, Deserialize, JsonSchema)]
+pub struct ResetInvocationParams {
+    pub truncate_from: Option<u32>,
+    #[serde(
+        default,
+        with = "serde_with::As::<Option<restate_serde_util::DurationString>>"
+    )]
+    #[schemars(with = "Option<String>")]
+    pub previous_attempt_retention: Option<Duration>,
+    pub apply_to_child_calls: Option<ResetInvocationApplyToChildInvocations>,
+    pub apply_to_pinned_deployment: Option<ResetInvocationApplyToPinnedDeployment>,
+}
+
+#[derive(Debug, Serialize, JsonSchema)]
+pub struct ResetInvocationResponse {
+    /// The new invocation epoch of the invocation.
+    pub new_invocation_epoch: InvocationEpoch,
+}
+
+generate_meta_api_error!(ResetInvocationError: [
+    InvocationNotFoundError,
+    InvocationClientError,
+    InvalidFieldError,
+    ResetInvocationUnsupportedError,
+    ResetInvocationNotRunningError
+]);
+
+/// Reset an invocation
+#[openapi(
+    summary = "Reset an invocation",
+    description = "Reset the given invocation, truncating the progress from the given journal entry index onward and resuming afterward.",
+    operation_id = "reset_invocation",
+    tags = "invocation",
+    parameters(
+        path(
+            name = "invocation_id",
+            description = "Invocation identifier.",
+            schema = "std::string::String"
+        ),
+        query(
+            name = "truncate_from",
+            description = "Journal entry index to truncate from, inclusive. The index MUST correspond to a command entry or to a signal notification, and it cannot be zero, otherwise this operation will fail. If not provided, it defaults to 1 (after the first entry).",
+            required = false,
+            style = "simple",
+            allow_empty_value = false,
+            schema = "u32",
+        ),
+        query(
+            name = "previous_attempt_retention",
+            description = "If set, it will override the configured completion_retention/journal_retention when the invocation was executed the first time. If none of the completion_retention/journal_retention are configured, and neither this previous_attempt_retention, then the previous attempt won't be retained at all. Can be configured using humantime format or ISO8601.",
+            required = false,
+            style = "simple",
+            allow_empty_value = false,
+            schema = String,
+        ),
+        query(
+            name = "apply_to_child_calls",
+            description = "What to do with children calls that have been created after the truncation point. By default, kills all the children calls. This doesn't apply to sends.",
+            required = false,
+            style = "simple",
+            allow_empty_value = false,
+            schema = ResetInvocationApplyToChildInvocations,
+        ),
+        query(
+            name = "apply_to_pinned_deployment",
+            description = "What to do with pinned deployment. By default, the current pinned deployment will be kept.",
+            required = false,
+            style = "simple",
+            allow_empty_value = false,
+            schema = ResetInvocationApplyToPinnedDeployment,
+        )
+    )
+)]
+pub async fn reset_invocation<V, IC>(
+    State(state): State<AdminServiceState<V, IC>>,
+    Path(invocation_id): Path<String>,
+    Query(ResetInvocationParams {
+        truncate_from,
+        previous_attempt_retention,
+        apply_to_child_calls,
+        apply_to_pinned_deployment,
+    }): Query<ResetInvocationParams>,
+) -> Result<Json<ResetInvocationResponse>, ResetInvocationError>
+where
+    IC: InvocationClient,
+{
+    let invocation_id = invocation_id
+        .parse::<InvocationId>()
+        .map_err(|e| InvalidFieldError("invocation_id", e.to_string()))?;
+
+    match state
+        .invocation_client
+        .reset_invocation(
+            PartitionProcessorRpcRequestId::new(),
+            invocation_id,
+            TruncateFrom::EntryIndex { entry_index: truncate_from.unwrap_or(1) },
+            previous_attempt_retention,
+            apply_to_child_calls.unwrap_or_default().into(),
+            apply_to_pinned_deployment.unwrap_or_default().into(),
+        )
+        .await
+        .map_err(InvocationClientError)?
+    {
+        client::ResetInvocationResponse::Ok { new_epoch } => Ok(ResetInvocationResponse {
+            new_invocation_epoch: new_epoch,
+        }
+            .into()),
+
+        client::ResetInvocationResponse::NotFound => {
+            Err(InvocationNotFoundError(invocation_id.to_string()))?
+        }
+        client::ResetInvocationResponse::Unsupported => {
+            Err(ResetInvocationUnsupportedError(invocation_id.to_string()))?
+        }
+        client::ResetInvocationResponse::NotRunning => {
+            Err(   ResetInvocationNotRunningError(invocation_id.to_string()))?
+        }
+        client::ResetInvocationResponse::BadIndex => {
+            Err(InvalidFieldError("truncate_from", "The index MUST correspond to a command entry or to a signal notification, and it cannot be zero.".to_owned()))?
+        }
+    }
+}

crates/admin/src/rest_api/mod.rs

@@ -106,6 +106,10 @@ where
             "/invocations/:invocation_id/restart",
             patch(openapi_handler!(invocations::restart_invocation)),
         )
+        .route(
+            "/invocations/:invocation_id/reset",
+            patch(openapi_handler!(invocations::reset_invocation)),
+        )
         .route(
             "/subscriptions",
             post(openapi_handler!(subscriptions::create_subscription)),

crates/core/src/worker_api/partition_processor_rpc_client.rs

@@ -21,7 +21,11 @@ use restate_types::identifiers::{
 use restate_types::invocation::client::{
     AttachInvocationResponse, CancelInvocationResponse, GetInvocationOutputResponse,
     InvocationClient, InvocationClientError, InvocationOutput, KillInvocationResponse,
-    PurgeInvocationResponse, RestartInvocationResponse, SubmittedInvocationNotification,
+    PurgeInvocationResponse, ResetInvocationResponse, RestartInvocationResponse,
+    SubmittedInvocationNotification,
+};
+use restate_types::invocation::reset::{
+    ApplyToChildInvocations, ApplyToPinnedDeployment, TruncateFrom,
 };
 use restate_types::invocation::restart::{ApplyToWorkflowRun, IfRunning};
 use restate_types::invocation::{
@@ -531,4 +535,36 @@ where
             }
         })
     }
+
+    async fn reset_invocation(
+        &self,
+        request_id: PartitionProcessorRpcRequestId,
+        invocation_id: InvocationId,
+        truncate_from: TruncateFrom,
+        previous_attempt_retention: Option<Duration>,
+        apply_to_child_invocations: ApplyToChildInvocations,
+        apply_to_pinned_deployment: ApplyToPinnedDeployment,
+    ) -> Result<ResetInvocationResponse, InvocationClientError> {
+        let response = self
+            .resolve_partition_id_and_send(
+                request_id,
+                PartitionProcessorRpcRequestInner::ResetInvocation {
+                    invocation_id,
+                    previous_attempt_retention,
+                    apply_to_child_calls: apply_to_child_invocations,
+                    apply_to_pinned_deployment,
+                    truncate_from,
+                },
+            )
+            .await?;
+
+        Ok(match response {
+            PartitionProcessorRpcResponse::ResetInvocation(reset_invocation_response) => {
+                reset_invocation_response.into()
+            }
+            _ => {
+                panic!("Expecting ResetInvocation rpc response")
+            }
+        })
+    }
 }

crates/invoker-impl/src/input_command.rs

@@ -33,7 +33,7 @@ pub(crate) struct InvokeCommand {
 pub(crate) enum InputCommand<SR> {
     Invoke(InvokeCommand),
     // TODO remove this when we remove journal v1
-    // Journal V1 doesn't support epochs nor trim and restart
+    // Journal V1 doesn't support epochs nor reset invocation
     Completion {
         partition: PartitionLeaderEpoch,
         invocation_id: InvocationId,

crates/partition-store/src/protobuf_types.rs

@@ -430,7 +430,7 @@ pub mod v1 {
                     deployment_id,
                     service_protocol_version,
                     current_invocation_epoch,
-                    trim_points,
+                    trim_points: truncation_points,
                     waiting_for_completions,
                     waiting_for_signal_indexes,
                     waiting_for_signal_names,
@@ -550,8 +550,8 @@ pub mod v1 {
                                 idempotency_key: idempotency_key.map(ByteString::from),
                                 hotfix_apply_cancellation_after_deployment_is_pinned,
                                 current_invocation_epoch,
-                                completion_range_epoch_map: CompletionRangeEpochMap::from_trim_points(
-                                    trim_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
+                                completion_range_epoch_map: CompletionRangeEpochMap::from_truncation_points(
+                                    truncation_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
                                 ),
                             },
                         ))
@@ -584,8 +584,8 @@ pub mod v1 {
                                 idempotency_key: idempotency_key.map(ByteString::from),
                                 hotfix_apply_cancellation_after_deployment_is_pinned,
                                 current_invocation_epoch,
-                                completion_range_epoch_map: CompletionRangeEpochMap::from_trim_points(
-                                    trim_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
+                                completion_range_epoch_map: CompletionRangeEpochMap::from_truncation_points(
+                                    truncation_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
                                 ),
                             },
                             waiting_for_notifications: waiting_for_completions
@@ -633,8 +633,8 @@ pub mod v1 {
                                     deployment_id,
                                     service_protocol_version,
                                 )?,
-                                completion_range_epoch_map: CompletionRangeEpochMap::from_trim_points(
-                                    trim_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
+                                completion_range_epoch_map: CompletionRangeEpochMap::from_truncation_points(
+                                    truncation_points.into_iter().map(|trim_point|(trim_point.completion_id, trim_point.invocation_epoch))
                                 ),
                             },
                         ))
@@ -842,7 +842,7 @@ pub mod v1 {
                             result: None,
                             hotfix_apply_cancellation_after_deployment_is_pinned,
                             current_invocation_epoch,
-                            trim_points: completion_range_epoch_map.into_trim_points_iter().into_iter().map(|(completion_id, invocation_epoch)| JournalTrimPoint {
+                            trim_points: completion_range_epoch_map.into_truncation_points_iter().into_iter().map(|(completion_id, invocation_epoch)| JournalTrimPoint {
                                 completion_id,
                                 invocation_epoch,
                             }).collect(),
@@ -938,7 +938,7 @@ pub mod v1 {
                             result: None,
                             hotfix_apply_cancellation_after_deployment_is_pinned,
                             current_invocation_epoch,
-                            trim_points: completion_range_epoch_map.into_trim_points_iter().into_iter().map(|(completion_id, invocation_epoch)| JournalTrimPoint {
+                            trim_points: completion_range_epoch_map.into_truncation_points_iter().into_iter().map(|(completion_id, invocation_epoch)| JournalTrimPoint {
                                 completion_id,
                                 invocation_epoch,
                             }).collect(),

crates/partition-store/src/tests/invocation_status_table_test/mod.rs

@@ -99,7 +99,7 @@ fn invoked_status(invocation_target: InvocationTarget) -> InvocationStatus {
         idempotency_key: None,
         hotfix_apply_cancellation_after_deployment_is_pinned: false,
         current_invocation_epoch: 1,
-        completion_range_epoch_map: CompletionRangeEpochMap::from_trim_points([(5, 1)]),
+        completion_range_epoch_map: CompletionRangeEpochMap::from_truncation_points([(5, 1)]),
     })
 }
 
@@ -126,7 +126,7 @@ fn suspended_status(invocation_target: InvocationTarget) -> InvocationStatus {
             idempotency_key: None,
             hotfix_apply_cancellation_after_deployment_is_pinned: false,
             current_invocation_epoch: 1,
-            completion_range_epoch_map: CompletionRangeEpochMap::from_trim_points([(5, 1)]),
+            completion_range_epoch_map: CompletionRangeEpochMap::from_truncation_points([(5, 1)]),
         },
         waiting_for_notifications: HashSet::default(),
     }