CDRIVER-2386 refactor retryable writes into shared function (#1576)

Adds a new private function `mongoc_cluster_run_retryable_write` and removes duplicate retry logic. No functional changes are expected.

* rename `is_retryable` to `can_retry`

Intended to better describe. `is_retryable_write` describes if the command is eligible for retry (regardless of whether retry occurred). `can_retry` is initially true if `is_retryable_write` is true. `can_retry` is set to false after a retry (can only retry once).

* simplify condition

`_mongoc_write_error_update_if_unsupported_storage_engine` only applies if the server does not support retryable writes. Apply regardless of whether a retry occurred.

* remove `_mongoc_client_retryable_write_command_with_stream`

---------

Co-authored-by: Ezra Chung <[email protected]>

Author: kevinAlbs

src/libmongoc/src/mongoc/mongoc-client.c

@@ -1554,109 +1554,6 @@ mongoc_client_command (mongoc_client_t *client,
 }
 
 
-static bool
-_mongoc_client_retryable_write_command_with_stream (mongoc_client_t *client,
-                                                    mongoc_cmd_parts_t *parts,
-                                                    mongoc_server_stream_t *server_stream,
-                                                    bson_t *reply,
-                                                    bson_error_t *error)
-{
-   mongoc_server_stream_t *retry_server_stream = NULL;
-   bson_iter_t txn_number_iter;
-   bool is_retryable = true;
-   bool ret;
-
-   ENTRY;
-
-   BSON_ASSERT_PARAM (client);
-   BSON_ASSERT (parts->is_retryable_write);
-
-   /* increment the transaction number for the first attempt of each retryable
-    * write command */
-   BSON_ASSERT (bson_iter_init_find (&txn_number_iter, parts->assembled.command, "txnNumber"));
-   bson_iter_overwrite_int64 (&txn_number_iter, ++parts->assembled.session->server_session->txn_number);
-
-   // Store the original error and reply if needed.
-   struct {
-      bson_t reply;
-      bson_error_t error;
-      bool set;
-   } original_error = {.reply = {0}, .error = {0}, false};
-
-retry:
-   ret = mongoc_cluster_run_command_monitored (&client->cluster, &parts->assembled, reply, error);
-
-   _mongoc_write_error_handle_labels (ret, error, reply, server_stream->sd);
-
-   if (is_retryable) {
-      _mongoc_write_error_update_if_unsupported_storage_engine (ret, error, reply);
-   }
-
-   /* If a retryable error is encountered and the write is retryable, select
-    * a new writable stream and retry. If server selection fails or the selected
-    * server does not support retryable writes, fall through and allow the
-    * original error to be reported. */
-   if (is_retryable && _mongoc_write_error_get_type (reply) == MONGOC_WRITE_ERR_RETRY) {
-      bson_error_t ignored_error;
-
-      // The write command may be retried at most once.
-      is_retryable = false;
-
-      {
-         mongoc_deprioritized_servers_t *const ds = mongoc_deprioritized_servers_new ();
-
-         mongoc_deprioritized_servers_add_if_sharded (ds, server_stream->topology_type, server_stream->sd);
-
-         BSON_ASSERT (!retry_server_stream);
-         retry_server_stream =
-            mongoc_cluster_stream_for_writes (&client->cluster, parts->assembled.session, ds, NULL, &ignored_error);
-
-         mongoc_deprioritized_servers_destroy (ds);
-      }
-
-      if (retry_server_stream) {
-         parts->assembled.server_stream = retry_server_stream;
-         {
-            // Store the original error and reply before retry.
-            BSON_ASSERT (!original_error.set); // Retry only happens once.
-            original_error.set = true;
-            bson_copy_to (reply, &original_error.reply);
-            if (error) {
-               original_error.error = *error;
-            }
-         }
-         bson_destroy (reply);
-         GOTO (retry);
-      }
-   }
-
-   if (retry_server_stream) {
-      mongoc_server_stream_cleanup (retry_server_stream);
-   }
-
-   // If a retry attempt fails with an error labeled NoWritesPerformed,
-   // drivers MUST return the original error.
-   if (original_error.set && mongoc_error_has_label (reply, "NoWritesPerformed")) {
-      if (error) {
-         *error = original_error.error;
-      }
-      bson_destroy (reply);
-      bson_copy_to (&original_error.reply, reply);
-   }
-
-   if (original_error.set) {
-      bson_destroy (&original_error.reply);
-   }
-
-   if (ret && error) {
-      /* if a retry succeeded, clear the initial error */
-      memset (error, 0, sizeof (bson_error_t));
-   }
-
-   RETURN (ret);
-}
-
-
 static bool
 _mongoc_client_retryable_read_command_with_stream (mongoc_client_t *client,
                                                    mongoc_cmd_parts_t *parts,
@@ -1750,7 +1647,17 @@ _mongoc_client_command_with_stream (mongoc_client_t *client,
    }
 
    if (parts->is_retryable_write) {
-      RETURN (_mongoc_client_retryable_write_command_with_stream (client, parts, server_stream, reply, error));
+      mongoc_server_stream_t *retry_server_stream = NULL;
+
+      bool ret = mongoc_cluster_run_retryable_write (
+         &client->cluster, &parts->assembled, true /* is_retryable */, &retry_server_stream, reply, error);
+
+      if (retry_server_stream) {
+         mongoc_server_stream_cleanup (retry_server_stream);
+         parts->assembled.server_stream = NULL;
+      }
+
+      RETURN (ret);
    }
 
    if (parts->is_retryable_read) {

src/libmongoc/src/mongoc/mongoc-cluster-private.h

@@ -186,6 +186,18 @@ mongoc_cluster_stream_valid (mongoc_cluster_t *cluster, mongoc_server_stream_t *
 bool
 mongoc_cluster_run_command_monitored (mongoc_cluster_t *cluster, mongoc_cmd_t *cmd, bson_t *reply, bson_error_t *error);
 
+// `mongoc_cluster_run_retryable_write` executes a write command and may apply retryable writes behavior.
+// `cmd->server_stream` is set to `*retry_server_stream` on retry. Otherwise, it is unmodified.
+// `*retry_server_stream` is set to a new stream on retry. The caller must call `mongoc_server_stream_cleanup`.
+// `*reply` must be uninitialized and is always initialized upon return. The caller must call `bson_destroy`.
+bool
+mongoc_cluster_run_retryable_write (mongoc_cluster_t *cluster,
+                                    mongoc_cmd_t *cmd,
+                                    bool is_retryable_write,
+                                    mongoc_server_stream_t **retry_server_stream,
+                                    bson_t *reply,
+                                    bson_error_t *error);
+
 bool
 mongoc_cluster_run_command_parts (mongoc_cluster_t *cluster,
                                   mongoc_server_stream_t *server_stream,

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

@@ -3636,3 +3636,99 @@ mcd_rpc_message_decompress_if_necessary (mcd_rpc_message *rpc, void **data, size
 
    return mcd_rpc_message_decompress (rpc, data, data_len);
 }
+
+bool
+mongoc_cluster_run_retryable_write (mongoc_cluster_t *cluster,
+                                    mongoc_cmd_t *cmd,
+                                    bool is_retryable_write,
+                                    mongoc_server_stream_t **retry_server_stream,
+                                    bson_t *reply,
+                                    bson_error_t *error)
+{
+   BSON_ASSERT_PARAM (cluster);
+   BSON_ASSERT_PARAM (cmd);
+   BSON_ASSERT_PARAM (retry_server_stream);
+   BSON_ASSERT_PARAM (reply);
+   BSON_ASSERT (error || true);
+
+   bool ret;
+   // `can_retry` is set to false on retry. A retry may only happen once.
+   bool can_retry = is_retryable_write;
+
+   // Increment the transaction number for the first attempt of each retryable write command.
+   if (is_retryable_write) {
+      bson_iter_t txn_number_iter;
+      BSON_ASSERT (bson_iter_init_find (&txn_number_iter, cmd->command, "txnNumber"));
+      bson_iter_overwrite_int64 (&txn_number_iter, ++cmd->session->server_session->txn_number);
+   }
+
+   // Store the original error and reply if needed.
+   struct {
+      bson_t reply;
+      bson_error_t error;
+      bool set;
+   } original_error = {.reply = {0}, .error = {0}, .set = false};
+
+   // Ensure `*retry_server_stream` is always valid or null.
+   *retry_server_stream = NULL;
+
+retry:
+   ret = mongoc_cluster_run_command_monitored (cluster, cmd, reply, error);
+
+   if (is_retryable_write) {
+      _mongoc_write_error_handle_labels (ret, error, reply, cmd->server_stream->sd);
+      _mongoc_write_error_update_if_unsupported_storage_engine (ret, error, reply);
+   }
+
+   // If a retryable error is encountered and the write is retryable, select a new writable stream and retry. If server
+   // selection fails or the selected server does not support retryable writes, fall through and allow the original
+   // error to be reported.
+   if (can_retry && _mongoc_write_error_get_type (reply) == MONGOC_WRITE_ERR_RETRY) {
+      bson_error_t ignored_error;
+
+      can_retry = false; // Only retry once.
+
+      // Select a server.
+      {
+         mongoc_deprioritized_servers_t *const ds = mongoc_deprioritized_servers_new ();
+
+         // If talking to a sharded cluster, deprioritize the just-used mongos to prefer a new mongos for the retry.
+         mongoc_deprioritized_servers_add_if_sharded (ds, cmd->server_stream->topology_type, cmd->server_stream->sd);
+
+         *retry_server_stream =
+            mongoc_cluster_stream_for_writes (cluster, cmd->session, ds, NULL /* reply */, &ignored_error);
+
+         mongoc_deprioritized_servers_destroy (ds);
+      }
+
+      if (*retry_server_stream) {
+         cmd->server_stream = *retry_server_stream; // Non-owning.
+         {
+            // Store the original error and reply before retry.
+            BSON_ASSERT (!original_error.set); // Retry only happens once.
+            original_error.set = true;
+            bson_copy_to (reply, &original_error.reply);
+            if (error) {
+               original_error.error = *error;
+            }
+         }
+         bson_destroy (reply);
+         GOTO (retry);
+      }
+   }
+
+   // If a retry attempt fails with an error labeled NoWritesPerformed, drivers MUST return the original error.
+   if (original_error.set && mongoc_error_has_label (reply, "NoWritesPerformed")) {
+      if (error) {
+         *error = original_error.error;
+      }
+      bson_destroy (reply);
+      bson_copy_to (&original_error.reply, reply);
+   }
+
+   if (original_error.set) {
+      bson_destroy (&original_error.reply);
+   }
+
+   RETURN (ret);
+}

src/libmongoc/src/mongoc/mongoc-collection.c

@@ -3130,13 +3130,11 @@ mongoc_collection_find_and_modify_with_opts (mongoc_collection_t *collection,
 {
    mongoc_cluster_t *cluster;
    mongoc_cmd_parts_t parts;
-   bool is_retryable;
    bson_iter_t iter;
    bson_iter_t inner;
    const char *name;
    bson_t ss_reply;
-   bson_t reply_local;
-   bson_t *reply_ptr;
+   bson_t reply_local = BSON_INITIALIZER;
    bool ret = false;
    bson_t command = BSON_INITIALIZER;
    mongoc_server_stream_t *server_stream = NULL;
@@ -3150,23 +3148,26 @@ mongoc_collection_find_and_modify_with_opts (mongoc_collection_t *collection,
    BSON_ASSERT_PARAM (query);
    BSON_ASSERT_PARAM (opts);
 
-   reply_ptr = reply ? reply : &reply_local;
+   if (reply) {
+      bson_init (reply);
+   } else {
+      // Caller did not pass an output `reply`. Use a local `reply` to determine if a server error is retryable.
+      reply = &reply_local;
+   }
    cluster = &collection->client->cluster;
 
    mongoc_cmd_parts_init (&parts, collection->client, collection->db, MONGOC_QUERY_NONE, &command);
    parts.is_read_command = true;
    parts.is_write_command = true;
 
-   bson_init (reply_ptr);
-
    if (!_mongoc_find_and_modify_appended_opts_parse (cluster->client, &opts->extra, &appended_opts, error)) {
       GOTO (done);
    }
 
    server_stream = mongoc_cluster_stream_for_writes (cluster, appended_opts.client_session, NULL, &ss_reply, error);
 
    if (!server_stream) {
-      bson_concat (reply_ptr, &ss_reply);
+      bson_concat (reply, &ss_reply);
       bson_destroy (&ss_reply);
       GOTO (done);
    }
@@ -3282,86 +3283,11 @@ mongoc_collection_find_and_modify_with_opts (mongoc_collection_t *collection,
       GOTO (done);
    }
 
-   is_retryable = parts.is_retryable_write;
-
-   /* increment the transaction number for the first attempt of each retryable
-    * write command */
-   if (is_retryable) {
-      bson_iter_t txn_number_iter;
-      BSON_ASSERT (bson_iter_init_find (&txn_number_iter, parts.assembled.command, "txnNumber"));
-      bson_iter_overwrite_int64 (&txn_number_iter, ++parts.assembled.session->server_session->txn_number);
-   }
-
-   // Store the original error and reply if needed.
-   struct {
-      bson_t reply;
-      bson_error_t error;
-      bool set;
-   } original_error = {.reply = {0}, .error = {0}, .set = false};
-
-retry:
-   bson_destroy (reply_ptr);
-   ret = mongoc_cluster_run_command_monitored (cluster, &parts.assembled, reply_ptr, error);
-
-   if (parts.is_retryable_write) {
-      _mongoc_write_error_handle_labels (ret, error, reply_ptr, server_stream->sd);
-   }
-
-   if (is_retryable) {
-      _mongoc_write_error_update_if_unsupported_storage_engine (ret, error, reply_ptr);
-   }
-
-   /* If a retryable error is encountered and the write is retryable, select
-    * a new writable stream and retry. If server selection fails or the selected
-    * server does not support retryable writes, fall through and allow the
-    * original error to be reported. */
-   if (is_retryable && _mongoc_write_error_get_type (reply_ptr) == MONGOC_WRITE_ERR_RETRY) {
-      bson_error_t ignored_error;
+   bson_destroy (reply);
+   ret = mongoc_cluster_run_retryable_write (
+      cluster, &parts.assembled, parts.is_retryable_write, &retry_server_stream, reply, error);
 
-      /* each write command may be retried at most once */
-      is_retryable = false;
-
-      {
-         mongoc_deprioritized_servers_t *const ds = mongoc_deprioritized_servers_new ();
-
-         mongoc_deprioritized_servers_add_if_sharded (ds, server_stream->topology_type, server_stream->sd);
-
-         retry_server_stream =
-            mongoc_cluster_stream_for_writes (cluster, parts.assembled.session, ds, NULL /* reply */, &ignored_error);
-
-         mongoc_deprioritized_servers_destroy (ds);
-      }
-
-      if (retry_server_stream) {
-         parts.assembled.server_stream = retry_server_stream;
-         {
-            // Store the original error and reply before retry.
-            BSON_ASSERT (!original_error.set); // Retry only happens once.
-            original_error.set = true;
-            bson_copy_to (reply_ptr, &original_error.reply);
-            if (error) {
-               original_error.error = *error;
-            }
-         }
-         GOTO (retry);
-      }
-   }
-
-   // If a retry attempt fails with an error labeled NoWritesPerformed,
-   // drivers MUST return the original error.
-   if (original_error.set && mongoc_error_has_label (reply_ptr, "NoWritesPerformed")) {
-      if (error) {
-         *error = original_error.error;
-      }
-      bson_destroy (reply_ptr);
-      bson_copy_to (&original_error.reply, reply_ptr);
-   }
-
-   if (original_error.set) {
-      bson_destroy (&original_error.reply);
-   }
-
-   if (bson_iter_init_find (&iter, reply_ptr, "writeConcernError") && BSON_ITER_HOLDS_DOCUMENT (&iter)) {
+   if (bson_iter_init_find (&iter, reply, "writeConcernError") && BSON_ITER_HOLDS_DOCUMENT (&iter)) {
       const char *errmsg = NULL;
       int32_t code = 0;
 
@@ -3388,9 +3314,7 @@ mongoc_collection_find_and_modify_with_opts (mongoc_collection_t *collection,
    }
    mongoc_cmd_parts_cleanup (&parts);
    bson_destroy (&command);
-   if (&reply_local == reply_ptr) {
-      bson_destroy (&reply_local);
-   }
+   bson_destroy (&reply_local);
    _mongoc_find_and_modify_appended_opts_cleanup (&appended_opts);
    RETURN (ret);
 }