CDRIVER-4718 store recoveryToken for Load Balanced topology (#1417)

* add regression test

* format mongoc-cluster.c

* store recoveryToken for loadBalanced topology

* make `_in_sharded_txn` static

* use `static`

Co-authored-by: Kyle Kloberdanz <[email protected]>

* removing redundant calls to `_mongoc_topology_get_type`

Co-authored-by: Kyle Kloberdanz <[email protected]>

Author: kevinAlbs

src/libmongoc/src/mongoc/mongoc-cluster.c

@@ -493,14 +493,32 @@ mongoc_cluster_run_command_opquery (mongoc_cluster_t *cluster,
    RETURN (ret);
 }
 
-bool
+static bool
 _in_sharded_txn (const mongoc_client_session_t *session)
 {
    return session && _mongoc_client_session_in_txn_or_ending (session) &&
           _mongoc_topology_get_type (session->client->topology) ==
              MONGOC_TOPOLOGY_SHARDED;
 }
 
+static bool
+_in_sharded_or_loadbalanced_txn (const mongoc_client_session_t *session)
+{
+   if (!session) {
+      return false;
+   }
+
+   if (!_mongoc_client_session_in_txn_or_ending (session)) {
+      return false;
+   }
+
+   mongoc_topology_description_type_t type =
+      _mongoc_topology_get_type (session->client->topology);
+
+   return (type == MONGOC_TOPOLOGY_SHARDED) ||
+          (type == MONGOC_TOPOLOGY_LOAD_BALANCED);
+}
+
 static void
 _handle_txn_error_labels (bool cmd_ret,
                           const bson_error_t *cmd_err,
@@ -655,7 +673,7 @@ mongoc_cluster_run_command_monitored (mongoc_cluster_t *cluster,
 
    _handle_txn_error_labels (retval, error, cmd, reply);
 
-   if (retval && _in_sharded_txn (cmd->session) &&
+   if (retval && _in_sharded_or_loadbalanced_txn (cmd->session) &&
        bson_iter_init_find (&iter, reply, "recoveryToken")) {
       bson_destroy (cmd->session->recovery_token);
       if (BSON_ITER_HOLDS_DOCUMENT (&iter)) {
@@ -893,7 +911,7 @@ _stream_run_hello (mongoc_cluster_t *cluster,
       if (negotiate_sasl_supported_mechs) {
          bsonParse (reply,
                     find (allOf (key ("ok"), isFalse), //
-                          do({
+                          do ({
                              /* hello response returned ok: 0. According to
                               * auth spec: "If the hello of the MongoDB
                               * Handshake fails with an error, drivers MUST

src/libmongoc/tests/test-mongoc-loadbalanced.c

@@ -823,6 +823,98 @@ skip_if_not_loadbalanced (void)
    return test_framework_is_loadbalanced () ? 1 : 0;
 }
 
+static void
+store_last_command_started_callback (const mongoc_apm_command_started_t *event)
+{
+   bson_t **last_command = mongoc_apm_command_started_get_context (event);
+   const bson_t *cmd = mongoc_apm_command_started_get_command (event);
+   bson_destroy (*last_command);
+   *last_command = bson_copy (cmd);
+}
+
+// `test_loadbalanced_sends_recoveryToken` is a regression test for
+// CDRIVER-4718. Ensure that a `recoveryToken` is included in the outgoing
+// `commitTransaction` and `abortTransaction` commands when connected to a load
+// balanced cluster.
+static void
+test_loadbalanced_sends_recoveryToken (void *unused)
+{
+   mongoc_client_t *client;
+   bson_error_t error;
+   bson_t *last_command = NULL;
+
+   BSON_UNUSED (unused);
+
+   client = test_framework_new_default_client ();
+   // Set a callback to store the most recent command started.
+   {
+      mongoc_apm_callbacks_t *cbs = mongoc_apm_callbacks_new ();
+      mongoc_apm_set_command_started_cb (cbs,
+                                         store_last_command_started_callback);
+      mongoc_client_set_apm_callbacks (client, cbs, &last_command);
+      mongoc_apm_callbacks_destroy (cbs);
+   }
+
+   mongoc_client_session_t *session =
+      mongoc_client_start_session (client, NULL /* opts */, &error);
+   ASSERT_OR_PRINT (session, error);
+
+   mongoc_collection_t *coll =
+      mongoc_client_get_collection (client, "db", "coll");
+
+   // Commit a transaction. Expect `commitTransaction` to include
+   // `recoveryToken`.
+   {
+      bool ok = mongoc_client_session_start_transaction (
+         session, NULL /* opts */, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      bson_t *insert_opts = tmp_bson ("{}");
+      ok = mongoc_client_session_append (session, insert_opts, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ok = mongoc_collection_insert_one (
+         coll, tmp_bson ("{}"), insert_opts, NULL, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ok = mongoc_client_session_commit_transaction (
+         session, NULL /* reply */, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ASSERT_MATCH (
+         last_command,
+         "{'commitTransaction': 1, 'recoveryToken': { '$exists': true } }");
+   }
+
+   // Abort a transaction. Expect `abortTransaction` to include
+   // `recoveryToken`.
+   {
+      bool ok = mongoc_client_session_start_transaction (
+         session, NULL /* opts */, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      bson_t *insert_opts = tmp_bson ("{}");
+      ok = mongoc_client_session_append (session, insert_opts, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ok = mongoc_collection_insert_one (
+         coll, tmp_bson ("{}"), insert_opts, NULL, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ok = mongoc_client_session_abort_transaction (session, &error);
+      ASSERT_OR_PRINT (ok, error);
+
+      ASSERT_MATCH (
+         last_command,
+         "{'abortTransaction': 1, 'recoveryToken': { '$exists': true } }");
+   }
+
+   mongoc_collection_destroy (coll);
+   mongoc_client_session_destroy (session);
+   mongoc_client_destroy (client);
+   bson_destroy (last_command);
+}
+
 void
 test_loadbalanced_install (TestSuite *suite)
 {
@@ -893,4 +985,11 @@ test_loadbalanced_install (TestSuite *suite)
       suite,
       "/loadbalanced/post_handshake_error_clears_pool",
       test_post_handshake_error_clears_pool);
+
+   TestSuite_AddFull (suite,
+                      "/loadbalanced/sends_recoveryToken",
+                      test_loadbalanced_sends_recoveryToken,
+                      NULL /* ctx */,
+                      NULL /* dtor */,
+                      skip_if_not_loadbalanced);
 }