Introduce graceful cancellation of invocations

This commit introduces a new TerminationFlavor::Cancel which will
gracefully terminate an invocation. The way it works is that only
leaf operations are being terminated with a CANCELED_INVOCATION_ERROR.
Leaf operations are all completable journal entries except for calls
and inboxed invocations. For pending calls in a journal, a cancellation
will propagate the cancellation to these calls.

This commit adds support for canceling inboxed, invoked, suspended and
virtual invocations.

This fixes restatedev#986 and restatedev#943.

crates/storage-api/src/outbox_table/mod.rs

@@ -10,7 +10,9 @@
 
 use crate::{GetFuture, PutFuture};
 use restate_types::identifiers::{FullInvocationId, IngressDispatcherId, PartitionId};
-use restate_types::invocation::{InvocationResponse, ResponseResult, ServiceInvocation};
+use restate_types::invocation::{
+    InvocationResponse, InvocationTermination, ResponseResult, ServiceInvocation,
+};
 use std::ops::Range;
 
 /// Types of outbox messages.
@@ -29,8 +31,8 @@ pub enum OutboxMessage {
         response: ResponseResult,
     },
 
-    /// Kill command to send to another partition processor
-    Kill(FullInvocationId),
+    /// Terminate invocation to send to another partition processor
+    TerminateInvocation(InvocationTermination),
 }
 
 pub trait OutboxTable {

crates/storage-proto/proto/dev/restate/storage/v1/domain.proto

@@ -14,6 +14,13 @@ message FullInvocationId {
     bytes invocation_uuid = 3;
 }
 
+message MaybeFullInvocationId {
+    oneof kind {
+        FullInvocationId full_invocation_id = 1;
+        bytes invocation_id = 2;
+    }
+}
+
 // ---------------------------------------------------------------------
 // Service Invocation
 // ---------------------------------------------------------------------
@@ -288,10 +295,7 @@ message OutboxMessage {
     }
 
     message OutboxServiceInvocationResponse {
-        oneof id {
-            FullInvocationId full_invocation_id = 1;
-            bytes invocation_id = 4;
-        }
+        MaybeFullInvocationId maybe_fid = 1;
         uint32 entry_index = 2;
         ResponseResult response_result = 3;
     }
@@ -303,14 +307,19 @@ message OutboxMessage {
     }
 
     message OutboxKill {
-        FullInvocationId full_invocation_id = 1;
+        MaybeFullInvocationId maybe_full_invocation_id = 1;
+    }
+
+    message OutboxCancel {
+        MaybeFullInvocationId maybe_full_invocation_id = 1;
     }
 
     oneof outbox_message {
         OutboxServiceInvocation service_invocation_case = 1;
         OutboxServiceInvocationResponse service_invocation_response = 2;
         OutboxIngressResponse ingress_response = 3;
         OutboxKill kill = 4;
+        OutboxCancel cancel = 5;
     }
 
 }

crates/storage-proto/src/lib.rs

@@ -24,38 +24,35 @@ pub mod storage {
                 Awakeable, BackgroundCall, ClearState, CompleteAwakeable, Custom, GetState, Invoke,
                 OutputStream, PollInputStream, SetState, Sleep,
             };
-            use crate::storage::v1::invocation_status::{
-                invoked, suspended, Free, Invoked, Suspended, Virtual,
-            };
+            use crate::storage::v1::invocation_status::{Free, Invoked, Suspended, Virtual};
             use crate::storage::v1::journal_entry::completion_result::{
                 Ack, Empty, Failure, Success,
             };
             use crate::storage::v1::journal_entry::{
                 completion_result, CompletionResult, Entry, Kind,
             };
             use crate::storage::v1::outbox_message::{
-                OutboxIngressResponse, OutboxKill, OutboxServiceInvocation,
+                OutboxCancel, OutboxIngressResponse, OutboxKill, OutboxServiceInvocation,
                 OutboxServiceInvocationResponse,
             };
             use crate::storage::v1::service_invocation_response_sink::{
                 Ingress, NewInvocation, PartitionProcessor, ResponseSink,
             };
             use crate::storage::v1::{
                 enriched_entry_header, invocation_resolution_result, invocation_status,
-                outbox_message, outbox_message::outbox_service_invocation_response,
-                response_result, span_relation, timer, BackgroundCallResolutionResult,
-                EnrichedEntryHeader, FullInvocationId, InboxEntry, InvocationResolutionResult,
-                InvocationStatus, JournalEntry, JournalMeta, OutboxMessage, ResponseResult,
-                SequencedTimer, ServiceInvocation, ServiceInvocationResponseSink, SpanContext,
-                SpanRelation, Timer,
+                maybe_full_invocation_id, outbox_message, response_result, span_relation, timer,
+                BackgroundCallResolutionResult, EnrichedEntryHeader, FullInvocationId, InboxEntry,
+                InvocationResolutionResult, InvocationStatus, JournalEntry, JournalMeta,
+                MaybeFullInvocationId, OutboxMessage, ResponseResult, SequencedTimer,
+                ServiceInvocation, ServiceInvocationResponseSink, SpanContext, SpanRelation, Timer,
             };
             use anyhow::anyhow;
             use bytes::{Buf, Bytes};
             use bytestring::ByteString;
             use opentelemetry_api::trace::TraceState;
             use restate_storage_api::StorageError;
             use restate_types::identifiers::ServiceId;
-            use restate_types::invocation::MaybeFullInvocationId;
+            use restate_types::invocation::{InvocationTermination, TerminationFlavor};
             use restate_types::time::MillisSinceEpoch;
             use std::collections::{HashSet, VecDeque};
             use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
@@ -592,6 +589,49 @@ pub mod storage {
                 }
             }
 
+            impl TryFrom<MaybeFullInvocationId> for restate_types::invocation::MaybeFullInvocationId {
+                type Error = ConversionError;
+
+                fn try_from(value: MaybeFullInvocationId) -> Result<Self, Self::Error> {
+                    match value.kind.ok_or(ConversionError::missing_field("kind"))? {
+                        maybe_full_invocation_id::Kind::FullInvocationId(fid) => {
+                            Ok(restate_types::invocation::MaybeFullInvocationId::Full(
+                                restate_types::identifiers::FullInvocationId::try_from(fid)?,
+                            ))
+                        }
+                        maybe_full_invocation_id::Kind::InvocationId(invocation_id) => {
+                            Ok(restate_types::invocation::MaybeFullInvocationId::Partial(
+                                restate_types::identifiers::InvocationId::from_slice(
+                                    &invocation_id,
+                                )
+                                .map_err(|e| ConversionError::invalid_data(e))?,
+                            ))
+                        }
+                    }
+                }
+            }
+
+            impl From<restate_types::invocation::MaybeFullInvocationId> for MaybeFullInvocationId {
+                fn from(value: restate_types::invocation::MaybeFullInvocationId) -> Self {
+                    match value {
+                        restate_types::invocation::MaybeFullInvocationId::Full(fid) => {
+                            MaybeFullInvocationId {
+                                kind: Some(maybe_full_invocation_id::Kind::FullInvocationId(
+                                    FullInvocationId::from(fid),
+                                )),
+                            }
+                        }
+                        restate_types::invocation::MaybeFullInvocationId::Partial(
+                            invocation_id,
+                        ) => MaybeFullInvocationId {
+                            kind: Some(maybe_full_invocation_id::Kind::InvocationId(
+                                Bytes::copy_from_slice(&invocation_id.as_bytes()),
+                            )),
+                        },
+                    }
+                }
+            }
+
             fn try_bytes_into_invocation_uuid(
                 bytes: Bytes,
             ) -> Result<restate_types::identifiers::InvocationUuid, ConversionError> {
@@ -1238,26 +1278,10 @@ pub mod storage {
                         ) => restate_storage_api::outbox_table::OutboxMessage::ServiceResponse(
                             restate_types::invocation::InvocationResponse {
                                 entry_index: invocation_response.entry_index,
-                                id: match invocation_response
-                                    .id
-                                    .ok_or(ConversionError::missing_field("id"))?
-                                {
-                                    outbox_service_invocation_response::Id::FullInvocationId(
-                                        fid,
-                                    ) => MaybeFullInvocationId::Full(
-                                        restate_types::identifiers::FullInvocationId::try_from(
-                                            fid,
-                                        )?,
-                                    ),
-                                    outbox_service_invocation_response::Id::InvocationId(
-                                        invocation_id_bytes,
-                                    ) => MaybeFullInvocationId::Partial(
-                                        restate_types::identifiers::InvocationId::from_slice(
-                                            &invocation_id_bytes,
-                                        )
-                                        .map_err(ConversionError::invalid_data)?,
-                                    ),
-                                },
+                                id: invocation_response
+                                    .maybe_fid
+                                    .ok_or(ConversionError::missing_field("maybe_fid"))?
+                                    .try_into()?,
                                 result: restate_types::invocation::ResponseResult::try_from(
                                     invocation_response
                                         .response_result
@@ -1284,11 +1308,27 @@ pub mod storage {
                             }
                         }
                         outbox_message::OutboxMessage::Kill(outbox_kill) => {
-                            let fid = outbox_kill
-                                .full_invocation_id
-                                .ok_or(ConversionError::missing_field("full_invocation_id"))?;
-                            restate_storage_api::outbox_table::OutboxMessage::Kill(
-                                restate_types::identifiers::FullInvocationId::try_from(fid)?,
+                            let maybe_fid = outbox_kill.maybe_full_invocation_id.ok_or(
+                                ConversionError::missing_field("maybe_full_invocation_id"),
+                            )?;
+                            restate_storage_api::outbox_table::OutboxMessage::TerminateInvocation(
+                                InvocationTermination::kill(
+                                    restate_types::invocation::MaybeFullInvocationId::try_from(
+                                        maybe_fid,
+                                    )?,
+                                ),
+                            )
+                        }
+                        outbox_message::OutboxMessage::Cancel(outbox_cancel) => {
+                            let maybe_fid = outbox_cancel.maybe_full_invocation_id.ok_or(
+                                ConversionError::missing_field("maybe_full_invocation_id"),
+                            )?;
+                            restate_storage_api::outbox_table::OutboxMessage::TerminateInvocation(
+                                InvocationTermination::cancel(
+                                    restate_types::invocation::MaybeFullInvocationId::try_from(
+                                        maybe_fid,
+                                    )?,
+                                ),
                             )
                         }
                     };
@@ -1314,18 +1354,9 @@ pub mod storage {
                         ) => outbox_message::OutboxMessage::ServiceInvocationResponse(
                             OutboxServiceInvocationResponse {
                                 entry_index: invocation_response.entry_index,
-                                id: Some(match invocation_response.id {
-                                    MaybeFullInvocationId::Partial(iid) => {
-                                        outbox_service_invocation_response::Id::InvocationId(
-                                            Bytes::copy_from_slice(&iid.as_bytes()),
-                                        )
-                                    }
-                                    MaybeFullInvocationId::Full(fid) => {
-                                        outbox_service_invocation_response::Id::FullInvocationId(
-                                            FullInvocationId::from(fid),
-                                        )
-                                    }
-                                }),
+                                maybe_fid: Some(MaybeFullInvocationId::from(
+                                    invocation_response.id,
+                                )),
                                 response_result: Some(ResponseResult::from(
                                     invocation_response.result,
                                 )),
@@ -1344,11 +1375,24 @@ pub mod storage {
                                 response_result: Some(ResponseResult::from(response)),
                             })
                         }
-                        restate_storage_api::outbox_table::OutboxMessage::Kill(fid) => {
-                            outbox_message::OutboxMessage::Kill(OutboxKill {
-                                full_invocation_id: Some(FullInvocationId::from(fid)),
-                            })
-                        }
+                        restate_storage_api::outbox_table::OutboxMessage::TerminateInvocation(
+                            invocation_termination,
+                        ) => match invocation_termination.flavor {
+                            TerminationFlavor::Kill => {
+                                outbox_message::OutboxMessage::Kill(OutboxKill {
+                                    maybe_full_invocation_id: Some(MaybeFullInvocationId::from(
+                                        invocation_termination.maybe_fid,
+                                    )),
+                                })
+                            }
+                            TerminationFlavor::Cancel => {
+                                outbox_message::OutboxMessage::Cancel(OutboxCancel {
+                                    maybe_full_invocation_id: Some(MaybeFullInvocationId::from(
+                                        invocation_termination.maybe_fid,
+                                    )),
+                                })
+                            }
+                        },
                     };
 
                     OutboxMessage {

crates/types/src/errors.rs

@@ -322,6 +322,11 @@ pub const KILLED_INVOCATION_ERROR: InvocationError = InvocationError::new_static
     "killed",
 );
 
+pub const CANCELED_INVOCATION_ERROR: InvocationError = InvocationError::new_static(
+    InvocationErrorCode::User(UserErrorCode::Cancelled),
+    "canceled",
+);
+
 #[cfg(feature = "tonic_conversions")]
 mod tonic_conversions_impl {
     use super::{InvocationError, InvocationErrorCode};

crates/types/src/invocation.rs

@@ -379,6 +379,38 @@ impl SpanRelation {
     }
 }
 
+/// Message to terminate an invocation.
+#[derive(Debug, Clone, PartialEq)]
+pub struct InvocationTermination {
+    pub maybe_fid: MaybeFullInvocationId,
+    pub flavor: TerminationFlavor,
+}
+
+impl InvocationTermination {
+    pub fn kill(maybe_fid: impl Into<MaybeFullInvocationId>) -> Self {
+        Self {
+            maybe_fid: maybe_fid.into(),
+            flavor: TerminationFlavor::Kill,
+        }
+    }
+
+    pub fn cancel(maybe_fid: impl Into<MaybeFullInvocationId>) -> Self {
+        Self {
+            maybe_fid: maybe_fid.into(),
+            flavor: TerminationFlavor::Cancel,
+        }
+    }
+}
+
+/// Flavor of the termination. Can be kill (hard stop) or graceful cancel.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum TerminationFlavor {
+    // hard termination, no clean up
+    Kill,
+    // graceful termination allowing the invocation to clean up
+    Cancel,
+}
+
 #[cfg(any(test, feature = "mocks"))]
 mod mocks {
     use super::*;

crates/worker/src/network_integration.rs

@@ -139,7 +139,7 @@ mod shuffle_integration {
     use restate_types::errors::InvocationError;
     use restate_types::identifiers::WithPartitionKey;
     use restate_types::identifiers::{PartitionId, PartitionKey, PeerId};
-    use restate_types::invocation::{MaybeFullInvocationId, ResponseResult};
+    use restate_types::invocation::ResponseResult;
     use restate_types::message::MessageIndex;
 
     #[derive(Debug)]
@@ -159,7 +159,9 @@ mod shuffle_integration {
                 shuffle::PartitionProcessorMessage::Response(response) => {
                     response.id.partition_key()
                 }
-                PartitionProcessorMessage::Kill(fid) => fid.partition_key(),
+                PartitionProcessorMessage::TerminateInvocation(invocation_termination) => {
+                    invocation_termination.maybe_fid.partition_key()
+                }
             }
         }
     }
@@ -192,9 +194,9 @@ mod shuffle_integration {
                         deduplication_source,
                     )
                 }
-                shuffle::PartitionProcessorMessage::Kill(fid) => {
+                shuffle::PartitionProcessorMessage::TerminateInvocation(invocation_termination) => {
                     partition::StateMachineAckCommand::dedup(
-                        partition::StateMachineCommand::Kill(MaybeFullInvocationId::from(fid)),
+                        partition::StateMachineCommand::TerminateInvocation(invocation_termination),
                         deduplication_source,
                     )
                 }

crates/worker/src/partition/shuffle.rs

@@ -12,7 +12,9 @@ use crate::partition::shuffle::state_machine::StateMachine;
 use futures::future::BoxFuture;
 use restate_storage_api::outbox_table::OutboxMessage;
 use restate_types::identifiers::{FullInvocationId, IngressDispatcherId, PartitionId, PeerId};
-use restate_types::invocation::{InvocationResponse, ResponseResult, ServiceInvocation};
+use restate_types::invocation::{
+    InvocationResponse, InvocationTermination, ResponseResult, ServiceInvocation,
+};
 use restate_types::message::{AckKind, MessageIndex};
 use std::time::Duration;
 use tokio::sync::mpsc;
@@ -53,7 +55,7 @@ pub(crate) struct ShuffleInput(pub(crate) AckKind);
 pub(crate) enum PartitionProcessorMessage {
     Invocation(ServiceInvocation),
     Response(InvocationResponse),
-    Kill(FullInvocationId),
+    TerminateInvocation(InvocationTermination),
 }
 
 #[derive(Debug, Clone)]
@@ -126,8 +128,10 @@ impl From<OutboxMessage> for ShuffleMessageDestination {
                     PartitionProcessorMessage::Invocation(invocation),
                 )
             }
-            OutboxMessage::Kill(fid) => {
-                ShuffleMessageDestination::PartitionProcessor(PartitionProcessorMessage::Kill(fid))
+            OutboxMessage::TerminateInvocation(invocation_termination) => {
+                ShuffleMessageDestination::PartitionProcessor(
+                    PartitionProcessorMessage::TerminateInvocation(invocation_termination),
+                )
             }
         }
     }