From d8f9e49b56a6436603e69a9d06ce3782adb626dc Mon Sep 17 00:00:00 2001
From: Ahmed Farghal <me@asoli.dev>
Date: Tue, 10 Sep 2024 08:51:53 +0100
Subject: [PATCH] Fix state_machine teardown in tests

Correctly shutdown rocksdb manager in those tests. This fixes the occasional segfaults in CI tests
---
 .../src/partition/state_machine/tests/delayed_send.rs  |  3 +++
 .../src/partition/state_machine/tests/idempotency.rs   | 10 ++++++++++
 .../src/partition/state_machine/tests/kill_cancel.rs   |  4 ++++
 .../src/partition/state_machine/tests/workflow.rs      |  4 ++++
 4 files changed, 21 insertions(+)

diff --git a/crates/worker/src/partition/state_machine/tests/delayed_send.rs b/crates/worker/src/partition/state_machine/tests/delayed_send.rs
index 1cc61eafff..18fcd25ab1 100644
--- a/crates/worker/src/partition/state_machine/tests/delayed_send.rs
+++ b/crates/worker/src/partition/state_machine/tests/delayed_send.rs
@@ -87,6 +87,7 @@ async fn send_with_delay() {
         test_env.storage.get_invocation_status(&invocation_id).await,
         ok(pat!(InvocationStatus::Invoked { .. }))
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -180,6 +181,7 @@ async fn send_with_delay_to_locked_virtual_object() {
             &invocation_target
         )))
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -280,4 +282,5 @@ async fn send_with_delay_and_idempotency_key() {
             ))))
         )
     );
+    test_env.shutdown().await;
 }
diff --git a/crates/worker/src/partition/state_machine/tests/idempotency.rs b/crates/worker/src/partition/state_machine/tests/idempotency.rs
index b758079e5a..76cc53b53c 100644
--- a/crates/worker/src/partition/state_machine/tests/idempotency.rs
+++ b/crates/worker/src/partition/state_machine/tests/idempotency.rs
@@ -134,6 +134,7 @@ async fn start_and_complete_idempotent_invocation() {
             response_result: eq(ResponseResult::Success(response_bytes))
         })))
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -246,6 +247,7 @@ async fn start_and_complete_idempotent_invocation_neo_table() {
         completed_invocation.completion_retention_duration,
         retention
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -326,6 +328,7 @@ async fn complete_already_completed_invocation() {
             .unwrap(),
         pat!(InvocationStatus::Free)
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -391,6 +394,7 @@ async fn known_invocation_id_but_missing_completion() {
             .unwrap(),
         pat!(InvocationStatus::Free)
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -502,6 +506,7 @@ async fn attach_with_service_invocation_command_while_executing() {
             ))))
         )
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -625,6 +630,7 @@ async fn attach_with_send_service_invocation() {
             ))))),
         )
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -737,6 +743,7 @@ async fn attach_inboxed_with_send_service_invocation() {
             ))))
         )
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -844,6 +851,7 @@ async fn attach_command() {
             ))))
         )
     );
+    test_env.shutdown().await;
 }
 
 // TODO remove this once we remove the old invocation status table
@@ -910,6 +918,7 @@ async fn timer_cleanup() {
             .unwrap(),
         none()
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -962,4 +971,5 @@ async fn purge_completed_idempotent_invocation() {
             .unwrap(),
         none()
     );
+    test_env.shutdown().await;
 }
diff --git a/crates/worker/src/partition/state_machine/tests/kill_cancel.rs b/crates/worker/src/partition/state_machine/tests/kill_cancel.rs
index 981db45750..8ac633060c 100644
--- a/crates/worker/src/partition/state_machine/tests/kill_cancel.rs
+++ b/crates/worker/src/partition/state_machine/tests/kill_cancel.rs
@@ -101,6 +101,7 @@ async fn kill_inboxed_invocation() -> anyhow::Result<()> {
         some((ge(0), outbox_message_matcher(caller_id)))
     );
 
+    test_env.shutdown().await;
     Ok(())
 }
 
@@ -210,6 +211,7 @@ async fn kill_call_tree() -> anyhow::Result<()> {
         )
     );
 
+    test_env.shutdown().await;
     Ok(())
 }
 
@@ -322,6 +324,7 @@ async fn cancel_invoked_invocation() -> Result<(), Error> {
         )
     );
 
+    test_env.shutdown().await;
     Ok(())
 }
 
@@ -437,6 +440,7 @@ async fn cancel_suspended_invocation() -> Result<(), Error> {
             }))
         )
     );
+    test_env.shutdown().await;
 
     Ok(())
 }
diff --git a/crates/worker/src/partition/state_machine/tests/workflow.rs b/crates/worker/src/partition/state_machine/tests/workflow.rs
index 268474867f..b5df04b6a9 100644
--- a/crates/worker/src/partition/state_machine/tests/workflow.rs
+++ b/crates/worker/src/partition/state_machine/tests/workflow.rs
@@ -197,6 +197,7 @@ async fn start_workflow_method() {
             }
         ))))
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -348,6 +349,7 @@ async fn attach_by_workflow_key() {
             }
         ))))
     );
+    test_env.shutdown().await;
 }
 
 // TODO remove this once we remove the old invocation status table
@@ -411,6 +413,7 @@ async fn timer_cleanup() {
             .unwrap(),
         pat!(VirtualObjectStatus::Unlocked)
     );
+    test_env.shutdown().await;
 }
 
 #[test(tokio::test)]
@@ -460,4 +463,5 @@ async fn purge_completed_workflow() {
             .unwrap(),
         pat!(VirtualObjectStatus::Unlocked)
     );
+    test_env.shutdown().await;
 }