From 6b80ac8e0db8e5b2cc0e561d5c885de5b9b1df1a Mon Sep 17 00:00:00 2001
From: Rimuksh Kansal <nihonrk@gmail.com>
Date: Mon, 10 Aug 2026 16:25:47 +0900
Subject: [PATCH] add heartbeat handle and drop for rust sdk

---
 core/sdk/src/clients/client.rs | 35 +++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/core/sdk/src/clients/client.rs b/core/sdk/src/clients/client.rs
index 98beb98300..82595a0bf3 100644
--- a/core/sdk/src/clients/client.rs
+++ b/core/sdk/src/clients/client.rs
@@ -39,8 +39,9 @@ use iggy_common::locking::{IggyRwLock, IggyRwLockFn};
 use iggy_common::{BinaryTransport, Client, HttpMethod, SystemClient};
 use iggy_common::{ConnectionStringUtils, DiagnosticEvent, Partitioner, TransportProtocol};
 use std::fmt::Debug;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use tokio::spawn;
+use tokio::task::JoinHandle;
 use tokio::time::sleep;
 use tracing::log::warn;
 use tracing::{debug, error, info};
@@ -64,6 +65,7 @@ pub struct IggyClient {
     pub(crate) client: IggyRwLock<ClientWrapper>,
     partitioner: Option<Arc<dyn Partitioner>>,
     pub(crate) encryptor: Option<Arc<EncryptorKind>>,
+    heartbeat_handle: Mutex<Option<JoinHandle<()>>>,
 }
 
 impl Default for IggyClient {
@@ -92,6 +94,7 @@ impl IggyClient {
             client,
             partitioner: None,
             encryptor: None,
+            heartbeat_handle: Mutex::new(None),
         }
     }
 
@@ -131,6 +134,7 @@ impl IggyClient {
             client,
             partitioner,
             encryptor,
+            heartbeat_handle: Mutex::new(None),
         }
     }
 
@@ -237,6 +241,19 @@ impl IggyClient {
     }
 }
 
+impl Drop for IggyClient {
+    fn drop(&mut self) {
+        let heartbeat_handle = self
+            .heartbeat_handle
+            .get_mut()
+            .unwrap_or_else(|error| error.into_inner())
+            .take();
+        if let Some(handle) = heartbeat_handle {
+            handle.abort();
+        }
+    }
+}
+
 #[async_trait]
 impl Client for IggyClient {
     async fn connect(&self) -> Result<(), IggyError> {
@@ -247,8 +264,20 @@ impl Client for IggyClient {
             heartbeat_interval = client.heartbeat_interval().await;
         }
 
+        let mut heartbeat_handle = self
+            .heartbeat_handle
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        if heartbeat_handle
+            .as_ref()
+            .is_some_and(|handle| !handle.is_finished())
+        {
+            return Ok(());
+        }
+
+        drop(heartbeat_handle.take());
         let client = self.client.clone();
-        spawn(async move {
+        *heartbeat_handle = Some(spawn(async move {
             loop {
                 debug!("Sending the heartbeat...");
                 if let Err(error) = client.read().await.ping().await {
@@ -269,7 +298,7 @@ impl Client for IggyClient {
                 }
                 sleep(heartbeat_interval.get_duration()).await
             }
-        });
+        }));
         Ok(())
     }