fix(dmq): consumer client can't connect more that once to server

jpraynaud · jpraynaud · commit c7c2c6384419 · 2025-08-28T17:58:27.000+02:00
diff --git a/internal/mithril-dmq/Cargo.toml b/internal/mithril-dmq/Cargo.toml
@@ -30,7 +30,7 @@ serde = { workspace = true }
 serde_bytes = "0.11.17"
 slog = { workspace = true }
 slog-scope = "4.4.0"
-tokio = { workspace = true, features = ["sync"] }
+tokio = { workspace = true, features = ["sync","rt-multi-thread"] }
 
 [dev-dependencies]
 mockall = { workspace = true }
diff --git a/internal/mithril-dmq/src/consumer/client/pallas.rs b/internal/mithril-dmq/src/consumer/client/pallas.rs
@@ -2,7 +2,7 @@ use std::{fmt::Debug, marker::PhantomData, path::PathBuf};
 
 use anyhow::{Context, anyhow};
 use pallas_network::{facades::DmqClient, miniprotocols::localmsgnotification::State};
-use slog::{Logger, debug};
+use slog::{Logger, debug, error};
 use tokio::sync::{Mutex, MutexGuard};
 
 use mithril_common::{
@@ -138,6 +138,18 @@ impl<M: TryFromBytes + Debug + Sync + Send> DmqConsumerClient<M> for DmqConsumer
     }
 }
 
+impl<M: TryFromBytes + Debug> Drop for DmqConsumerClientPallas<M> {
+    fn drop(&mut self) {
+        tokio::task::block_in_place(|| {
+            tokio::runtime::Handle::current().block_on(async {
+                if let Err(e) = self.drop_client().await {
+                    error!(self.logger, "Failed to drop DMQ consumer client: {}", e);
+                }
+            });
+        });
+    }
+}
+
 #[cfg(all(test, unix))]
 mod tests {
 
@@ -236,7 +248,7 @@ mod tests {
         })
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn pallas_dmq_consumer_publisher_succeeds_when_messages_are_available() {
         let socket_path = create_temp_dir(current_function!()).join("node.socket");
         let reply_messages = fake_msgs();
@@ -293,7 +305,7 @@ mod tests {
         result.expect_err("Should have timed out");
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn pallas_dmq_consumer_client_is_dropped_when_returning_error() {
         let socket_path = create_temp_dir(current_function!()).join("node.socket");
         let reply_messages = fake_msgs();
diff --git a/internal/mithril-dmq/src/consumer/server/pallas.rs b/internal/mithril-dmq/src/consumer/server/pallas.rs
@@ -91,8 +91,6 @@ impl DmqConsumerServerPallas {
     }
 
     /// Drops the current `DmqServer`, if it exists.
-    // TODO: remove allow dead code
-    #[allow(dead_code)]
     async fn drop_server(&self) -> StdResult<()> {
         debug!(
             self.logger,
@@ -179,9 +177,9 @@ impl DmqConsumerServerPallas {
                         }
                         Err(err) => {
                             error!(self.logger, "Failed to process message"; "error" => ?err);
-                            /* if let Err(drop_err) = self.drop_server().await {
+                            if let Err(drop_err) = self.drop_server().await {
                                 error!(self.logger, "Failed to drop DMQ consumer server"; "error" => ?drop_err);
-                            } */
+                            }
                         }
                     }
                 }
@@ -197,15 +195,9 @@ impl DmqConsumerServer for DmqConsumerServerPallas {
             self.logger,
             "Waiting for message received from the DMQ network"
         );
+
         let mut server_guard = self.get_server().await?;
         let server = server_guard.as_mut().ok_or(anyhow!("DMQ server does not exist"))?;
-
-        debug!(
-            self.logger,
-            "DMQ Server state: {:?}",
-            server.msg_notification().state()
-        );
-
         let request = server
             .msg_notification()
             .recv_next_request()
@@ -253,7 +245,6 @@ impl DmqConsumerServer for DmqConsumerServerPallas {
         Ok(())
     }
 
-    /// Runs the DMQ publisher server, processing messages in a loop.
     async fn run(&self) -> StdResult<()> {
         info!(
             self.logger,
@@ -273,6 +264,18 @@ impl DmqConsumerServer for DmqConsumerServerPallas {
     }
 }
 
+impl Drop for DmqConsumerServerPallas {
+    fn drop(&mut self) {
+        tokio::task::block_in_place(|| {
+            tokio::runtime::Handle::current().block_on(async {
+                if let Err(e) = self.drop_server().await {
+                    error!(self.logger, "Failed to drop DMQ consumer server: {}", e);
+                }
+            });
+        });
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::{sync::Arc, time::Duration};
@@ -306,7 +309,7 @@ mod tests {
         }
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn pallas_dmq_consumer_server_non_blocking_success() {
         let (stop_tx, stop_rx) = watch::channel(());
         let (signature_dmq_tx, signature_dmq_rx) = unbounded_channel::<DmqMessage>();
@@ -322,7 +325,8 @@ mod tests {
         let message = fake_msg();
         let client = tokio::spawn({
             async move {
-                tokio::time::sleep(Duration::from_secs(1)).await;
+                // sleep to avoid refused connection from the server
+                tokio::time::sleep(Duration::from_millis(10)).await;
 
                 // client setup
                 let mut client = DmqClient::connect(socket_path.clone(), 0).await.unwrap();
@@ -364,7 +368,7 @@ mod tests {
         assert_eq!(vec![message], messages_received);
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn pallas_dmq_consumer_server_blocking_success() {
         let (stop_tx, stop_rx) = watch::channel(());
         let (signature_dmq_tx, signature_dmq_rx) = unbounded_channel::<DmqMessage>();
@@ -380,7 +384,8 @@ mod tests {
         let message = fake_msg();
         let client = tokio::spawn({
             async move {
-                tokio::time::sleep(Duration::from_secs(1)).await;
+                // sleep to avoid refused connection from the server
+                tokio::time::sleep(Duration::from_millis(10)).await;
 
                 // client setup
                 let mut client = DmqClient::connect(socket_path.clone(), 0).await.unwrap();
@@ -422,10 +427,10 @@ mod tests {
         assert_eq!(vec![message], messages_received);
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn pallas_dmq_consumer_server_blocking_blocks_when_no_message_available() {
         let (_stop_tx, stop_rx) = watch::channel(());
-        let (signature_dmq_tx, signature_dmq_rx) = unbounded_channel::<DmqMessage>();
+        let (_signature_dmq_tx, signature_dmq_rx) = unbounded_channel::<DmqMessage>();
         let socket_path = create_temp_dir(current_function!()).join("node.socket");
         let cardano_network = CardanoNetwork::TestNet(0);
         let dmq_consumer_server = Arc::new(DmqConsumerServerPallas::new(
@@ -437,6 +442,9 @@ mod tests {
         dmq_consumer_server.register_receiver(signature_dmq_rx).await.unwrap();
         let client = tokio::spawn({
             async move {
+                // sleep to avoid refused connection from the server
+                tokio::time::sleep(Duration::from_millis(10)).await;
+
                 // client setup
                 let mut client = DmqClient::connect(socket_path.clone(), 0).await.unwrap();
 
@@ -451,10 +459,9 @@ mod tests {
                     localmsgnotification::State::BusyBlocking
                 );
 
-                client_msg.recv_next_reply().await.unwrap();
+                let _ = client_msg.recv_next_reply().await;
             }
         });
-        let _signature_dmq_tx_clone = signature_dmq_tx.clone();
 
         let result = tokio::select!(
             _res = sleep(Duration::from_millis(1000)) => {Err(anyhow!("Timeout"))},
diff --git a/internal/mithril-dmq/tests/consumer_client_server.rs b/internal/mithril-dmq/tests/consumer_client_server.rs
@@ -19,10 +19,8 @@ async fn create_fake_msg(bytes: &[u8]) -> DmqMessage {
     let dmq_builder = DmqMessageBuilder::new(
         {
             let (kes_signature, operational_certificate) = KesSignerFake::dummy_signature();
-            let kes_signer = KesSignerFake::new(vec![
-                Ok((kes_signature, operational_certificate.clone())),
-                Ok((kes_signature, operational_certificate.clone())), // TODO: remove this line once the hack of KES signature is removed in DMQ message builder
-            ]);
+            let kes_signer =
+                KesSignerFake::new(vec![Ok((kes_signature, operational_certificate.clone()))]);
 
             Arc::new(kes_signer)
         },
@@ -33,7 +31,7 @@ async fn create_fake_msg(bytes: &[u8]) -> DmqMessage {
     dmq_builder.build(&message.to_bytes_vec().unwrap()).await.unwrap()
 }
 
-#[tokio::test]
+#[tokio::test(flavor = "multi_thread")]
 async fn dmq_consumer_client_server() {
     let cardano_network = CardanoNetwork::TestNet(0);
     let socket_path =
@@ -89,7 +87,7 @@ async fn dmq_consumer_client_server() {
     );
 
     // Sleep to avoid refused connection from the server
-    tokio::time::sleep(std::time::Duration::from_millis(3000)).await;
+    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
 
     // Start a second client, receive messages
     let client = tokio::spawn({
