From fa5eea2132c44a7c1b1711d2646598b9091c9bfa Mon Sep 17 00:00:00 2001
From: Abraham Egnor <abraham.egnor@mongodb.com>
Date: Fri, 14 Feb 2025 13:15:45 -0500
Subject: [PATCH] RUST-2155 Precreate cleanup task for endSessions (#1305)

---
 src/client.rs | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index f99d6606c..c92ad72a5 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -131,6 +131,7 @@ struct ClientInner {
     session_pool: ServerSessionPool,
     shutdown: Shutdown,
     dropped: AtomicBool,
+    end_sessions_token: std::sync::Mutex<AsyncDropToken>,
     #[cfg(feature = "in-use-encryption")]
     csfle: tokio::sync::RwLock<Option<csfle::ClientState>>,
     #[cfg(test)]
@@ -159,6 +160,18 @@ impl Client {
     pub fn with_options(options: ClientOptions) -> Result<Self> {
         options.validate()?;
 
+        // Spawn a cleanup task, similar to register_async_drop
+        let (cleanup_tx, cleanup_rx) = tokio::sync::oneshot::channel::<BoxFuture<'static, ()>>();
+        crate::runtime::spawn(async move {
+            // If the cleanup channel is closed, that task was dropped.
+            if let Ok(cleanup) = cleanup_rx.await {
+                cleanup.await;
+            }
+        });
+        let end_sessions_token = std::sync::Mutex::new(AsyncDropToken {
+            tx: Some(cleanup_tx),
+        });
+
         let inner = TrackingArc::new(ClientInner {
             topology: Topology::new(options.clone())?,
             session_pool: ServerSessionPool::new(),
@@ -168,6 +181,7 @@ impl Client {
                 executed: AtomicBool::new(false),
             },
             dropped: AtomicBool::new(false),
+            end_sessions_token,
             #[cfg(feature = "in-use-encryption")]
             csfle: Default::default(),
             #[cfg(test)]
@@ -682,9 +696,13 @@ impl Drop for Client {
             // this cycle.
             self.inner.dropped.store(true, Ordering::SeqCst);
             let client = self.clone();
-            crate::runtime::spawn(async move {
-                client.end_all_sessions().await;
-            });
+            self.inner
+                .end_sessions_token
+                .lock()
+                .unwrap()
+                .spawn(async move {
+                    client.end_all_sessions().await;
+                });
         }
     }
 }