Implement Cancel restart/Kill restart feature

.github/workflows/ci.yml

@@ -117,6 +117,24 @@ jobs:
         RUST_LOG=info,restate_invoker=trace,restate_ingress_http=trace,restate_bifrost=trace,restate_log_server=trace,restate_core::partitions=trace,restate=debug
       testArtifactOutput: sdk-java-kafka-next-gen-integration-test-report
 
+  sdk-java-invocation-status-killed:
+    name: Run SDK-Java integration tests with InvocationStatusKilled
+    permissions:
+      contents: read
+      issues: read
+      checks: write
+      pull-requests: write
+      actions: read
+    secrets: inherit
+    needs: docker
+    uses: restatedev/sdk-java/.github/workflows/integration.yaml@main
+    with:
+      restateCommit: ${{ github.event.pull_request.head.sha || github.sha }}
+      envVars: |
+        RESTATE_WORKER__EXPERIMENTAL_FEATURE_INVOCATION_STATUS_KILLED=true
+        RUST_LOG=info,restate_invoker=trace,restate_ingress_http=trace,restate_bifrost=trace,restate_log_server=trace,restate_core::partitions=trace,restate=debug
+      testArtifactOutput: sdk-java-invocation-status-killed-integration-test-report
+
   sdk-python:
     name: Run SDK-Python integration tests
     permissions:

cli/src/clients/admin_interface.rs

@@ -41,7 +41,12 @@ pub trait AdminClientInterface {
 
     async fn purge_invocation(&self, id: &str) -> reqwest::Result<Envelope<()>>;
 
-    async fn cancel_invocation(&self, id: &str, kill: bool) -> reqwest::Result<Envelope<()>>;
+    async fn cancel_invocation(
+        &self,
+        id: &str,
+        kill: bool,
+        retry: bool,
+    ) -> reqwest::Result<Envelope<()>>;
 
     async fn patch_state(
         &self,
@@ -132,15 +137,25 @@ impl AdminClientInterface for AdminClient {
         self.run(reqwest::Method::DELETE, url).await
     }
 
-    async fn cancel_invocation(&self, id: &str, kill: bool) -> reqwest::Result<Envelope<()>> {
+    async fn cancel_invocation(
+        &self,
+        id: &str,
+        kill: bool,
+        retry: bool,
+    ) -> reqwest::Result<Envelope<()>> {
         let mut url = self
             .base_url
             .join(&format!("/invocations/{}", id))
             .expect("Bad url!");
 
         url.set_query(Some(&format!(
             "mode={}",
-            if kill { "kill" } else { "cancel" }
+            match (kill, retry) {
+                (false, false) => "cancel",
+                (false, true) => "cancel-and-restart",
+                (true, false) => "kill",
+                (true, true) => "kill-and-restart",
+            }
         )));
 
         self.run(reqwest::Method::DELETE, url).await

cli/src/clients/datafusion_helpers.rs

@@ -128,6 +128,7 @@ pub enum InvocationState {
     Running,
     Suspended,
     BackingOff,
+    Killed,
     Completed,
 }
 
@@ -142,6 +143,7 @@ impl FromStr for InvocationState {
             "suspended" => Self::Suspended,
             "backing-off" => Self::BackingOff,
             "completed" => Self::Completed,
+            "killed" => Self::Killed,
             _ => Self::Unknown,
         })
     }
@@ -157,6 +159,7 @@ impl Display for InvocationState {
             InvocationState::Running => write!(f, "running"),
             InvocationState::Suspended => write!(f, "suspended"),
             InvocationState::BackingOff => write!(f, "backing-off"),
+            InvocationState::Killed => write!(f, "killed"),
             InvocationState::Completed => write!(f, "completed"),
         }
     }

cli/src/commands/invocations/cancel.rs

@@ -37,6 +37,9 @@ pub struct Cancel {
     /// Ungracefully kill the invocation and its children
     #[clap(long)]
     kill: bool,
+    /// After cancelling/killing, restart the invocation using the same input.
+    #[clap(long, alias = "retry")]
+    restart: bool,
 }
 
 pub async fn run_cancel(State(env): State<CliEnv>, opts: &Cancel) -> Result<()> {
@@ -67,16 +70,19 @@ pub async fn run_cancel(State(env): State<CliEnv>, opts: &Cancel) -> Result<()>
     // Get the invocation and confirm
     let prompt = format!(
         "Are you sure you want to {} these invocations?",
-        if opts.kill {
-            Styled(Style::Danger, "kill")
-        } else {
-            Styled(Style::Warn, "cancel")
-        },
+        match (opts.kill, opts.restart) {
+            (false, false) => Styled(Style::Warn, "cancel"),
+            (false, true) => Styled(Style::Warn, "cancel and restart"),
+            (true, false) => Styled(Style::Danger, "kill"),
+            (true, true) => Styled(Style::Danger, "kill and restart"),
+        }
     );
     confirm_or_exit(&prompt)?;
 
     for inv in invocations {
-        let result = client.cancel_invocation(&inv.id, opts.kill).await?;
+        let result = client
+            .cancel_invocation(&inv.id, opts.kill, opts.restart)
+            .await?;
         let _ = result.success_or_error()?;
     }
 

cli/src/ui/invocations.rs

@@ -113,6 +113,7 @@ pub fn invocation_status_style(status: InvocationState) -> Style {
         InvocationState::Suspended => DStyle::new().dim(),
         InvocationState::BackingOff => DStyle::new().red(),
         InvocationState::Completed => DStyle::new().blue(),
+        InvocationState::Killed => DStyle::new().red(),
     }
 }
 

crates/admin/src/rest_api/invocations.rs

@@ -31,6 +31,10 @@ pub enum DeletionMode {
     Kill,
     #[serde(alias = "purge")]
     Purge,
+    #[serde(alias = "kill-and-restart")]
+    KillAndRestart,
+    #[serde(alias = "cancel-and-restart")]
+    CancelAndRestart,
 }
 #[derive(Debug, Default, Deserialize, JsonSchema)]
 pub struct DeleteInvocationParams {
@@ -90,6 +94,12 @@ pub async fn delete_invocation<V>(
             Command::TerminateInvocation(InvocationTermination::kill(invocation_id))
         }
         DeletionMode::Purge => Command::PurgeInvocation(PurgeInvocationRequest { invocation_id }),
+        DeletionMode::CancelAndRestart => {
+            Command::TerminateInvocation(InvocationTermination::cancel_and_restart(invocation_id))
+        }
+        DeletionMode::KillAndRestart => {
+            Command::TerminateInvocation(InvocationTermination::kill_and_restart(invocation_id))
+        }
     };
 
     let partition_key = invocation_id.partition_key();

crates/invoker-api/src/handle.rs

@@ -60,6 +60,8 @@ pub trait InvokerHandle<SR> {
         &mut self,
         partition_leader_epoch: PartitionLeaderEpoch,
         invocation_id: InvocationId,
+        // If true, acknowledge the abort. This will generate a Failed effect
+        acknowledge: bool,
     ) -> impl Future<Output = Result<(), NotRunningError>> + Send;
 
     fn register_partition(

crates/invoker-api/src/lib.rs

@@ -128,6 +128,7 @@ pub mod test_util {
             &mut self,
             _partition_leader_epoch: PartitionLeaderEpoch,
             _invocation_id: InvocationId,
+            _acknowledge: bool,
         ) -> Result<(), NotRunningError> {
             Ok(())
         }

crates/invoker-impl/src/input_command.rs

@@ -45,6 +45,7 @@ pub(crate) enum InputCommand<SR> {
     Abort {
         partition: PartitionLeaderEpoch,
         invocation_id: InvocationId,
+        acknowledge: bool,
     },
 
     /// Command used to clean up internal state when a partition leader is going away
@@ -129,11 +130,13 @@ impl<SR: Send> restate_invoker_api::InvokerHandle<SR> for InvokerHandle<SR> {
         &mut self,
         partition: PartitionLeaderEpoch,
         invocation_id: InvocationId,
+        acknowledge: bool,
     ) -> Result<(), NotRunningError> {
         self.input
             .send(InputCommand::Abort {
                 partition,
                 invocation_id,
+                acknowledge,
             })
             .map_err(|_| NotRunningError)
     }

crates/invoker-impl/src/lib.rs

@@ -61,6 +61,7 @@ pub use input_command::ChannelStatusReader;
 pub use input_command::InvokerHandle;
 use restate_service_client::{AssumeRoleCacheMode, ServiceClient};
 use restate_types::deployment::PinnedDeployment;
+use restate_types::errors::KILLED_INVOCATION_ERROR;
 use restate_types::invocation::InvocationTarget;
 use restate_types::schema::service::ServiceMetadataResolver;
 
@@ -351,8 +352,8 @@ where
                         self.handle_register_partition(partition, partition_key_range,
                                 storage_reader, sender);
                     },
-                    InputCommand::Abort { partition, invocation_id } => {
-                        self.handle_abort_invocation(partition, invocation_id);
+                    InputCommand::Abort { partition, invocation_id,  acknowledge } => {
+                        self.handle_abort_invocation(partition, invocation_id, acknowledge).await;
                     }
                     InputCommand::AbortAllPartition { partition } => {
                         self.handle_abort_partition(partition);
@@ -808,12 +809,13 @@ where
             restate.invoker.partition_leader_epoch = ?partition,
         )
     )]
-    fn handle_abort_invocation(
+    async fn handle_abort_invocation(
         &mut self,
         partition: PartitionLeaderEpoch,
         invocation_id: InvocationId,
+        acknowledge: bool,
     ) {
-        if let Some((_, _, mut ism)) = self
+        if let Some((tx, _, mut ism)) = self
             .invocation_state_machine_manager
             .remove_invocation(partition, &invocation_id)
         {
@@ -823,6 +825,14 @@ where
             ism.abort();
             self.quota.unreserve_slot();
             self.status_store.on_end(&partition, &invocation_id);
+            if acknowledge {
+                let _ = tx
+                    .send(Effect {
+                        invocation_id,
+                        kind: EffectKind::Failed(KILLED_INVOCATION_ERROR),
+                    })
+                    .await;
+            }
         } else {
             trace!("Ignoring Abort command because there is no matching partition/invocation");
         }
@@ -1415,7 +1425,9 @@ mod tests {
         assert_eq!(*available_slots, 1);
 
         // Abort the invocation
-        service_inner.handle_abort_invocation(MOCK_PARTITION, invocation_id);
+        service_inner
+            .handle_abort_invocation(MOCK_PARTITION, invocation_id, false)
+            .await;
 
         // Check the quota
         let_assert!(InvokerConcurrencyQuota::Limited { available_slots } = &service_inner.quota);

crates/partition-store/src/invocation_status_table/mod.rs

@@ -17,11 +17,11 @@ use futures::Stream;
 use futures_util::stream;
 use restate_rocksdb::RocksDbPerfGuard;
 use restate_storage_api::invocation_status_table::{
-    InvocationStatus, InvocationStatusTable, InvocationStatusV1, ReadOnlyInvocationStatusTable,
+    InvocationStatus, InvocationStatusTable, InvocationStatusV1,
+    InvokedOrKilledInvocationStatusLite, ReadOnlyInvocationStatusTable,
 };
 use restate_storage_api::{Result, StorageError};
 use restate_types::identifiers::{InvocationId, InvocationUuid, PartitionKey, WithPartitionKey};
-use restate_types::invocation::InvocationTarget;
 use restate_types::storage::StorageCodec;
 use std::ops::RangeInclusive;
 use tracing::trace;
@@ -169,7 +169,7 @@ fn delete_invocation_status<S: StorageAccess>(storage: &mut S, invocation_id: &I
 fn invoked_invocations<S: StorageAccess>(
     storage: &mut S,
     partition_key_range: RangeInclusive<PartitionKey>,
-) -> Vec<Result<(InvocationId, InvocationTarget)>> {
+) -> Vec<Result<InvokedOrKilledInvocationStatusLite>> {
     let _x = RocksDbPerfGuard::new("invoked-invocations");
     let mut invocations = storage.for_each_key_value_in_place(
         FullScanPartitionKeyRange::<InvocationStatusKeyV1>(partition_key_range.clone()),
@@ -239,12 +239,16 @@ fn all_invocation_status<S: StorageAccess>(
 fn read_invoked_v1_full_invocation_id(
     mut k: &mut &[u8],
     v: &mut &[u8],
-) -> Result<Option<(InvocationId, InvocationTarget)>> {
+) -> Result<Option<InvokedOrKilledInvocationStatusLite>> {
     let invocation_id = invocation_id_from_v1_key_bytes(&mut k)?;
     let invocation_status = StorageCodec::decode::<InvocationStatusV1, _>(v)
         .map_err(|err| StorageError::Generic(err.into()))?;
     if let InvocationStatus::Invoked(invocation_meta) = invocation_status.0 {
-        Ok(Some((invocation_id, invocation_meta.invocation_target)))
+        Ok(Some(InvokedOrKilledInvocationStatusLite {
+            invocation_id,
+            invocation_target: invocation_meta.invocation_target,
+            is_invoked: true,
+        }))
     } else {
         Ok(None)
     }
@@ -253,13 +257,23 @@ fn read_invoked_v1_full_invocation_id(
 fn read_invoked_full_invocation_id(
     mut k: &mut &[u8],
     v: &mut &[u8],
-) -> Result<Option<(InvocationId, InvocationTarget)>> {
+) -> Result<Option<InvokedOrKilledInvocationStatusLite>> {
     // TODO this can be improved by simply parsing InvocationTarget and the Status enum
     let invocation_id = invocation_id_from_key_bytes(&mut k)?;
     let invocation_status = StorageCodec::decode::<InvocationStatus, _>(v)
         .map_err(|err| StorageError::Generic(err.into()))?;
     if let InvocationStatus::Invoked(invocation_meta) = invocation_status {
-        Ok(Some((invocation_id, invocation_meta.invocation_target)))
+        Ok(Some(InvokedOrKilledInvocationStatusLite {
+            invocation_id,
+            invocation_target: invocation_meta.invocation_target,
+            is_invoked: true,
+        }))
+    } else if let InvocationStatus::Killed(invocation_meta) = invocation_status {
+        Ok(Some(InvokedOrKilledInvocationStatusLite {
+            invocation_id,
+            invocation_target: invocation_meta.invocation_target,
+            is_invoked: false,
+        }))
     } else {
         Ok(None)
     }
@@ -274,9 +288,9 @@ impl ReadOnlyInvocationStatusTable for PartitionStore {
         get_invocation_status(self, invocation_id)
     }
 
-    fn all_invoked_invocations(
+    fn all_invoked_or_killed_invocations(
         &mut self,
-    ) -> impl Stream<Item = Result<(InvocationId, InvocationTarget)>> + Send {
+    ) -> impl Stream<Item = Result<InvokedOrKilledInvocationStatusLite>> + Send {
         stream::iter(invoked_invocations(
             self,
             self.partition_key_range().clone(),
@@ -300,9 +314,9 @@ impl<'a> ReadOnlyInvocationStatusTable for PartitionStoreTransaction<'a> {
         try_migrate_and_get_invocation_status(self, invocation_id)
     }
 
-    fn all_invoked_invocations(
+    fn all_invoked_or_killed_invocations(
         &mut self,
-    ) -> impl Stream<Item = Result<(InvocationId, InvocationTarget)>> + Send {
+    ) -> impl Stream<Item = Result<InvokedOrKilledInvocationStatusLite>> + Send {
         stream::iter(invoked_invocations(
             self,
             self.partition_key_range().clone(),