Return the new invocation epoch in restart_invocation API

Author: slinkydeveloper

crates/admin/src/rest_api/invocations.rs

@@ -12,19 +12,20 @@ use super::error::*;
 use crate::generate_meta_api_error;
 use crate::rest_api::create_envelope_header;
 use crate::state::AdminServiceState;
+use axum::Json;
 use axum::extract::{Path, Query, State};
 use axum::http::StatusCode;
 use okapi_operation::*;
 use restate_types::identifiers::{InvocationId, PartitionProcessorRpcRequestId, WithPartitionKey};
 use restate_types::invocation::client::{
-    CancelInvocationResponse, InvocationClient, KillInvocationResponse, PurgeInvocationResponse,
-    RestartInvocationResponse,
+    self, CancelInvocationResponse, InvocationClient, KillInvocationResponse,
+    PurgeInvocationResponse,
 };
 use restate_types::invocation::{
     InvocationEpoch, InvocationTermination, PurgeInvocationRequest, TerminationFlavor, restart,
 };
 use restate_wal_protocol::{Command, Envelope};
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use std::time::Duration;
 use tracing::warn;
@@ -343,6 +344,12 @@ where
     Ok(())
 }
 
+#[derive(Debug, Serialize, JsonSchema)]
+pub struct RestartInvocationResponse {
+    /// The new invocation epoch of the invocation.
+    pub new_invocation_epoch: InvocationEpoch,
+}
+
 /// What to do if the invocation is still running. By default, the running invocation will be killed.
 #[derive(Default, Debug, Deserialize, JsonSchema)]
 #[serde(rename_all = "snake_case")]
@@ -462,7 +469,7 @@ pub async fn restart_invocation<V, IC>(
         previous_attempt_retention,
         apply_to_workflow_run,
     }): Query<RestartInvocationParams>,
-) -> Result<(), RestartInvocationError>
+) -> Result<Json<RestartInvocationResponse>, RestartInvocationError>
 where
     IC: InvocationClient,
 {
@@ -482,23 +489,24 @@ where
         .await
         .map_err(InvocationClientError)?
     {
-        RestartInvocationResponse::Ok => {}
-        RestartInvocationResponse::NotFound => {
+        client::RestartInvocationResponse::Ok { new_epoch } => Ok(RestartInvocationResponse {
+            new_invocation_epoch: new_epoch,
+        }
+        .into()),
+        client::RestartInvocationResponse::NotFound => {
             Err(InvocationNotFoundError(invocation_id.to_string()))?
         }
-        RestartInvocationResponse::StillRunning => Err(RestartInvocationStillRunningError(
-            invocation_id.to_string(),
-        ))?,
-        RestartInvocationResponse::Unsupported => {
+        client::RestartInvocationResponse::StillRunning => Err(
+            RestartInvocationStillRunningError(invocation_id.to_string()),
+        )?,
+        client::RestartInvocationResponse::Unsupported => {
             Err(RestartInvocationUnsupportedError(invocation_id.to_string()))?
         }
-        RestartInvocationResponse::MissingInput => Err(RestartInvocationMissingInputError(
-            invocation_id.to_string(),
-        ))?,
-        RestartInvocationResponse::NotStarted => {
+        client::RestartInvocationResponse::MissingInput => Err(
+            RestartInvocationMissingInputError(invocation_id.to_string()),
+        )?,
+        client::RestartInvocationResponse::NotStarted => {
             Err(RestartInvocationNotStartedError(invocation_id.to_string()))?
         }
-    };
-
-    Ok(())
+    }
 }

crates/types/src/invocation/client.rs

@@ -108,7 +108,7 @@ pub enum PurgeInvocationResponse {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum RestartInvocationResponse {
-    Ok,
+    Ok { new_epoch: InvocationEpoch },
     NotFound,
     StillRunning,
     Unsupported,

crates/types/src/net/partition_processor.rs

@@ -227,7 +227,7 @@ impl From<PurgeInvocationResponse> for PurgeInvocationRpcResponse {
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum RestartInvocationRpcResponse {
-    Ok,
+    Ok { new_epoch: InvocationEpoch },
     NotFound,
     StillRunning,
     Unsupported,
@@ -238,7 +238,9 @@ pub enum RestartInvocationRpcResponse {
 impl From<RestartInvocationRpcResponse> for RestartInvocationResponse {
     fn from(value: RestartInvocationRpcResponse) -> Self {
         match value {
-            RestartInvocationRpcResponse::Ok => RestartInvocationResponse::Ok,
+            RestartInvocationRpcResponse::Ok { new_epoch } => {
+                RestartInvocationResponse::Ok { new_epoch }
+            }
             RestartInvocationRpcResponse::NotFound => RestartInvocationResponse::NotFound,
             RestartInvocationRpcResponse::StillRunning => RestartInvocationResponse::StillRunning,
             RestartInvocationRpcResponse::Unsupported => RestartInvocationResponse::Unsupported,
@@ -251,7 +253,9 @@ impl From<RestartInvocationRpcResponse> for RestartInvocationResponse {
 impl From<RestartInvocationResponse> for RestartInvocationRpcResponse {
     fn from(value: RestartInvocationResponse) -> Self {
         match value {
-            RestartInvocationResponse::Ok => RestartInvocationRpcResponse::Ok,
+            RestartInvocationResponse::Ok { new_epoch } => {
+                RestartInvocationRpcResponse::Ok { new_epoch }
+            }
             RestartInvocationResponse::NotFound => RestartInvocationRpcResponse::NotFound,
             RestartInvocationResponse::StillRunning => RestartInvocationRpcResponse::StillRunning,
             RestartInvocationResponse::Unsupported => RestartInvocationRpcResponse::Unsupported,

crates/worker/src/partition/state_machine/lifecycle/restart.rs

@@ -315,7 +315,7 @@ where
         .await?;
 
         // Reply to the listener, restart went well
-        ctx.reply_to_restart_invocation(response_sink, RestartInvocationResponse::Ok);
+        ctx.reply_to_restart_invocation(response_sink, RestartInvocationResponse::Ok { new_epoch });
 
         Ok(())
     }
@@ -478,7 +478,7 @@ mod tests {
             all!(
                 contains(pat!(Action::ForwardRestartInvocationResponse {
                     request_id: eq(restart_request_id),
-                    response: eq(RestartInvocationResponse::Ok)
+                    response: eq(RestartInvocationResponse::Ok { new_epoch: 1 })
                 })),
                 contains(pat!(Action::Invoke {
                     invocation_id: eq(invocation_id),
@@ -603,7 +603,7 @@ mod tests {
                 // Verify the restart response is sent
                 contains(pat!(Action::ForwardRestartInvocationResponse {
                     request_id: eq(restart_request_id),
-                    response: eq(RestartInvocationResponse::Ok)
+                    response: eq(RestartInvocationResponse::Ok { new_epoch: 1 })
                 })),
                 // Verify the invocation is restarted
                 contains(pat!(Action::Invoke {
@@ -787,7 +787,7 @@ mod tests {
             actions,
             contains(pat!(Action::ForwardRestartInvocationResponse {
                 request_id: eq(restart_request_id),
-                response: eq(RestartInvocationResponse::Ok)
+                response: eq(RestartInvocationResponse::Ok { new_epoch: 1 })
             }))
         );
 
@@ -927,7 +927,7 @@ mod tests {
                 // Verify the restart response is sent
                 contains(pat!(Action::ForwardRestartInvocationResponse {
                     request_id: eq(restart_request_id),
-                    response: eq(RestartInvocationResponse::Ok)
+                    response: eq(RestartInvocationResponse::Ok { new_epoch: 1 })
                 })),
                 // Verify the invocation is restarted
                 contains(pat!(Action::Invoke {