From 61efc26ce8aeb61a858bf210b4b229c7f47f116c Mon Sep 17 00:00:00 2001
From: Ishan Bhanuka <bhanuka.ishan@gmail.com>
Date: Thu, 3 Oct 2024 21:26:24 +0530
Subject: [PATCH] Cleanup websocket tasks on disconnect if it is not alive

---
 nautilus_core/network/src/websocket.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/nautilus_core/network/src/websocket.rs b/nautilus_core/network/src/websocket.rs
index 79c15b81748..f8aecfc9054 100644
--- a/nautilus_core/network/src/websocket.rs
+++ b/nautilus_core/network/src/websocket.rs
@@ -444,7 +444,12 @@ impl WebSocketClient {
                         }
                         break;
                     }
-                    (true, false) => break,
+                    // Close the heartbeat task on disconnect if the connection is already closed
+                    (true, false) => {
+                        tracing::debug!("Inner client is disconnected");
+                        tracing::debug!("Shutting down inner client to clean up running tasks");
+                        inner.shutdown().await
+                    }
                     _ => (),
                 }
             }