Mindless plumbing

Author: slinkydeveloper

Cargo.lock

@@ -21,9 +21,12 @@ use restate_types::identifiers::{
 use restate_types::invocation::client::{
     AttachInvocationResponse, CancelInvocationResponse, GetInvocationOutputResponse,
     InvocationClient, InvocationClientError, InvocationOutput, KillInvocationResponse,
-    PurgeInvocationResponse, SubmittedInvocationNotification,
+    PurgeInvocationResponse, RestartInvocationResponse, SubmittedInvocationNotification,
+};
+use restate_types::invocation::restart::{ApplyToWorkflowRun, IfRunning};
+use restate_types::invocation::{
+    InvocationEpoch, InvocationQuery, InvocationRequest, InvocationResponse,
 };
-use restate_types::invocation::{InvocationQuery, InvocationRequest, InvocationResponse};
 use restate_types::journal_v2::Signal;
 use restate_types::live::Live;
 use restate_types::net::codec::EncodeError;
@@ -32,6 +35,7 @@ use restate_types::net::partition_processor::{
     PartitionProcessorRpcRequest, PartitionProcessorRpcRequestInner, PartitionProcessorRpcResponse,
 };
 use restate_types::partition_table::{FindPartition, PartitionTable, PartitionTableError};
+use std::time::Duration;
 use tracing::trace;
 
 #[derive(Debug, thiserror::Error)]
@@ -450,11 +454,15 @@ where
         &self,
         request_id: PartitionProcessorRpcRequestId,
         invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
     ) -> Result<PurgeInvocationResponse, InvocationClientError> {
         let response = self
             .resolve_partition_id_and_send(
                 request_id,
-                PartitionProcessorRpcRequestInner::PurgeInvocation { invocation_id },
+                PartitionProcessorRpcRequestInner::PurgeInvocation {
+                    invocation_id,
+                    invocation_epoch,
+                },
             )
             .await?;
 
@@ -472,11 +480,15 @@ where
         &self,
         request_id: PartitionProcessorRpcRequestId,
         invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
     ) -> Result<PurgeInvocationResponse, InvocationClientError> {
         let response = self
             .resolve_partition_id_and_send(
                 request_id,
-                PartitionProcessorRpcRequestInner::PurgeJournal { invocation_id },
+                PartitionProcessorRpcRequestInner::PurgeJournal {
+                    invocation_id,
+                    invocation_epoch,
+                },
             )
             .await?;
 
@@ -489,4 +501,34 @@ where
             }
         })
     }
+
+    async fn restart_invocation(
+        &self,
+        request_id: PartitionProcessorRpcRequestId,
+        invocation_id: InvocationId,
+        if_running: IfRunning,
+        previous_attempt_retention: Option<Duration>,
+        apply_to_workflow_run: ApplyToWorkflowRun,
+    ) -> Result<RestartInvocationResponse, InvocationClientError> {
+        let response = self
+            .resolve_partition_id_and_send(
+                request_id,
+                PartitionProcessorRpcRequestInner::RestartInvocation {
+                    invocation_id,
+                    if_running,
+                    previous_attempt_retention,
+                    apply_to_workflow_run,
+                },
+            )
+            .await?;
+
+        Ok(match response {
+            PartitionProcessorRpcResponse::RestartInvocation(purge_invocation_response) => {
+                purge_invocation_response.into()
+            }
+            _ => {
+                panic!("Expecting RestartInvocation rpc response")
+            }
+        })
+    }
 }

crates/core/src/worker_api/partition_processor_rpc_client.rs

@@ -175,6 +175,7 @@ pub mod codes {
         PROTOCOL_VIOLATION 571 "Protocol violation",
         CONFLICT 409 "Conflict",
         NOT_READY 470 "Not ready",
+        RESTARTED 471 "Restarted",
     );
 }
 
@@ -295,6 +296,11 @@ impl From<anyhow::Error> for InvocationError {
 pub const KILLED_INVOCATION_ERROR: InvocationError =
     InvocationError::new_static(codes::ABORTED, "killed");
 
+pub const RESTARTED_INVOCATION_ERROR: InvocationError = InvocationError::new_static(
+    codes::RESTARTED,
+    "The invocation was restarted. You can re-attach to it to retrieve the new result.",
+);
+
 // TODO: Once we want to distinguish server side cancellations from user code returning the
 //  UserErrorCode::Cancelled, we need to add a new RestateErrorCode.
 pub const CANCELED_INVOCATION_ERROR: InvocationError =

crates/types/src/errors.rs

@@ -10,10 +10,14 @@
 
 use crate::errors::InvocationError;
 use crate::identifiers::{InvocationId, PartitionProcessorRpcRequestId};
-use crate::invocation::{InvocationQuery, InvocationRequest, InvocationResponse, InvocationTarget};
+use crate::invocation::restart::{ApplyToWorkflowRun, IfRunning};
+use crate::invocation::{
+    InvocationEpoch, InvocationQuery, InvocationRequest, InvocationResponse, InvocationTarget,
+};
 use crate::journal_v2::Signal;
 use crate::time::MillisSinceEpoch;
 use bytes::Bytes;
+use std::time::Duration;
 
 #[derive(Debug, thiserror::Error)]
 #[error("{inner}")]
@@ -102,6 +106,16 @@ pub enum PurgeInvocationResponse {
     NotCompleted,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RestartInvocationResponse {
+    Ok,
+    NotFound,
+    StillRunning,
+    Unsupported,
+    MissingInput,
+    NotStarted,
+}
+
 /// This trait provides the functionalities to interact with Restate invocations.
 pub trait InvocationClient {
     /// Append the invocation to the log, waiting for the PP to emit [`SubmittedInvocationNotification`] when the command is processed.
@@ -161,17 +175,29 @@ pub trait InvocationClient {
         invocation_id: InvocationId,
     ) -> impl Future<Output = Result<KillInvocationResponse, InvocationClientError>> + Send;
 
-    /// Purge the given invocation.
+    /// Purge the given invocation. If the invocation_epoch is the latest, all the previous epochs will be purged as well.
     fn purge_invocation(
         &self,
         request_id: PartitionProcessorRpcRequestId,
         invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
     ) -> impl Future<Output = Result<PurgeInvocationResponse, InvocationClientError>> + Send;
 
-    /// Purge the given invocation journal.
+    /// Purge the given invocation journal. If the invocation_epoch is the latest, all the previous epochs will be purged as well.
     fn purge_journal(
         &self,
         request_id: PartitionProcessorRpcRequestId,
         invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
     ) -> impl Future<Output = Result<PurgeInvocationResponse, InvocationClientError>> + Send;
+
+    /// See [`crate::invocation::restart::Request`].
+    fn restart_invocation(
+        &self,
+        request_id: PartitionProcessorRpcRequestId,
+        invocation_id: InvocationId,
+        if_running: IfRunning,
+        previous_attempt_retention: Option<Duration>,
+        apply_to_workflow_run: ApplyToWorkflowRun,
+    ) -> impl Future<Output = Result<RestartInvocationResponse, InvocationClientError>> + Send;
 }

crates/types/src/invocation/client.rs

@@ -75,6 +75,12 @@ impl From<CompletionType> for NotificationType {
     }
 }
 
+impl From<CompletionType> for EntryType {
+    fn from(value: CompletionType) -> Self {
+        EntryType::Notification(value.into())
+    }
+}
+
 #[enum_dispatch]
 pub trait NotificationMetadata {
     fn id(&self) -> NotificationId;

crates/types/src/journal_v2/notification.rs

@@ -13,13 +13,15 @@ use crate::identifiers::{
 };
 use crate::invocation::client::{
     CancelInvocationResponse, InvocationOutput, KillInvocationResponse, PurgeInvocationResponse,
-    SubmittedInvocationNotification,
+    RestartInvocationResponse, SubmittedInvocationNotification,
 };
-use crate::invocation::{InvocationQuery, InvocationRequest, InvocationResponse};
+use crate::invocation::restart::{ApplyToWorkflowRun, IfRunning};
+use crate::invocation::{InvocationEpoch, InvocationQuery, InvocationRequest, InvocationResponse};
 use crate::journal_v2::Signal;
 use crate::net::ServiceTag;
 use crate::net::{default_wire_codec, define_rpc, define_service};
 use serde::{Deserialize, Serialize};
+use std::time::Duration;
 
 pub struct PartitionLeaderService;
 
@@ -68,10 +70,26 @@ pub enum PartitionProcessorRpcRequestInner {
     GetInvocationOutput(InvocationQuery, GetInvocationOutputResponseMode),
     AppendInvocationResponse(InvocationResponse),
     AppendSignal(InvocationId, Signal),
-    CancelInvocation { invocation_id: InvocationId },
-    KillInvocation { invocation_id: InvocationId },
-    PurgeInvocation { invocation_id: InvocationId },
-    PurgeJournal { invocation_id: InvocationId },
+    CancelInvocation {
+        invocation_id: InvocationId,
+    },
+    KillInvocation {
+        invocation_id: InvocationId,
+    },
+    PurgeInvocation {
+        invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
+    },
+    PurgeJournal {
+        invocation_id: InvocationId,
+        invocation_epoch: InvocationEpoch,
+    },
+    RestartInvocation {
+        invocation_id: InvocationId,
+        if_running: IfRunning,
+        previous_attempt_retention: Option<Duration>,
+        apply_to_workflow_run: ApplyToWorkflowRun,
+    },
 }
 
 impl WithPartitionKey for PartitionProcessorRpcRequestInner {
@@ -87,10 +105,13 @@ impl WithPartitionKey for PartitionProcessorRpcRequestInner {
             PartitionProcessorRpcRequestInner::KillInvocation { invocation_id } => {
                 invocation_id.partition_key()
             }
-            PartitionProcessorRpcRequestInner::PurgeInvocation { invocation_id } => {
+            PartitionProcessorRpcRequestInner::PurgeInvocation { invocation_id, .. } => {
                 invocation_id.partition_key()
             }
-            PartitionProcessorRpcRequestInner::PurgeJournal { invocation_id } => {
+            PartitionProcessorRpcRequestInner::PurgeJournal { invocation_id, .. } => {
+                invocation_id.partition_key()
+            }
+            PartitionProcessorRpcRequestInner::RestartInvocation { invocation_id, .. } => {
                 invocation_id.partition_key()
             }
         }
@@ -204,6 +225,42 @@ impl From<PurgeInvocationResponse> for PurgeInvocationRpcResponse {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum RestartInvocationRpcResponse {
+    Ok,
+    NotFound,
+    StillRunning,
+    Unsupported,
+    MissingInput,
+    NotStarted,
+}
+
+impl From<RestartInvocationRpcResponse> for RestartInvocationResponse {
+    fn from(value: RestartInvocationRpcResponse) -> Self {
+        match value {
+            RestartInvocationRpcResponse::Ok => RestartInvocationResponse::Ok,
+            RestartInvocationRpcResponse::NotFound => RestartInvocationResponse::NotFound,
+            RestartInvocationRpcResponse::StillRunning => RestartInvocationResponse::StillRunning,
+            RestartInvocationRpcResponse::Unsupported => RestartInvocationResponse::Unsupported,
+            RestartInvocationRpcResponse::MissingInput => RestartInvocationResponse::MissingInput,
+            RestartInvocationRpcResponse::NotStarted => RestartInvocationResponse::NotStarted,
+        }
+    }
+}
+
+impl From<RestartInvocationResponse> for RestartInvocationRpcResponse {
+    fn from(value: RestartInvocationResponse) -> Self {
+        match value {
+            RestartInvocationResponse::Ok => RestartInvocationRpcResponse::Ok,
+            RestartInvocationResponse::NotFound => RestartInvocationRpcResponse::NotFound,
+            RestartInvocationResponse::StillRunning => RestartInvocationRpcResponse::StillRunning,
+            RestartInvocationResponse::Unsupported => RestartInvocationRpcResponse::Unsupported,
+            RestartInvocationResponse::MissingInput => RestartInvocationRpcResponse::MissingInput,
+            RestartInvocationResponse::NotStarted => RestartInvocationRpcResponse::NotStarted,
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum PartitionProcessorRpcResponse {
     Appended,
@@ -216,4 +273,5 @@ pub enum PartitionProcessorRpcResponse {
     KillInvocation(KillInvocationRpcResponse),
     PurgeInvocation(PurgeInvocationRpcResponse),
     PurgeJournal(PurgeInvocationRpcResponse),
+    RestartInvocation(RestartInvocationRpcResponse),
 }

crates/types/src/net/partition_processor.rs

@@ -13,7 +13,7 @@ use restate_types::GenerationalNodeId;
 use restate_types::identifiers::{LeaderEpoch, PartitionId, PartitionKey, WithPartitionKey};
 use restate_types::invocation::{
     AttachInvocationRequest, GetInvocationOutputResponse, InvocationResponse,
-    InvocationTermination, NotifySignalRequest, PurgeInvocationRequest, ServiceInvocation,
+    InvocationTermination, NotifySignalRequest, PurgeInvocationRequest, ServiceInvocation, restart,
 };
 use restate_types::logs::{HasRecordKeys, Keys, MatchKeyQuery};
 use restate_types::message::MessageIndex;
@@ -148,6 +148,8 @@ pub enum Command {
     ProxyThrough(ServiceInvocation),
     /// Attach to an existing invocation
     AttachInvocation(AttachInvocationRequest),
+    /// Restart an invocation
+    RestartInvocation(restart::Request),
 
     // -- Partition processor events for PP
     /// Invoker is reporting effect(s) from an ongoing invocation.
@@ -203,6 +205,7 @@ impl HasRecordKeys for Envelope {
             }
             Command::PurgeInvocation(purge) => Keys::Single(purge.invocation_id.partition_key()),
             Command::PurgeJournal(purge) => Keys::Single(purge.invocation_id.partition_key()),
+            Command::RestartInvocation(restart) => Keys::Single(restart.partition_key()),
             Command::Invoke(invoke) => Keys::Single(invoke.partition_key()),
             // todo: Remove this, or pass the partition key range but filter based on partition-id
             // on read if needed.
@@ -316,6 +319,7 @@ mod envelope {
         NotifyGetInvocationOutputResponse = 13, // bilrost
         NotifySignal = 14,                      // protobuf
         PurgeJournal = 15,                      // bilrost
+        RestartInvocation = 16,                 // serde
     }
 
     #[derive(bilrost::Message)]
@@ -439,6 +443,10 @@ mod envelope {
             Command::PurgeJournal(value) => {
                 (CommandKind::PurgeJournal, Field::encode_bilrost(value))
             }
+            Command::RestartInvocation(value) => (
+                CommandKind::RestartInvocation,
+                Field::encode_serde(StorageCodecKind::FlexbuffersSerde, value),
+            ),
             Command::Invoke(value) => {
                 let value = protobuf::ServiceInvocation::from(value.clone());
                 (CommandKind::Invoke, Field::encode_protobuf(&value))
@@ -528,6 +536,10 @@ mod envelope {
                 codec_or_error!(envelope.command, StorageCodecKind::Bilrost);
                 Command::PurgeJournal(envelope.command.decode_bilrost()?)
             }
+            CommandKind::RestartInvocation => {
+                codec_or_error!(envelope.command, StorageCodecKind::FlexbuffersSerde);
+                Command::RestartInvocation(envelope.command.decode_serde()?)
+            }
             CommandKind::Invoke => {
                 codec_or_error!(envelope.command, StorageCodecKind::Protobuf);
                 let value: protobuf::ServiceInvocation = envelope.command.decode_protobuf()?;

crates/wal-protocol/src/lib.rs

@@ -501,6 +501,16 @@ impl LeaderState {
                     )));
                 }
             }
+            Action::ForwardRestartInvocationResponse {
+                request_id,
+                response,
+            } => {
+                if let Some(response_tx) = self.awaiting_rpc_actions.remove(&request_id) {
+                    response_tx.send(Ok(PartitionProcessorRpcResponse::RestartInvocation(
+                        response.into(),
+                    )));
+                }
+            }
         }
 
         Ok(())