CDRIVER-3230 fix opquery batch splitting

OP_QUERY batch splitting was not taking command options into account
when computing the size to split. So, for example, a write concern
document could put the batch over the maxBsonObjectSize + 16k limit.

Author: kevinAlbs

src/libmongoc/src/mongoc/mongoc-write-command-private.h

@@ -37,6 +37,11 @@ struct _mongoc_crud_opts_t;
 #define MONGOC_WRITE_COMMAND_INSERT 1
 #define MONGOC_WRITE_COMMAND_UPDATE 2
 
+/* MongoDB has a extra allowance to allow updating 16mb document, as the update
+ * operators would otherwise overflow the 16mb object limit. See SERVER-10643
+ * for context. */
+#define BSON_OBJECT_ALLOWANCE (16 * 1024)
+
 struct _mongoc_bulk_write_flags_t {
    bool ordered;
    bool bypass_document_validation;

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

@@ -418,7 +418,7 @@ _mongoc_write_command_will_overflow (uint32_t len_so_far,
    /* max BSON object size + 16k bytes.
     * server guarantees there is enough room: SERVER-10643
     */
-   int32_t max_cmd_size = max_bson_size + 16384;
+   int32_t max_cmd_size = max_bson_size + BSON_OBJECT_ALLOWANCE;
 
    BSON_ASSERT (max_bson_size);
 
@@ -469,10 +469,6 @@ _mongoc_write_opmsg (mongoc_write_command_t *command,
    BSON_ASSERT (server_stream);
    BSON_ASSERT (collection);
 
-/* MongoDB has a extra allowance to allow updating 16mb document,
- * as the update operators would otherwise overflow the 16mb object limit
- */
-#define BSON_OBJECT_ALLOWANCE (16 * 1024)
    max_bson_obj_size = mongoc_server_stream_max_bson_obj_size (server_stream);
    max_msg_size = mongoc_server_stream_max_msg_size (server_stream);
    max_document_count =
@@ -683,6 +679,41 @@ _append_array_from_command (mongoc_write_command_t *command, bson_t *bson)
    bson_reader_destroy (reader);
 }
 
+/* Assemble the base @cmd with all of the command options.
+ * @parts is always initialized, even on error.
+ * This is called twice in _mongoc_write_opquery.
+ * Once with no payload documents, to determine the total size. And once with
+ * payload documents, to send the final command. */
+static bool
+_assemble_cmd (bson_t *cmd,
+               mongoc_write_command_t *command,
+               mongoc_client_t *client,
+               mongoc_server_stream_t *server_stream,
+               const char *database,
+               const mongoc_write_concern_t *write_concern,
+               mongoc_cmd_parts_t *parts,
+               bson_error_t *error)
+{
+   bool ret;
+   bson_iter_t iter;
+
+   mongoc_cmd_parts_init (parts, client, database, MONGOC_QUERY_NONE, cmd);
+   parts->is_write_command = true;
+   parts->assembled.operation_id = command->operation_id;
+
+   ret = mongoc_cmd_parts_set_write_concern (
+      parts, write_concern, server_stream->sd->max_wire_version, error);
+   if (ret) {
+      BSON_ASSERT (bson_iter_init (&iter, &command->cmd_opts));
+      ret = mongoc_cmd_parts_append_opts (
+         parts, &iter, server_stream->sd->max_wire_version, error);
+   }
+   if (ret) {
+      ret = mongoc_cmd_parts_assemble (parts, server_stream, error);
+   }
+   return ret;
+}
+
 static void
 _mongoc_write_opquery (mongoc_write_command_t *command,
                        mongoc_client_t *client,
@@ -695,12 +726,10 @@ _mongoc_write_opquery (mongoc_write_command_t *command,
                        bson_error_t *error)
 {
    mongoc_cmd_parts_t parts;
-   bson_iter_t iter;
    const char *key;
    uint32_t len = 0;
    bson_t ar;
    bson_t cmd;
-   bson_t reply;
    char str[16];
    bool has_more;
    bool ret = false;
@@ -732,10 +761,30 @@ _mongoc_write_opquery (mongoc_write_command_t *command,
    i = 0;
 
    _mongoc_write_command_init (&cmd, command, collection);
+   /* If any part of assembling failed, return with failure. */
+   if (!_assemble_cmd (&cmd,
+                       command,
+                       client,
+                       server_stream,
+                       database,
+                       write_concern,
+                       &parts,
+                       error)) {
+      result->failed = true;
+      bson_destroy (&cmd);
+      mongoc_cmd_parts_cleanup (&parts);
+      EXIT;
+   }
 
-   /* 1 byte to specify array type, 1 byte for field name's null terminator */
-   overhead = cmd.len + 2 + gCommandFieldLens[command->type];
-
+   /* Use the assembled command to compute the overhead, since it may be a new
+    * BSON document with options applied. If no options were applied, then
+    * parts.assembled.command points to cmd. The constant 2 is due to 1 byte to
+    * specify array type and 1 byte for field name's null terminator. */
+   overhead =
+      parts.assembled.command->len + 2 + gCommandFieldLens[command->type];
+   /* Toss out the assembled command, we'll assemble again after adding all of
+    * the payload documents. */
+   mongoc_cmd_parts_cleanup (&parts);
 
    reader = bson_reader_new_from_data (command->payload.data + data_offset,
                                        command->payload.len - data_offset);
@@ -773,36 +822,20 @@ _mongoc_write_opquery (mongoc_write_command_t *command,
          data_offset += len;
       }
    } else {
-      mongoc_cmd_parts_init (&parts, client, database, MONGOC_QUERY_NONE, &cmd);
-      parts.is_write_command = true;
-      parts.assembled.operation_id = command->operation_id;
-      if (!mongoc_cmd_parts_set_write_concern (
-             &parts,
-             write_concern,
-             server_stream->sd->max_wire_version,
-             error)) {
-         bson_reader_destroy (reader);
-         bson_destroy (&cmd);
-         mongoc_cmd_parts_cleanup (&parts);
-         EXIT;
-      }
+      bson_t reply;
 
-      BSON_ASSERT (bson_iter_init (&iter, &command->cmd_opts));
-      if (!mongoc_cmd_parts_append_opts (
-             &parts, &iter, server_stream->sd->max_wire_version, error)) {
-         bson_reader_destroy (reader);
-         bson_destroy (&cmd);
-         mongoc_cmd_parts_cleanup (&parts);
-         EXIT;
-      }
-
-      ret = mongoc_cmd_parts_assemble (&parts, server_stream, error);
+      ret = _assemble_cmd (&cmd,
+                           command,
+                           client,
+                           server_stream,
+                           database,
+                           write_concern,
+                           &parts,
+                           error);
       if (ret) {
          ret = mongoc_cluster_run_command_monitored (
             &client->cluster, &parts.assembled, &reply, error);
       } else {
-         /* assembling failed */
-         result->must_stop = true;
          bson_init (&reply);
       }
 

src/libmongoc/tests/mock_server/request.c

@@ -131,6 +131,7 @@ request_new (const mongoc_buffer_t *buffer,
 const bson_t *
 request_get_doc (const request_t *request, int n)
 {
+   BSON_ASSERT (request);
    return _mongoc_array_index (&request->docs, const bson_t *, n);
 }
 

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

@@ -3458,6 +3458,7 @@ test_client_reset_sessions (void)
    /* Add an autoresponder for endSessions to unblock the test. */
    mock_server_auto_endsessions (server);
 
+   bson_destroy (&opts);
    bson_destroy (&lsid);
    request_destroy (request);
    future_destroy (future);

src/libmongoc/tests/test-mongoc-write-commands.c

@@ -10,6 +10,9 @@
 
 #include "test-libmongoc.h"
 #include "test-conveniences.h"
+#include "mock_server/mock-server.h"
+#include "mock_server/future.h"
+#include "mock_server/future-functions.h"
 
 
 static void
@@ -369,6 +372,92 @@ test_bypass_not_sent ()
    mongoc_client_destroy (client);
 }
 
+static void
+test_split_opquery_with_options (void)
+{
+   mock_server_t *server;
+   mongoc_client_t *client;
+   mongoc_collection_t *coll;
+   bson_t **docs;
+   int i;
+   bson_error_t error;
+   future_t *future;
+   request_t *request;
+   const bson_t *insert;
+   bson_t opts;
+   mongoc_write_concern_t *wc;
+   int n_docs;
+
+   /* Use a reduced maxBsonObjectSize, and wire version for OP_QUERY */
+   const char *ismaster = "{'ok': 1.0,"
+                          " 'ismaster': true,"
+                          " 'minWireVersion': 0,"
+                          " 'maxWireVersion': 5,"
+                          " 'maxBsonObjectSize': 100}";
+
+   server = mock_server_new ();
+   mock_server_auto_ismaster (server, ismaster);
+   mock_server_run (server);
+
+   /* Create an insert with two batches. Because of the reduced
+   * maxBsonObjectSize, each document must be less than 100 bytes.
+   * Because of the hardcoded allowance (see SERVER-10643), and our current
+   * incorrect batching logic (see CDRIVER-3310) the complete insert
+   * command can be can be 16k + 100 bytes.
+   * After CDRIVER-3310, update this test, since the allowance will not be
+   * taken into consideration for document batching.
+   * So create enough documents to fill up at least one batch.
+   */
+   n_docs = ((BSON_OBJECT_ALLOWANCE) / tmp_bson ("{ '_id': 1 }")->len) +
+            1; /* inexact, but errs towards more than enough documents. */
+   docs = bson_malloc (sizeof (bson_t *) * n_docs);
+   for (i = 0; i < n_docs; i++) {
+      docs[i] = BCON_NEW ("_id", BCON_INT64 (i));
+   }
+
+   client = mongoc_client_new_from_uri (mock_server_get_uri (server));
+   coll = mongoc_client_get_collection (client, "db", "coll");
+
+   /* Add a write concern, to ensure that it is taken into account during
+    * splitting. */
+   bson_init (&opts);
+   wc = mongoc_write_concern_new ();
+   mongoc_write_concern_set_wmajority (wc, 100);
+   mongoc_write_concern_append (wc, &opts);
+
+   future = future_collection_insert_many (
+      coll, (const bson_t **) docs, n_docs, &opts, NULL, &error);
+   /* Mock server recieves first insert. */
+   request = mock_server_receives_request (server);
+   BSON_ASSERT (request);
+   insert = request_get_doc (request, 0);
+   /* The total command size is just a hair under BSON_OBJECT_ALLOWANCE (16384)
+    * + 100 */
+   BSON_ASSERT (insert->len == 16482);
+   mock_server_replies_ok_and_destroys (request);
+
+   /* Mock server recieves second insert. */
+   request = mock_server_receives_request (server);
+   BSON_ASSERT (request);
+   insert = request_get_doc (request, 0);
+   /* The size doesn't really matter for the purpose of the test, but check it
+    * anyway. */
+   BSON_ASSERT (insert->len == 10433);
+   mock_server_replies_ok_and_destroys (request);
+   BSON_ASSERT (future_get_bool (future));
+
+   future_destroy (future);
+   for (i = 0; i < n_docs; i++) {
+      bson_destroy (docs[i]);
+   }
+   bson_free (docs);
+   bson_destroy (&opts);
+   mongoc_collection_destroy (coll);
+   mongoc_write_concern_destroy (wc);
+   mongoc_client_destroy (client);
+   mock_server_destroy (server);
+}
+
 void
 test_write_command_install (TestSuite *suite)
 {
@@ -383,4 +472,7 @@ test_write_command_install (TestSuite *suite)
                       NULL,
                       NULL,
                       test_framework_skip_if_max_wire_version_less_than_4);
+   TestSuite_AddMockServerTest (suite,
+                                "/WriteCommand/split_opquery_with_options",
+                                test_split_opquery_with_options);
 }