Fix recovery token logic and add test

Author: Nathan Blinn

src/client/executor.rs

@@ -411,10 +411,8 @@ impl Client {
                         self.update_cluster_time(&r, session).await;
                         if r.is_success() {
                             // Retrieve recovery token from successful response.
-                            if let Some(ref mut session) = session {
-                                if session.in_transaction() && is_sharded {
-                                    session.transaction.recovery_token = r.recovery_token();
-                                }
+                            if is_sharded {
+                                Client::update_recovery_token(&r, session).await;
                             }
 
                             Ok(CommandResult {
@@ -461,7 +459,10 @@ impl Client {
                                         }))
                                     }
                                     // for ok: 1 just return the original deserialization error.
-                                    _ => Err(deserialize_error),
+                                    _ => {
+                                        Client::update_recovery_token(&error_response, session).await;
+                                        Err(deserialize_error)
+                                    },
                                 }
                             }
                             // We failed to deserialize even that, so just return the original
@@ -640,6 +641,17 @@ impl Client {
             }
         }
     }
+
+    async fn update_recovery_token<T: Response>(
+        response: &T,
+        session: &mut Option<&mut ClientSession>,
+    ) {
+        if let Some(ref mut session) = session {
+            if session.in_transaction() {
+                session.transaction.recovery_token = response.recovery_token().cloned();
+            }
+        }
+    }
 }
 
 impl Error {

src/operation/mod.rs

@@ -154,7 +154,7 @@ pub(crate) trait Response: Sized {
     fn at_cluster_time(&self) -> Option<Timestamp>;
 
     /// The `recoveryToken` field of the response.
-    fn recovery_token(&self) -> Option<Document>;
+    fn recovery_token(&self) -> Option<&Document>;
 
     /// Convert into the body of the response.
     fn into_body(self) -> Self::Body;
@@ -202,8 +202,8 @@ impl<T: DeserializeOwned> Response for CommandResponse<T> {
         self.at_cluster_time
     }
 
-    fn recovery_token(&self) -> Option<Document> {
-        self.recovery_token.clone()
+    fn recovery_token(&self) -> Option<&Document> {
+        self.recovery_token.as_ref()
     }
 
     fn into_body(self) -> Self::Body {
@@ -238,7 +238,7 @@ impl<T: DeserializeOwned> Response for CursorResponse<T> {
         self.response.body.cursor.at_cluster_time
     }
 
-    fn recovery_token(&self) -> Option<Document> {
+    fn recovery_token(&self) -> Option<&Document> {
         self.response.recovery_token()
     }
 

src/operation/run_command/mod.rs

@@ -138,8 +138,8 @@ impl super::Response for Response {
             .ok()
     }
 
-    fn recovery_token(&self) -> Option<Document> {
-        self.recovery_token.clone()
+    fn recovery_token(&self) -> Option<&Document> {
+        self.recovery_token.as_ref()
     }
 
     fn into_body(self) -> Self::Body {

src/test/spec/transactions.rs

@@ -5,6 +5,7 @@ use crate::{
     bson::{doc, serde_helpers::serialize_u64_as_i32, Document},
     client::session::TransactionState,
     test::{run_spec_test, TestClient, LOCK},
+    Collection,
 };
 
 use super::{run_unified_format_test, run_v2_test};
@@ -92,3 +93,56 @@ async fn client_errors() {
     assert!(result.is_err());
     assert_eq!(session.transaction.state, TransactionState::InProgress);
 }
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test(flavor = "multi_thread"))]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+#[function_name::named]
+async fn deserialize_recovery_token() {
+    let _guard: RwLockReadGuard<()> = LOCK.run_concurrently().await;
+
+    #[derive(Debug, Serialize)]
+    struct A {
+        num: i32,
+    }
+
+    #[derive(Debug, Deserialize)]
+    struct B {
+        str: String,
+    }
+
+    let client = TestClient::new().await;
+    if !client.is_sharded() || client.server_version_lt(4, 2) {
+        return;
+    }
+
+    let mut session = client.start_session(None).await.unwrap();
+
+    // Insert a document with schema A.
+    client
+        .database(function_name!())
+        .collection::<Document>(function_name!())
+        .drop(None)
+        .await
+        .unwrap();
+    client
+        .database(function_name!())
+        .create_collection(function_name!(), None)
+        .await
+        .unwrap();
+    let coll = client
+        .database(function_name!())
+        .collection(function_name!());
+    coll.insert_one(A { num: 4 }, None).await.unwrap();
+
+    // Attempt to execute Find on a document with schema B.
+    let coll: Collection<B> = client
+        .database(function_name!())
+        .collection(function_name!());
+    session.start_transaction(None).await.unwrap();
+    assert!(session.transaction.recovery_token.is_none());
+    let result = coll.find_one_with_session(None, None, &mut session).await;
+    assert!(result.is_err()); // Assert that the deserialization failed.
+
+    // Nevertheless, the recovery token should have been retrieved from the ok: 1 response.
+    assert!(session.transaction.recovery_token.is_some());
+}