RUST-1608 Clean shutdown for Client (#920)

Author: abr-egn

.evergreen/check-cargo-deny.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 
 set -o errexit
+set -o xtrace
 
 source ./.evergreen/env.sh
 

Cargo.toml

@@ -180,6 +180,7 @@ features = ["v4"]
 anyhow = { version = "1.0", features = ["backtrace"] }
 approx = "0.5.1"
 async_once = "0.2.6"
+backtrace = { version = "0.3.68" }
 ctrlc = "3.2.2"
 function_name = "0.2.1"
 futures = "0.3"

src/client.rs

@@ -6,13 +6,18 @@ pub mod options;
 pub mod session;
 
 use std::{
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Mutex as SyncMutex,
+    },
     time::{Duration, Instant},
 };
 
 #[cfg(feature = "in-use-encryption-unstable")]
 pub use self::csfle::client_builder::*;
 use derivative::Derivative;
+use futures_core::{future::BoxFuture, Future};
+use futures_util::{future::join_all, FutureExt};
 
 #[cfg(test)]
 use crate::options::ServerAddress;
@@ -36,6 +41,7 @@ use crate::{
     db::Database,
     error::{Error, ErrorKind, Result},
     event::command::{handle_command_event, CommandEvent},
+    id_set::IdSet,
     operation::{AggregateTarget, ListDatabases},
     options::{
         ClientOptions,
@@ -47,6 +53,7 @@ use crate::{
     },
     results::DatabaseSpecification,
     sdam::{server_selection, SelectedServer, Topology},
+    tracking_arc::TrackingArc,
     ClientSession,
 };
 
@@ -103,9 +110,25 @@ const DEFAULT_SERVER_SELECTION_TIMEOUT: Duration = Duration::from_secs(30);
 /// driver does not set ``tcp_keepalive_intvl``. See the
 /// [MongoDB Diagnostics FAQ keepalive section](https://www.mongodb.com/docs/manual/faq/diagnostics/#does-tcp-keepalive-time-affect-mongodb-deployments)
 /// for instructions on setting these values at the system level.
-#[derive(Clone, Debug)]
+///
+/// ## Clean shutdown
+/// Because Rust has no async equivalent of `Drop`, values that require server-side cleanup when
+/// dropped spawn a new async task to perform that cleanup.  This can cause two potential issues:
+///
+/// * Drop tasks pending or in progress when the async runtime shuts down may not complete, causing
+///   server-side resources to not be freed.
+/// * Drop tasks may run at an arbitrary time even after no `Client` values exist, making it hard to
+///   reason about associated resources (e.g. event handlers).
+///
+/// To address these issues, we highly recommend you use [`Client::shutdown`] in the termination
+/// path of your application.  This will ensure that outstanding resources have been cleaned up and
+/// terminate internal worker tasks before returning.  Please note that `shutdown` will wait for
+/// _all_ outstanding resource handles to be dropped, so they must either have been dropped before
+/// calling `shutdown` or in a concurrent task; see the documentation of `shutdown` for more
+/// details.
+#[derive(Debug, Clone)]
 pub struct Client {
-    inner: Arc<ClientInner>,
+    inner: TrackingArc<ClientInner>,
 }
 
 #[allow(dead_code, unreachable_code, clippy::diverging_sub_expression)]
@@ -124,10 +147,17 @@ struct ClientInner {
     topology: Topology,
     options: ClientOptions,
     session_pool: ServerSessionPool,
+    shutdown: Shutdown,
     #[cfg(feature = "in-use-encryption-unstable")]
     csfle: tokio::sync::RwLock<Option<csfle::ClientState>>,
 }
 
+#[derive(Debug)]
+struct Shutdown {
+    pending_drops: SyncMutex<IdSet<crate::runtime::AsyncJoinHandle<()>>>,
+    executed: AtomicBool,
+}
+
 impl Client {
     /// Creates a new `Client` connected to the cluster specified by `uri`. `uri` must be a valid
     /// MongoDB connection string.
@@ -144,12 +174,16 @@ impl Client {
     pub fn with_options(options: ClientOptions) -> Result<Self> {
         options.validate()?;
 
-        let inner = Arc::new(ClientInner {
+        let inner = TrackingArc::new(ClientInner {
             topology: Topology::new(options.clone())?,
             session_pool: ServerSessionPool::new(),
+            options,
+            shutdown: Shutdown {
+                pending_drops: SyncMutex::new(IdSet::new()),
+                executed: AtomicBool::new(false),
+            },
             #[cfg(feature = "in-use-encryption-unstable")]
             csfle: Default::default(),
-            options,
         });
         Ok(Self { inner })
     }
@@ -461,6 +495,123 @@ impl Client {
             .await
     }
 
+    pub(crate) fn register_async_drop(&self) -> AsyncDropToken {
+        let (cleanup_tx, cleanup_rx) = tokio::sync::oneshot::channel::<BoxFuture<'static, ()>>();
+        let (id_tx, id_rx) = tokio::sync::oneshot::channel::<crate::id_set::Id>();
+        let weak = self.weak();
+        let handle = crate::runtime::spawn(async move {
+            // Unwrap safety: the id is sent immediately after task creation, with no
+            // await points in between.
+            let id = id_rx.await.unwrap();
+            // If the cleanup channel is closed, that task was dropped.
+            let cleanup = if let Ok(f) = cleanup_rx.await {
+                f
+            } else {
+                return;
+            };
+            cleanup.await;
+            if let Some(client) = weak.upgrade() {
+                client
+                    .inner
+                    .shutdown
+                    .pending_drops
+                    .lock()
+                    .unwrap()
+                    .remove(&id);
+            }
+        });
+        let id = self
+            .inner
+            .shutdown
+            .pending_drops
+            .lock()
+            .unwrap()
+            .insert(handle);
+        let _ = id_tx.send(id);
+        AsyncDropToken {
+            tx: Some(cleanup_tx),
+        }
+    }
+
+    /// Shut down this `Client`, terminating background thread workers and closing connections.
+    /// This will wait for any live handles to server-side resources (see below) to be
+    /// dropped and any associated server-side operations to finish.
+    ///
+    /// IMPORTANT: Any live resource handles that are not dropped will cause this method to wait
+    /// indefinitely.  It's strongly recommended to structure your usage to avoid this, e.g. by
+    /// only using those types in shorter-lived scopes than the `Client`.  If this is not possible,
+    /// see [`shutdown_immediate`](Client::shutdown_immediate).  For example:
+    ///
+    /// ```rust
+    /// # use mongodb::{Client, GridFsBucket, error::Result};
+    /// async fn upload_data(bucket: &GridFsBucket) {
+    ///   let stream = bucket.open_upload_stream("test", None);
+    ///    // .. write to the stream ..
+    /// }
+    ///
+    /// # async fn run() -> Result<()> {
+    /// let client = Client::with_uri_str("mongodb://example.com").await?;
+    /// let bucket = client.database("test").gridfs_bucket(None);
+    /// upload_data(&bucket).await;
+    /// client.shutdown().await;
+    /// // Background cleanup work from `upload_data` is guaranteed to have run.
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    /// If the handle is used in the same scope as `shutdown`, explicit `drop` may be needed:
+    ///
+    /// ```rust
+    /// # use mongodb::{Client, error::Result};
+    /// # async fn run() -> Result<()> {
+    /// let client = Client::with_uri_str("mongodb://example.com").await?;
+    /// let bucket = client.database("test").gridfs_bucket(None);
+    /// let stream = bucket.open_upload_stream("test", None);
+    /// // .. write to the stream ..
+    /// drop(stream);
+    /// client.shutdown().await;
+    /// // Background cleanup work for `stream` is guaranteed to have run.
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    /// Calling any methods on  clones of this `Client` or derived handles after this will return
+    /// errors.
+    ///
+    /// Handles to server-side resources are `Cursor`, `SessionCursor`, `Session`, or
+    /// `GridFsUploadStream`.
+    pub async fn shutdown(self) {
+        // Subtle bug: if this is inlined into the `join_all(..)` call, Rust will extend the
+        // lifetime of the temporary unnamed `MutexLock` until the end of the *statement*,
+        // causing the lock to be held for the duration of the join, which deadlocks.
+        let pending = self.inner.shutdown.pending_drops.lock().unwrap().extract();
+        join_all(pending).await;
+        self.shutdown_immediate().await;
+    }
+
+    /// Shut down this `Client`, terminating background thread workers and closing connections.
+    /// This does *not* wait for other pending resources to be cleaned up, which may cause both
+    /// client-side errors and server-side resource leaks. Calling any methods on clones of this
+    /// `Client` or derived handles after this will return errors.
+    ///
+    /// ```rust
+    /// # use mongodb::{Client, error::Result};
+    /// # async fn run() -> Result<()> {
+    /// let client = Client::with_uri_str("mongodb://example.com").await?;
+    /// let bucket = client.database("test").gridfs_bucket(None);
+    /// let stream = bucket.open_upload_stream("test", None);
+    /// // .. write to the stream ..
+    /// client.shutdown_immediate().await;
+    /// // Background cleanup work for `stream` may or may not have run.
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn shutdown_immediate(self) {
+        self.inner.topology.shutdown().await;
+        // This has to happen last to allow pending cleanup to execute commands.
+        self.inner.shutdown.executed.store(true, Ordering::SeqCst);
+    }
+
     /// Check in a server session to the server session pool. The session will be discarded if it is
     /// expired or dirty.
     pub(crate) async fn check_in_server_session(&self, session: ServerSession) {
@@ -630,10 +781,9 @@ impl Client {
         }
     }
 
-    #[cfg(feature = "in-use-encryption-unstable")]
     pub(crate) fn weak(&self) -> WeakClient {
         WeakClient {
-            inner: Arc::downgrade(&self.inner),
+            inner: TrackingArc::downgrade(&self.inner),
         }
     }
 
@@ -653,16 +803,35 @@ impl Client {
     }
 }
 
-#[cfg(feature = "in-use-encryption-unstable")]
 #[derive(Clone, Debug)]
 pub(crate) struct WeakClient {
-    inner: std::sync::Weak<ClientInner>,
+    inner: crate::tracking_arc::Weak<ClientInner>,
 }
 
-#[cfg(feature = "in-use-encryption-unstable")]
 impl WeakClient {
-    #[allow(dead_code)]
     pub(crate) fn upgrade(&self) -> Option<Client> {
         self.inner.upgrade().map(|inner| Client { inner })
     }
 }
+
+#[derive(Derivative)]
+#[derivative(Debug)]
+pub(crate) struct AsyncDropToken {
+    #[derivative(Debug = "ignore")]
+    tx: Option<tokio::sync::oneshot::Sender<BoxFuture<'static, ()>>>,
+}
+
+impl AsyncDropToken {
+    pub(crate) fn spawn(&mut self, fut: impl Future<Output = ()> + Send + 'static) {
+        if let Some(tx) = self.tx.take() {
+            let _ = tx.send(fut.boxed());
+        } else {
+            #[cfg(debug_assertions)]
+            panic!("exhausted AsyncDropToken");
+        }
+    }
+
+    pub(crate) fn take(&mut self) -> Self {
+        Self { tx: self.tx.take() }
+    }
+}

src/client/executor.rs

@@ -6,7 +6,11 @@ use futures_core::future::BoxFuture;
 use lazy_static::lazy_static;
 use serde::de::DeserializeOwned;
 
-use std::{collections::HashSet, sync::Arc, time::Instant};
+use std::{
+    collections::HashSet,
+    sync::{atomic::Ordering, Arc},
+    time::Instant,
+};
 
 use super::{session::TransactionState, Client, ClientSession};
 use crate::{
@@ -53,6 +57,7 @@ use crate::{
     options::{ChangeStreamOptions, SelectionCriteria},
     sdam::{HandshakePhase, SelectedServer, ServerType, TopologyType, TransactionSupportStatus},
     selection_criteria::ReadPreference,
+    tracking_arc::TrackingArc,
     ClusterTime,
 };
 
@@ -99,6 +104,9 @@ impl Client {
         op: T,
         session: impl Into<Option<&mut ClientSession>>,
     ) -> Result<ExecutionDetails<T>> {
+        if self.inner.shutdown.executed.load(Ordering::SeqCst) {
+            return Err(ErrorKind::Shutdown.into());
+        }
         Box::pin(async {
             // TODO RUST-9: allow unacknowledged write concerns
             if !op.is_acknowledged() {
@@ -109,7 +117,7 @@ impl Client {
             }
             let session = session.into();
             if let Some(session) = &session {
-                if !Arc::ptr_eq(&self.inner, &session.client().inner) {
+                if !TrackingArc::ptr_eq(&self.inner, &session.client().inner) {
                     return Err(ErrorKind::InvalidArgument {
                         message: "the session provided to an operation must be created from the \
                                   same client as the collection/database"

src/client/session.rs

@@ -18,15 +18,14 @@ use crate::{
     error::{ErrorKind, Result},
     operation::{AbortTransaction, CommitTransaction, Operation},
     options::{SessionOptions, TransactionOptions},
-    runtime,
     sdam::{ServerInfo, TransactionSupportStatus},
     selection_criteria::SelectionCriteria,
     Client,
 };
 pub use cluster_time::ClusterTime;
 pub(super) use pool::ServerSessionPool;
 
-use super::options::ServerAddress;
+use super::{options::ServerAddress, AsyncDropToken};
 
 lazy_static! {
     pub(crate) static ref SESSIONS_UNSUPPORTED_COMMANDS: HashSet<&'static str> = {
@@ -107,6 +106,7 @@ pub struct ClientSession {
     client: Client,
     is_implicit: bool,
     options: Option<SessionOptions>,
+    drop_token: AsyncDropToken,
     pub(crate) transaction: Transaction,
     pub(crate) snapshot_time: Option<Timestamp>,
     pub(crate) operation_time: Option<Timestamp>,
@@ -212,6 +212,7 @@ impl ClientSession {
         let timeout = client.inner.topology.logical_session_timeout();
         let server_session = client.inner.session_pool.check_out(timeout).await;
         Self {
+            drop_token: client.register_async_drop(),
             client,
             server_session,
             cluster_time: None,
@@ -694,6 +695,7 @@ impl From<DroppedClientSession> for ClientSession {
         Self {
             cluster_time: dropped_session.cluster_time,
             server_session: dropped_session.server_session,
+            drop_token: dropped_session.client.register_async_drop(),
             client: dropped_session.client,
             is_implicit: dropped_session.is_implicit,
             options: dropped_session.options,
@@ -719,14 +721,14 @@ impl Drop for ClientSession {
                 snapshot_time: self.snapshot_time,
                 operation_time: self.operation_time,
             };
-            runtime::execute(async move {
+            self.drop_token.spawn(async move {
                 let mut session: ClientSession = dropped_session.into();
                 let _result = session.abort_transaction().await;
             });
         } else {
             let client = self.client.clone();
             let server_session = self.server_session.clone();
-            runtime::execute(async move {
+            self.drop_token.spawn(async move {
                 client.check_in_server_session(server_session).await;
             });
         }