CDRIVER-4363 Support multiple OP_MSG doc sequences (#1578)

Support sending an OP_MSG with two document sequences (OP_MSG Section with payloadType=1). Intended for upcoming support of the `bulkWrite` command. Drivers are expected to send document sequences for both `nsInfo` and `ops` fields.

* remove unnecessary `_mongoc_get_documents_field_name`

The payloadType=1 identifier is already available in `mongoc_cmd_t::payload_identifier`

* define a `mongoc_cmd_payload_t` struct

* define `MONGOC_CMD_PAYLOADS_COUNT_MAX`

---------

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

Author: kevinAlbs

src/libmongoc/src/mongoc/mongoc-apm.c

@@ -31,7 +31,8 @@ static bson_oid_t kObjectIdZero = {{0}};
 static void
 append_documents_from_cmd (const mongoc_cmd_t *cmd, mongoc_apm_command_started_t *event)
 {
-   if (!cmd->payload || !cmd->payload_size) {
+   // If there are no document sequences (OP_MSG Section with payloadType=1), return the command unchanged.
+   if (cmd->payloads_count == 0) {
       return;
    }
 

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

@@ -1010,8 +1010,8 @@ append_bson_range_opts (bson_t *bson_range_opts, const mongoc_client_encryption_
 static void
 _prep_for_auto_encryption (const mongoc_cmd_t *cmd, bson_t *out)
 {
-   /* If there is no type=1 payload, return the command unchanged. */
-   if (!cmd->payload || !cmd->payload_size) {
+   // If there are no document sequences (OP_MSG Section with payloadType=1), return the command unchanged.
+   if (cmd->payloads_count == 0) {
       BSON_ASSERT (bson_init_static (out, bson_get_data (cmd->command), cmd->command->len));
       return;
    }
@@ -1226,10 +1226,9 @@ _mongoc_cse_auto_encrypt (mongoc_client_t *client_encrypted,
 
    /* Create the modified cmd_t. */
    memcpy (encrypted_cmd, cmd, sizeof (mongoc_cmd_t));
-   /* Modify the mongoc_cmd_t and clear the payload, since
-    * _mongoc_cse_auto_encrypt converted the payload into an embedded array. */
-   encrypted_cmd->payload = NULL;
-   encrypted_cmd->payload_size = 0;
+   /* Modify the mongoc_cmd_t and clear the payloads, since
+    * _mongoc_cse_auto_encrypt converted the payloads into an embedded array. */
+   encrypted_cmd->payloads_count = 0;
    encrypted_cmd->command = encrypted;
 
    ret = true;

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

@@ -3227,24 +3227,28 @@ _mongoc_cluster_run_opmsg_send (
       message_length += mcd_rpc_header_set_response_to (rpc, 0);
       message_length += mcd_rpc_header_set_op_code (rpc, MONGOC_OP_CODE_MSG);
 
-      mcd_rpc_op_msg_set_sections_count (rpc, cmd->payload ? 2u : 1u);
+      BSON_ASSERT (cmd->payloads_count <= MONGOC_CMD_PAYLOADS_COUNT_MAX);
+      // Reserve one section for the body (kind 0) and any needed sections for document sequences (kind 1)
+      mcd_rpc_op_msg_set_sections_count (rpc, 1u + cmd->payloads_count);
 
       message_length += mcd_rpc_op_msg_set_flag_bits (rpc, flags);
       message_length += mcd_rpc_op_msg_section_set_kind (rpc, 0u, 0);
       message_length += mcd_rpc_op_msg_section_set_body (rpc, 0u, bson_get_data (cmd->command));
 
-      if (cmd->payload) {
-         BSON_ASSERT (bson_in_range_signed (size_t, cmd->payload_size));
+      for (size_t i = 0; i < cmd->payloads_count; i++) {
+         const mongoc_cmd_payload_t payload = cmd->payloads[i];
 
-         const size_t section_length =
-            sizeof (int32_t) + strlen (cmd->payload_identifier) + 1u + (size_t) cmd->payload_size;
+         BSON_ASSERT (bson_in_range_signed (size_t, payload.size));
+
+         const size_t section_length = sizeof (int32_t) + strlen (payload.identifier) + 1u + (size_t) payload.size;
          BSON_ASSERT (bson_in_range_unsigned (int32_t, section_length));
 
-         message_length += mcd_rpc_op_msg_section_set_kind (rpc, 1u, 1);
-         message_length += mcd_rpc_op_msg_section_set_length (rpc, 1u, (int32_t) section_length);
-         message_length += mcd_rpc_op_msg_section_set_identifier (rpc, 1u, cmd->payload_identifier);
+         size_t section_idx = 1u + i;
+         message_length += mcd_rpc_op_msg_section_set_kind (rpc, section_idx, 1);
+         message_length += mcd_rpc_op_msg_section_set_length (rpc, section_idx, (int32_t) section_length);
+         message_length += mcd_rpc_op_msg_section_set_identifier (rpc, section_idx, payload.identifier);
          message_length +=
-            mcd_rpc_op_msg_section_set_document_sequence (rpc, 1u, cmd->payload, (size_t) cmd->payload_size);
+            mcd_rpc_op_msg_section_set_document_sequence (rpc, section_idx, payload.documents, (size_t) payload.size);
       }
 
       mcd_rpc_message_set_length (rpc, message_length);

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

@@ -51,14 +51,24 @@ typedef enum {
    MONGOC_CMD_PARTS_ALLOW_TXN_NUMBER_NO
 } mongoc_cmd_parts_allow_txn_number_t;
 
+// `mongoc_cmd_payload_t` represents a document sequence (OP_MSG Section with payloadType=1).
+typedef struct {
+   int32_t size;
+   const char *identifier;
+   const uint8_t *documents;
+} mongoc_cmd_payload_t;
+
+// OP_MSG supports any number of document sequences. Increase array size to support more document sequences.
+#define MONGOC_CMD_PAYLOADS_COUNT_MAX 2
+
 typedef struct _mongoc_cmd_t {
    const char *db_name;
    mongoc_query_flags_t query_flags;
    const bson_t *command;
    const char *command_name;
-   const uint8_t *payload;
-   int32_t payload_size;
-   const char *payload_identifier;
+   size_t payloads_count;
+   // `payloads[i]` may be read only when `0 <= i < payloads_count`.
+   mongoc_cmd_payload_t payloads[MONGOC_CMD_PAYLOADS_COUNT_MAX];
    mongoc_server_stream_t *server_stream;
    int64_t operation_id;
    mongoc_client_session_t *session;

src/libmongoc/src/mongoc/mongoc-cmd.c

@@ -55,8 +55,7 @@ mongoc_cmd_parts_init (mongoc_cmd_parts_t *parts,
    parts->assembled.command = NULL;
    parts->assembled.query_flags = MONGOC_QUERY_NONE;
    parts->assembled.op_msg_is_exhaust = false;
-   parts->assembled.payload_identifier = NULL;
-   parts->assembled.payload = NULL;
+   parts->assembled.payloads_count = 0;
    parts->assembled.session = NULL;
    parts->assembled.is_acknowledged = true;
    parts->assembled.is_txn_finish = false;
@@ -960,27 +959,11 @@ mongoc_cmd_is_compressible (mongoc_cmd_t *cmd)
           !!strcasecmp (cmd->command_name, "createuser") && !!strcasecmp (cmd->command_name, "updateuser");
 }
 
-/*--------------------------------------------------------------------------
- *
- * _mongoc_cmd_append_payload_as_array --
- *    Append a write command payload as an array in a BSON document.
- *    Used by APM and Client-Side Encryption
- *
- * Arguments:
- *    cmd The mongoc_cmd_t, which may contain a payload to be appended.
- *    out A bson_t, which will be appended to if @cmd->payload is set.
- *
- * Pre-conditions:
- *    - @out is initialized.
- *    - cmd has a payload (i.e. is a write command).
- *
- * Post-conditions:
- *    - If @cmd->payload is set, then @out is appended to with the payload
- *      field's name ("documents" if insert, "updates" if update,
- *      "deletes" if delete) an the payload as a BSON array.
- *
- *--------------------------------------------------------------------------
- */
+
+//`_mongoc_cmd_append_payload_as_array` appends document seqence payloads as BSON arrays.
+// `cmd` must contain one or more document sequence payloads (`cmd->payloads_count` > 0).
+// `out` must be initialized by the caller.
+// Used by APM and In-Use Encryption (document sequences are not supported for auto encryption).
 void
 _mongoc_cmd_append_payload_as_array (const mongoc_cmd_t *cmd, bson_t *out)
 {
@@ -990,25 +973,28 @@ _mongoc_cmd_append_payload_as_array (const mongoc_cmd_t *cmd, bson_t *out)
    const char *field_name;
    bson_array_builder_t *bson;
 
-   BSON_ASSERT (cmd->payload && cmd->payload_size);
+   BSON_ASSERT (cmd->payloads_count > 0);
+   BSON_ASSERT (cmd->payloads_count <= MONGOC_CMD_PAYLOADS_COUNT_MAX);
 
-   /* make array from outgoing OP_MSG payload type 1 on an "insert",
-    * "update", or "delete" command. */
-   field_name = _mongoc_get_documents_field_name (cmd->command_name);
-   BSON_ASSERT (field_name);
-   BSON_ASSERT (BSON_APPEND_ARRAY_BUILDER_BEGIN (out, field_name, &bson));
+   for (size_t i = 0; i < cmd->payloads_count; i++) {
+      BSON_ASSERT (cmd->payloads[i].documents && cmd->payloads[i].size);
 
-   pos = cmd->payload;
-   while (pos < cmd->payload + cmd->payload_size) {
-      memcpy (&doc_len, pos, sizeof (doc_len));
-      doc_len = BSON_UINT32_FROM_LE (doc_len);
-      BSON_ASSERT (bson_init_static (&doc, pos, (size_t) doc_len));
-      bson_array_builder_append_document (bson, &doc);
+      // Create a BSON array from a document sequence (OP_MSG Section with payloadType=1).
+      field_name = cmd->payloads[i].identifier;
+      BSON_ASSERT (field_name);
+      BSON_ASSERT (BSON_APPEND_ARRAY_BUILDER_BEGIN (out, field_name, &bson));
 
-      pos += doc_len;
-   }
+      pos = cmd->payloads[i].documents;
+      while (pos < cmd->payloads[i].documents + cmd->payloads[i].size) {
+         memcpy (&doc_len, pos, sizeof (doc_len));
+         doc_len = BSON_UINT32_FROM_LE (doc_len);
+         BSON_ASSERT (bson_init_static (&doc, pos, (size_t) doc_len));
+         bson_array_builder_append_document (bson, &doc);
 
-   bson_append_array_builder_end (out, bson);
+         pos += doc_len;
+      }
+      bson_append_array_builder_end (out, bson);
+   }
 }
 
 /*--------------------------------------------------------------------------

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

@@ -82,9 +82,6 @@ _mongoc_get_real_time_ms (void);
 const char *
 _mongoc_get_command_name (const bson_t *command);
 
-const char *
-_mongoc_get_documents_field_name (const char *command_name);
-
 bool
 _mongoc_lookup_bool (const bson_t *bson, const char *key, bool default_value);
 

src/libmongoc/src/mongoc/mongoc-util.c

@@ -163,25 +163,6 @@ _mongoc_get_command_name (const bson_t *command)
    return name;
 }
 
-
-const char *
-_mongoc_get_documents_field_name (const char *command_name)
-{
-   if (!strcmp (command_name, "insert")) {
-      return "documents";
-   }
-
-   if (!strcmp (command_name, "update")) {
-      return "updates";
-   }
-
-   if (!strcmp (command_name, "delete")) {
-      return "deletes";
-   }
-
-   return NULL;
-}
-
 bool
 _mongoc_lookup_bool (const bson_t *bson, const char *key, bool default_value)
 {

src/libmongoc/src/mongoc/mongoc-write-command.c

@@ -664,11 +664,13 @@ _mongoc_write_opmsg (mongoc_write_command_t *command,
       }
 
       if (ship_it) {
+         parts.assembled.payloads_count = 1;
+         mongoc_cmd_payload_t *const payload = &parts.assembled.payloads[0];
          /* Seek past the document offset we have already sent */
-         parts.assembled.payload = command->payload.data + payload_total_offset;
+         payload->documents = command->payload.data + payload_total_offset;
          /* Only send the documents up to this size */
-         parts.assembled.payload_size = payload_batch_size;
-         parts.assembled.payload_identifier = gCommandFields[command->type];
+         payload->size = payload_batch_size;
+         payload->identifier = gCommandFields[command->type];
 
 
          mongoc_server_stream_t *new_retry_server_stream = NULL;

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

@@ -23,6 +23,7 @@
 #include "mock_server/mock-server.h"
 #include "mock_server/future-functions.h"
 #include "test-libmongoc.h"
+#include "mongoc-cluster-private.h"
 
 #undef MONGOC_LOG_DOMAIN
 #define MONGOC_LOG_DOMAIN "cmd-test-options"
@@ -68,9 +69,88 @@ test_client_cmd_options (void)
    mock_server_destroy (server);
 }
 
+static void
+capture_last_command (const mongoc_apm_command_started_t *event)
+{
+   bson_t *last_captured = mongoc_apm_command_started_get_context (event);
+   bson_destroy (last_captured);
+   const bson_t *cmd = mongoc_apm_command_started_get_command (event);
+   bson_copy_to (cmd, last_captured);
+}
+
+// `test_cmd_with_two_payload1` tests sending an OP_MSG with two document sequence payloads (payloadType=1).
+static void
+test_cmd_with_two_payload1 (void *ctx)
+{
+   BSON_UNUSED (ctx);
+   mongoc_client_t *client = test_framework_new_default_client ();
+
+   bson_t last_captured = BSON_INITIALIZER;
+   // Set callback to capture the last command.
+   {
+      mongoc_apm_callbacks_t *cbs = mongoc_apm_callbacks_new ();
+      mongoc_apm_set_command_started_cb (cbs, capture_last_command);
+      mongoc_client_set_apm_callbacks (client, cbs, &last_captured);
+      mongoc_apm_callbacks_destroy (cbs);
+   }
+
+   mongoc_cluster_t *cluster = &client->cluster;
+   bson_error_t error;
+
+   // Use `bulkWrite`. Currently, only the `bulkWrite` command supports two document sequence payloads.
+   bson_t *payload0 = tmp_bson (BSON_STR ({"bulkWrite" : 1}));
+   bson_t *op = tmp_bson (BSON_STR ({"insert" : 0, "document" : {}}));
+   bson_t *nsInfo = tmp_bson (BSON_STR ({"ns" : "db.coll"}));
+
+   // Create the `mongoc_cmd_t`.
+   mongoc_cmd_parts_t parts;
+   mongoc_cmd_parts_init (&parts, client, "admin", MONGOC_QUERY_NONE, payload0);
+   mongoc_server_stream_t *server_stream = mongoc_cluster_stream_for_writes (
+      cluster, NULL /* session */, NULL /* deprioritized servers */, NULL /* reply */, &error);
+   ASSERT_OR_PRINT (server_stream, error);
+   bool ok = mongoc_cmd_parts_assemble (&parts, server_stream, &error);
+   ASSERT_OR_PRINT (ok, error);
+
+   parts.assembled.payloads_count = 2;
+
+   // Set `ops` as a payload1 (of one document)
+   parts.assembled.payloads[0].identifier = "ops";
+   parts.assembled.payloads[0].documents = bson_get_data (op);
+   parts.assembled.payloads[0].size = op->len;
+
+   // Set `nsInfo` as a payload1 (of one document)
+   parts.assembled.payloads[1].identifier = "nsInfo";
+   parts.assembled.payloads[1].documents = bson_get_data (nsInfo);
+   parts.assembled.payloads[1].size = nsInfo->len;
+
+   // Run the command.
+   bson_t reply;
+   ok = mongoc_cluster_run_command_monitored (cluster, &parts.assembled, &reply, &error);
+   ASSERT_OR_PRINT (ok, error);
+   ASSERT_MATCH (&reply, BSON_STR ({"ok" : 1}));
+
+   // Check that document sequences are converted to a BSON arrays for command monitoring.
+   ASSERT_MATCH (
+      &last_captured,
+      BSON_STR ({"bulkWrite" : 1, "ops" : [ {"insert" : 0, "document" : {}} ], "nsInfo" : [ {"ns" : "db.coll"} ]}));
+
+   bson_destroy (&reply);
+   mongoc_server_stream_cleanup (server_stream);
+   mongoc_cmd_parts_cleanup (&parts);
+   mongoc_client_destroy (client);
+   // Destroy `last_captured` after `client`. `mongoc_client_destroy` sends an `endSessions` command.
+   bson_destroy (&last_captured);
+}
 
 void
 test_client_cmd_install (TestSuite *suite)
 {
    TestSuite_AddMockServerTest (suite, "/Client/cmd/options", test_client_cmd_options);
+   TestSuite_AddFull (suite,
+                      "/cmd/with_two_payload1",
+                      test_cmd_with_two_payload1,
+                      NULL /* dtor */,
+                      NULL /* ctx */,
+                      test_framework_skip_if_max_wire_version_less_than_25 // require server 8.0 for `bulkWrite`
+   );
 }