Add some unit test coverage of journal table v2 feature, especially the problematic situations.

Fix some incorrectly handled corner cases, such as deletion of journal when the pinned deployment is not set yet.

Author: slinkydeveloper

crates/worker/src/partition/state_machine/entries/mod.rs

@@ -405,9 +405,11 @@ impl<CMD> ApplyJournalCommandEffect<'_, CMD> {
 
 #[cfg(test)]
 mod tests {
+    use crate::partition::state_machine::ExperimentalFeature;
     use crate::partition::state_machine::tests::fixtures::invoker_entry_effect;
     use crate::partition::state_machine::tests::{TestEnv, fixtures, matchers};
     use bytes::Bytes;
+    use enumset::EnumSet;
     use googletest::prelude::*;
     use restate_storage_api::invocation_status_table::ReadInvocationStatusTable;
     use restate_types::identifiers::{InvocationId, ServiceId};
@@ -416,10 +418,15 @@ mod tests {
     };
     use restate_types::journal_v2::{CallCommand, CallRequest};
     use restate_wal_protocol::Command;
+    use rstest::rstest;
 
+    #[rstest]
     #[restate_core::test]
-    async fn update_journal_and_commands_length() {
-        let mut test_env = TestEnv::create().await;
+    async fn update_journal_and_commands_length(
+        #[values(ExperimentalFeature::UseJournalTableV2AsDefault.into(), EnumSet::empty())]
+        features: EnumSet<ExperimentalFeature>,
+    ) {
+        let mut test_env = TestEnv::create_with_experimental_features(features).await;
         let invocation_id = fixtures::mock_start_invocation(&mut test_env).await;
         fixtures::mock_pinned_deployment_v5(&mut test_env, invocation_id).await;
 

crates/worker/src/partition/state_machine/entries/notification.rs

@@ -201,9 +201,11 @@ where
 mod tests {
     use super::*;
 
+    use crate::partition::state_machine::ExperimentalFeature;
     use crate::partition::state_machine::tests::{TestEnv, fixtures, matchers};
     use bytes::Bytes;
     use bytestring::ByteString;
+    use enumset::EnumSet;
     use googletest::prelude::*;
     use restate_service_protocol_v4::entry_codec::ServiceProtocolV4Codec;
     use restate_storage_api::invocation_status_table::{
@@ -262,9 +264,13 @@ mod tests {
         test_env.shutdown().await;
     }
 
+    #[rstest]
     #[restate_core::test]
-    async fn notify_signal_received_before_pinned_deployment() {
-        let mut test_env = TestEnv::create().await;
+    async fn notify_signal_received_before_pinned_deployment(
+        #[values(ExperimentalFeature::UseJournalTableV2AsDefault.into(), EnumSet::empty())]
+        features: EnumSet<ExperimentalFeature>,
+    ) {
+        let mut test_env = TestEnv::create_with_experimental_features(features).await;
         let invocation_id = fixtures::mock_start_invocation(&mut test_env).await;
 
         // Send signal notification before pinned deployment

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

@@ -117,16 +117,17 @@ where
 mod tests {
     use super::*;
 
-    use crate::partition::state_machine::Action;
     use crate::partition::state_machine::tests::{TestEnv, fixtures, matchers};
+    use crate::partition::state_machine::{Action, ExperimentalFeature};
     use crate::partition::types::InvokerEffectKind;
     use assert2::assert;
-    use googletest::pat;
-    use googletest::prelude::{assert_that, contains, eq, ge, not, some};
+    use enumset::EnumSet;
+    use googletest::prelude::*;
     use restate_invoker_api::Effect;
     use restate_storage_api::invocation_status_table::{
         InvocationStatus, ReadInvocationStatusTable,
     };
+    use restate_storage_api::journal_table_v2;
     use restate_storage_api::outbox_table::ReadOutboxTable;
     use restate_types::deployment::PinnedDeployment;
     use restate_types::errors::CANCELED_INVOCATION_ERROR;
@@ -136,10 +137,11 @@ mod tests {
         InvocationTarget, InvocationTermination, JournalCompletionTarget, NotifySignalRequest,
         ResponseResult, ServiceInvocation, ServiceInvocationResponseSink,
     };
-    use restate_types::journal_v2::CANCEL_SIGNAL;
+    use restate_types::journal_v2::{CANCEL_SIGNAL, CommandType, Entry, EntryMetadata, EntryType};
     use restate_types::service_protocol::ServiceProtocolVersion;
     use restate_types::time::MillisSinceEpoch;
     use restate_wal_protocol::Command;
+    use rstest::rstest;
 
     #[restate_core::test]
     async fn cancel_invoked_invocation() {
@@ -191,7 +193,8 @@ mod tests {
     }
 
     #[restate_core::test]
-    async fn cancel_invoked_invocation_without_pinned_deployment() {
+    async fn cancel_invoked_invocation_without_pinned_deployment_without_journal_table_v2_default()
+    {
         let mut test_env = TestEnv::create().await;
         let invocation_id = fixtures::mock_start_invocation(&mut test_env).await;
 
@@ -233,8 +236,46 @@ mod tests {
     }
 
     #[restate_core::test]
-    async fn cancel_scheduled_invocation_through_notify_signal() -> anyhow::Result<()> {
-        let mut test_env = TestEnv::create().await;
+    async fn cancel_invoked_invocation_without_pinned_deployment_with_journal_table_v2_default() {
+        let mut test_env = TestEnv::create_with_experimental_features(
+            ExperimentalFeature::UseJournalTableV2AsDefault.into(),
+        )
+        .await;
+        let invocation_id = fixtures::mock_start_invocation(&mut test_env).await;
+
+        // Send signal notification before pinning the deployment
+        let actions = test_env
+            .apply(Command::NotifySignal(NotifySignalRequest {
+                invocation_id,
+                signal: CANCEL_SIGNAL.try_into().unwrap(),
+            }))
+            .await;
+        assert_that!(
+            actions,
+            contains(matchers::actions::forward_notification(
+                invocation_id,
+                CANCEL_SIGNAL.clone()
+            ))
+        );
+
+        assert_that!(
+            test_env.read_journal_to_vec(invocation_id, 2).await,
+            elements_are![
+                property!(Entry.ty(), eq(EntryType::Command(CommandType::Input))),
+                matchers::entry_eq(CANCEL_SIGNAL),
+            ]
+        );
+
+        test_env.shutdown().await;
+    }
+
+    #[rstest]
+    #[restate_core::test]
+    async fn cancel_scheduled_invocation_through_notify_signal(
+        #[values(ExperimentalFeature::UseJournalTableV2AsDefault.into(), EnumSet::empty())]
+        features: EnumSet<ExperimentalFeature>,
+    ) -> anyhow::Result<()> {
+        let mut test_env = TestEnv::create_with_experimental_features(features).await;
 
         let invocation_id = InvocationId::mock_random();
         let rpc_id = PartitionProcessorRpcRequestId::new();
@@ -277,13 +318,39 @@ mod tests {
             .await?;
         assert!(let InvocationStatus::Free = current_invocation_status);
 
+        // Both journal table v1 and v2 are empty
+        assert!(
+            journal_table::ReadJournalTable::get_journal_entry(
+                &mut test_env.storage,
+                &invocation_id,
+                0
+            )
+            .await
+            .unwrap()
+            .is_none()
+        );
+        assert!(
+            journal_table_v2::ReadJournalTable::get_journal_entry(
+                &mut test_env.storage,
+                invocation_id,
+                0
+            )
+            .await
+            .unwrap()
+            .is_none()
+        );
+
         test_env.shutdown().await;
         Ok(())
     }
 
+    #[rstest]
     #[restate_core::test]
-    async fn cancel_inboxed_invocation_through_notify_signal() -> anyhow::Result<()> {
-        let mut test_env = TestEnv::create().await;
+    async fn cancel_inboxed_invocation_through_notify_signal(
+        #[values(ExperimentalFeature::UseJournalTableV2AsDefault.into(), EnumSet::empty())]
+        features: EnumSet<ExperimentalFeature>,
+    ) -> anyhow::Result<()> {
+        let mut test_env = TestEnv::create_with_experimental_features(features).await;
 
         let invocation_target = InvocationTarget::mock_virtual_object();
         let invocation_id = InvocationId::mock_generate(&invocation_target);
@@ -335,6 +402,28 @@ mod tests {
         // assert that invocation status was removed
         assert!(let InvocationStatus::Free = current_invocation_status);
 
+        // Both journal table v1 and v2 are empty
+        assert!(
+            journal_table::ReadJournalTable::get_journal_entry(
+                &mut test_env.storage,
+                &inboxed_id,
+                0
+            )
+            .await
+            .unwrap()
+            .is_none()
+        );
+        assert!(
+            journal_table_v2::ReadJournalTable::get_journal_entry(
+                &mut test_env.storage,
+                inboxed_id,
+                0
+            )
+            .await
+            .unwrap()
+            .is_none()
+        );
+
         assert_that!(
             actions,
             contains(

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

@@ -24,7 +24,6 @@ use restate_types::invocation::client::PurgeInvocationResponse;
 use restate_types::invocation::{
     InvocationMutationResponseSink, InvocationTargetType, WorkflowHandlerType,
 };
-use restate_types::service_protocol::ServiceProtocolVersion;
 use tracing::trace;
 
 pub struct OnPurgeCommand {
@@ -57,10 +56,9 @@ where
                 pinned_deployment,
                 ..
             }) => {
-                let should_remove_journal_table_v2 =
-                    pinned_deployment.as_ref().is_some_and(|pinned_deployment| {
-                        pinned_deployment.service_protocol_version >= ServiceProtocolVersion::V4
-                    });
+                let pinned_service_protocol_version = pinned_deployment
+                    .as_ref()
+                    .map(|pd| pd.service_protocol_version);
 
                 ctx.do_free_invocation(invocation_id)?;
 
@@ -92,7 +90,7 @@ where
                     ctx.do_drop_journal(
                         invocation_id,
                         journal_metadata.length,
-                        should_remove_journal_table_v2,
+                        pinned_service_protocol_version,
                     )
                     .await?;
                 }

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

@@ -18,7 +18,6 @@ use restate_storage_api::journal_table_v2::WriteJournalTable;
 use restate_types::identifiers::InvocationId;
 use restate_types::invocation::InvocationMutationResponseSink;
 use restate_types::invocation::client::PurgeInvocationResponse;
-use restate_types::service_protocol::ServiceProtocolVersion;
 use tracing::trace;
 
 pub struct OnPurgeJournalCommand {
@@ -42,19 +41,17 @@ where
         } = self;
         match ctx.get_invocation_status(&invocation_id).await? {
             InvocationStatus::Completed(mut completed) => {
-                let should_remove_journal_table_v2 = completed
+                let pinned_service_protocol_version = completed
                     .pinned_deployment
                     .as_ref()
-                    .is_some_and(|pinned_deployment| {
-                        pinned_deployment.service_protocol_version >= ServiceProtocolVersion::V4
-                    });
+                    .map(|pd| pd.service_protocol_version);
 
                 // If journal is not empty, clean it up
                 if completed.journal_metadata.length != 0 {
                     ctx.do_drop_journal(
                         invocation_id,
                         completed.journal_metadata.length,
-                        should_remove_journal_table_v2,
+                        pinned_service_protocol_version,
                     )
                     .await?;
                 }
@@ -113,6 +110,7 @@ mod tests {
         InvocationTarget, PurgeInvocationRequest, ServiceInvocation, ServiceInvocationResponseSink,
     };
     use restate_types::journal_v2::{CommandType, OutputCommand, OutputResult};
+    use restate_types::service_protocol::ServiceProtocolVersion;
     use restate_wal_protocol::Command;
     use std::time::Duration;
 

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

@@ -280,6 +280,7 @@ where
 mod tests {
     use super::*;
 
+    use crate::partition::state_machine::ExperimentalFeature;
     use crate::partition::state_machine::tests::{TestEnv, fixtures, matchers};
     use googletest::prelude::*;
     use restate_storage_api::invocation_status_table::{
@@ -291,8 +292,8 @@ mod tests {
     };
     use restate_types::invocation::client::RestartAsNewInvocationResponse;
     use restate_types::invocation::{
-        IngressInvocationResponseSink, InvocationTarget, NotifySignalRequest,
-        RestartAsNewInvocationRequest, ServiceInvocation,
+        IngressInvocationResponseSink, InvocationTarget, InvocationTermination,
+        NotifySignalRequest, RestartAsNewInvocationRequest, ServiceInvocation, TerminationFlavor,
     };
     use restate_types::journal_v2::{
         CommandType, CompletionType, NotificationType, OutputCommand, OutputResult, Signal,
@@ -377,6 +378,82 @@ mod tests {
         test_env.shutdown().await;
     }
 
+    #[restate_core::test]
+    async fn restart_killed_invocation() {
+        // This works only when using journal table v2 as default!
+        // The corner case with journal table v1 is handled by the rpc handler instead.
+        let mut test_env = TestEnv::create_with_experimental_features(
+            ExperimentalFeature::UseJournalTableV2AsDefault.into(),
+        )
+        .await;
+
+        // Start invocation, then kill it
+        let invocation_target = InvocationTarget::mock_virtual_object();
+        let original_invocation_id = InvocationId::generate(&invocation_target, None);
+        let _ = test_env
+            .apply_multiple([
+                Command::Invoke(Box::new(ServiceInvocation {
+                    invocation_id: original_invocation_id,
+                    invocation_target: invocation_target.clone(),
+                    completion_retention_duration: Duration::from_secs(120),
+                    journal_retention_duration: Duration::from_secs(120),
+                    ..ServiceInvocation::mock()
+                })),
+                Command::TerminateInvocation(InvocationTermination {
+                    invocation_id: original_invocation_id,
+                    flavor: TerminationFlavor::Kill,
+                    response_sink: None,
+                }),
+            ])
+            .await;
+
+        // Restart as new with copy_prefix_up_to_index_included = 0
+        let new_id = InvocationId::mock_generate(&invocation_target);
+        let request_id = PartitionProcessorRpcRequestId::new();
+        let actions = test_env
+            .apply(Command::RestartAsNewInvocation(
+                RestartAsNewInvocationRequest {
+                    invocation_id: original_invocation_id,
+                    new_invocation_id: new_id,
+                    copy_prefix_up_to_index_included: 0,
+                    patch_deployment_id: None,
+                    response_sink: Some(InvocationMutationResponseSink::Ingress(
+                        IngressInvocationResponseSink { request_id },
+                    )),
+                },
+            ))
+            .await;
+
+        // We should invoke the new invocation and send OK back
+        assert_that!(
+            actions,
+            all!(
+                contains(matchers::actions::invoke_for_id(new_id)),
+                contains(pat!(Action::ForwardRestartAsNewInvocationResponse {
+                    request_id: eq(request_id),
+                    response: eq(RestartAsNewInvocationResponse::Ok {
+                        new_invocation_id: new_id
+                    })
+                }))
+            )
+        );
+
+        assert_that!(
+            test_env
+                .storage
+                .get_invocation_status(&new_id)
+                .await
+                .unwrap(),
+            all!(
+                matchers::storage::is_variant(InvocationStatusDiscriminants::Invoked),
+                matchers::storage::has_journal_length(1),
+                matchers::storage::has_commands(1)
+            )
+        );
+
+        test_env.shutdown().await;
+    }
+
     #[restate_core::test]
     async fn restart_copy_input_only() {
         let mut test_env = TestEnv::create().await;