From bad438b5ad43778a8231c2d52ce09391d8d36fd7 Mon Sep 17 00:00:00 2001
From: Andy Duplain <trojanfoe@gmail.com>
Date: Wed, 20 May 2026 12:24:20 +0100
Subject: [PATCH] Address review comments.

---
 Cargo.lock                                 | 1 +
 common/client-core/surb-storage/src/lib.rs | 7 ++++---
 nym-sqlx-pool-guard/Cargo.toml             | 1 +
 nym-sqlx-pool-guard/src/lib.rs             | 6 +++++-
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0eadd13981..dbdc7f892d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8353,6 +8353,7 @@ dependencies = [
  "tokio",
  "tracing",
  "tracing-subscriber",
+ "tracing-test",
  "windows 0.61.3",
 ]
 
diff --git a/common/client-core/surb-storage/src/lib.rs b/common/client-core/surb-storage/src/lib.rs
index f88a746ad2..f79127b422 100644
--- a/common/client-core/surb-storage/src/lib.rs
+++ b/common/client-core/surb-storage/src/lib.rs
@@ -56,10 +56,11 @@ where
 
         info!("PersistentReplyStorage is flushing all reply-related data to underlying storage");
         if let Err(err) = self.backend.flush_surb_storage(&mem_state).await {
-            error!("failed to flush our reply-related data to the persistent storage: {err}")
-        } else {
-            info!("Data flush is complete")
+            error!("failed to flush our reply-related data to the persistent storage: {err}");
+            self.backend.stop_storage_session().await.ok();
+            return;
         }
+        info!("Data flush is complete");
 
         if let Err(err) = self.backend.stop_storage_session().await {
             error!("failed to properly stop the storage session - {err}. We might not be able to smoothly restore it")
diff --git a/nym-sqlx-pool-guard/Cargo.toml b/nym-sqlx-pool-guard/Cargo.toml
index 85622863ad..e4b86ad2c5 100644
--- a/nym-sqlx-pool-guard/Cargo.toml
+++ b/nym-sqlx-pool-guard/Cargo.toml
@@ -42,3 +42,4 @@ windows = { version = "0.61", features = [
 [dev-dependencies]
 tempfile.workspace = true
 tracing-subscriber.workspace = true
+tracing-test.workspace = true
diff --git a/nym-sqlx-pool-guard/src/lib.rs b/nym-sqlx-pool-guard/src/lib.rs
index d1b3808bd9..8ded2cfb61 100644
--- a/nym-sqlx-pool-guard/src/lib.rs
+++ b/nym-sqlx-pool-guard/src/lib.rs
@@ -180,6 +180,7 @@ mod tests {
         ConnectOptions, Executor,
         sqlite::{SqliteAutoVacuum, SqliteSynchronous},
     };
+    use tracing_test::traced_test;
 
     use super::*;
 
@@ -219,6 +220,7 @@ mod tests {
         tokio::fs::remove_file(database_path).await.unwrap();
     }
 
+    #[traced_test]
     #[tokio::test]
     async fn test_clone_drop_no_warning() {
         // Cloning the guard and dropping the clone should not warn because the original is still alive.
@@ -238,9 +240,11 @@ mod tests {
         {
             let _clone = guard.clone();
             assert_eq!(Arc::strong_count(&guard.inner), 2);
-            // _clone drops here - should NOT warn since guard still holds a reference
         }
         assert_eq!(Arc::strong_count(&guard.inner), 1);
+        assert!(!logs_contain(
+            "SqlitePoolGuard dropped without explicit close"
+        ));
 
         guard.close().await;
         tokio::fs::remove_file(database_path).await.unwrap();