RUST-384 Close connection which was dropped during command execution (#214)

Author: Lesiuk

src/cmap/conn/mod.rs

@@ -61,6 +61,11 @@ pub(crate) struct Connection {
     /// currently checked into the pool, this will be None.
     pub(super) pool: Option<Weak<ConnectionPoolInner>>,
 
+    /// Whether or not a command is currently being run on this connection. This is set to `true`
+    /// right before sending bytes to the server and set back to `false` once a full response has
+    /// been read.
+    command_executing: bool,
+
     stream: AsyncStream,
 
     #[derivative(Debug = "ignore")]
@@ -85,6 +90,7 @@ impl Connection {
             id,
             generation,
             pool: None,
+            command_executing: false,
             ready_and_available_time: None,
             stream: AsyncStream::connect(stream_options).await?,
             address,
@@ -207,9 +213,13 @@ impl Connection {
         request_id: impl Into<Option<i32>>,
     ) -> Result<CommandResponse> {
         let message = Message::with_command(command, request_id.into());
+
+        self.command_executing = true;
         message.write_to(&mut self.stream).await?;
 
         let response_message = Message::read_from(&mut self.stream).await?;
+        self.command_executing = false;
+
         CommandResponse::new(self.address.clone(), response_message)
     }
 
@@ -253,24 +263,30 @@ impl Connection {
 
 impl Drop for Connection {
     fn drop(&mut self) {
-        // If the connection has a weak reference to a pool, that means that the connection is being
-        // dropped when it's checked out. If the pool is still alive, it should check itself back
-        // in. Otherwise, the connection should close itself and emit a ConnectionClosed event
-        // (because the `close_and_drop` helper was not called explicitly).
-        //
-        // If the connection does not have a weak reference to a pool, then the connection is being
-        // dropped while it's not checked out. This means that the pool called the `close_and_drop`
-        // helper explicitly, so we don't add it back to the pool or emit any events.
-        if let Some(ref weak_pool_ref) = self.pool {
-            if let Some(strong_pool_ref) = weak_pool_ref.upgrade() {
-                let dropped_connection_state = self.take();
-                RUNTIME.execute(async move {
-                    strong_pool_ref
-                        .check_in(dropped_connection_state.into())
-                        .await;
-                });
-            } else {
-                self.close(ConnectionClosedReason::PoolClosed);
+        if self.command_executing {
+            self.close(ConnectionClosedReason::Dropped);
+        } else {
+            // If the connection has a weak reference to a pool, that means that the connection is
+            // being dropped when it's checked out. If the pool is still alive, it
+            // should check itself back in. Otherwise, the connection should close
+            // itself and emit a ConnectionClosed event (because the `close_and_drop`
+            // helper was not called explicitly).
+            //
+            // If the connection does not have a weak reference to a pool, then the connection is
+            // being dropped while it's not checked out. This means that the pool called
+            // the `close_and_drop` helper explicitly, so we don't add it back to the
+            // pool or emit any events.
+            if let Some(ref weak_pool_ref) = self.pool {
+                if let Some(strong_pool_ref) = weak_pool_ref.upgrade() {
+                    let dropped_connection_state = self.take();
+                    RUNTIME.execute(async move {
+                        strong_pool_ref
+                            .check_in(dropped_connection_state.into())
+                            .await;
+                    });
+                } else {
+                    self.close(ConnectionClosedReason::PoolClosed);
+                }
             }
         }
     }
@@ -318,6 +334,7 @@ impl From<DroppedConnectionState> for Connection {
             id: state.id,
             address: state.address.clone(),
             generation: state.generation,
+            command_executing: false,
             stream: std::mem::replace(&mut state.stream, AsyncStream::Null),
             handler: state.handler.take(),
             stream_description: state.stream_description.take(),

src/event/cmap.rs

@@ -117,6 +117,9 @@ pub enum ConnectionClosedReason {
     /// An error occurred while using the connection.
     Error,
 
+    /// The connection was dropped during read or write.
+    Dropped,
+
     /// The pool that the connection belongs to has been closed.
     PoolClosed,
 }

src/test/client.rs

@@ -9,6 +9,7 @@ use crate::{
     selection_criteria::{ReadPreference, ReadPreferenceOptions, SelectionCriteria},
     test::{util::TestClient, CLIENT_OPTIONS, LOCK},
     Client,
+    RUNTIME,
 };
 
 #[derive(Debug, Deserialize)]
@@ -57,6 +58,48 @@ async fn metadata_sent_in_handshake() {
     assert_eq!(metadata.client.driver.name, "mrd");
 }
 
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+#[function_name::named]
+async fn connection_drop_during_read() {
+    let _guard = LOCK.run_concurrently().await;
+
+    let options = CLIENT_OPTIONS.clone();
+
+    let client = Client::with_options(options.clone()).unwrap();
+    let db = client.database("test");
+
+    db.collection(function_name!())
+        .insert_one(doc! { "x": 1 }, None)
+        .await
+        .unwrap();
+
+    let _: Result<_, _> = RUNTIME
+        .timeout(
+            Duration::from_millis(50),
+            db.run_command(
+                doc! {
+                    "count": function_name!(),
+                    "query": {
+                        "$where": "sleep(100) && true"
+                    }
+                },
+                None,
+            ),
+        )
+        .await;
+
+    RUNTIME.delay_for(Duration::from_millis(200)).await;
+
+    let is_master_response = db.run_command(doc! { "isMaster": 1 }, None).await;
+
+    // Ensure that the response to `isMaster` is read, not the response to `count`.
+    assert!(is_master_response
+        .ok()
+        .and_then(|value| value.get("ismaster").and_then(|value| value.as_bool()))
+        .is_some());
+}
+
 #[cfg_attr(feature = "tokio-runtime", tokio::test)]
 #[cfg_attr(feature = "async-std-runtime", async_std::test)]
 async fn server_selection_timeout_message() {