refactor: build commands with connection

Author: durran

src/cursor/client_bulk_write_cursor.ts

@@ -1,8 +1,10 @@
-import type { Document } from '../bson';
+import { type Document } from 'bson';
+
 import { type ClientBulkWriteCursorResponse } from '../cmap/wire_protocol/responses';
 import { MongoClientBulkWriteCursorError } from '../error';
 import type { MongoClient } from '../mongo_client';
 import { ClientBulkWriteOperation } from '../operations/client_bulk_write/client_bulk_write';
+import { type ClientBulkWriteCommandBuilder } from '../operations/client_bulk_write/command_builder';
 import { type ClientBulkWriteOptions } from '../operations/client_bulk_write/common';
 import { executeOperation } from '../operations/execute_operation';
 import type { ClientSession } from '../sessions';
@@ -24,17 +26,21 @@ export interface ClientBulkWriteCursorOptions
  * @internal
  */
 export class ClientBulkWriteCursor extends AbstractCursor {
-  public readonly command: Document;
+  commandBuilder: ClientBulkWriteCommandBuilder;
   /** @internal */
   private cursorResponse?: ClientBulkWriteCursorResponse;
   /** @internal */
   private clientBulkWriteOptions: ClientBulkWriteOptions;
 
   /** @internal */
-  constructor(client: MongoClient, command: Document, options: ClientBulkWriteOptions = {}) {
+  constructor(
+    client: MongoClient,
+    commandBuilder: ClientBulkWriteCommandBuilder,
+    options: ClientBulkWriteOptions = {}
+  ) {
     super(client, new MongoDBNamespace('admin', '$cmd'), options);
 
-    this.command = command;
+    this.commandBuilder = commandBuilder;
     this.clientBulkWriteOptions = options;
   }
 
@@ -49,17 +55,24 @@ export class ClientBulkWriteCursor extends AbstractCursor {
     );
   }
 
+  /**
+   * Get the last set of operations the cursor executed.
+   */
+  get operations(): Document[] {
+    return this.commandBuilder.lastOperations;
+  }
+
   clone(): ClientBulkWriteCursor {
     const clonedOptions = mergeOptions({}, this.clientBulkWriteOptions);
     delete clonedOptions.session;
-    return new ClientBulkWriteCursor(this.client, this.command, {
+    return new ClientBulkWriteCursor(this.client, this.commandBuilder, {
       ...clonedOptions
     });
   }
 
   /** @internal */
   async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
-    const clientBulkWriteOperation = new ClientBulkWriteOperation(this.command, {
+    const clientBulkWriteOperation = new ClientBulkWriteOperation(this.commandBuilder, {
       ...this.clientBulkWriteOptions,
       ...this.cursorOptions,
       session

src/operations/client_bulk_write/client_bulk_write.ts

@@ -1,28 +1,28 @@
-import { type Document } from 'bson';
-
+import { ServerType } from '../../beta';
 import { ClientBulkWriteCursorResponse } from '../../cmap/wire_protocol/responses';
 import type { Server } from '../../sdam/server';
 import type { ClientSession } from '../../sessions';
 import { MongoDBNamespace } from '../../utils';
 import { CommandOperation } from '../command';
 import { Aspect, defineAspects } from '../operation';
+import { type ClientBulkWriteCommandBuilder } from './command_builder';
 import { type ClientBulkWriteOptions } from './common';
 
 /**
  * Executes a single client bulk write operation within a potential batch.
  * @internal
  */
 export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCursorResponse> {
-  command: Document;
+  commandBuilder: ClientBulkWriteCommandBuilder;
   override options: ClientBulkWriteOptions;
 
   override get commandName() {
     return 'bulkWrite' as const;
   }
 
-  constructor(command: Document, options: ClientBulkWriteOptions) {
+  constructor(commandBuilder: ClientBulkWriteCommandBuilder, options: ClientBulkWriteOptions) {
     super(undefined, options);
-    this.command = command;
+    this.commandBuilder = commandBuilder;
     this.options = options;
     this.ns = new MongoDBNamespace('admin', '$cmd');
   }
@@ -37,9 +37,34 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
     server: Server,
     session: ClientSession | undefined
   ): Promise<ClientBulkWriteCursorResponse> {
-    return await super.executeCommand(server, session, this.command, ClientBulkWriteCursorResponse);
+    let command;
+
+    if (server.description.type === ServerType.LoadBalancer) {
+      // Checkout a connection to build the command.
+      const connection = await server.pool.checkOut();
+      command = this.commandBuilder.buildBatch(
+        connection.hello?.maxMessageSizeBytes,
+        connection.hello?.maxWriteBatchSize
+      );
+      // Pin the connection to the session so it get used to execute the command and we do not
+      // perform a double check-in/check-out.
+      session?.pin(connection);
+    } else {
+      // At this point we have a server and the auto connect code has already
+      // run in executeOperation, so the server description will be populated.
+      // We can use that to build the command.
+      command = this.commandBuilder.buildBatch(
+        server.description.maxMessageSizeBytes,
+        server.description.maxWriteBatchSize
+      );
+    }
+    return await super.executeCommand(server, session, command, ClientBulkWriteCursorResponse);
   }
 }
 
 // Skipping the collation as it goes on the individual ops.
-defineAspects(ClientBulkWriteOperation, [Aspect.WRITE_OPERATION, Aspect.SKIP_COLLATION]);
+defineAspects(ClientBulkWriteOperation, [
+  Aspect.WRITE_OPERATION,
+  Aspect.SKIP_COLLATION,
+  Aspect.CURSOR_CREATING
+]);

src/operations/client_bulk_write/command_builder.ts

@@ -38,6 +38,8 @@ export class ClientBulkWriteCommandBuilder {
   models: AnyClientBulkWriteModel[];
   options: ClientBulkWriteOptions;
   pkFactory: PkFactory;
+  currentModelIndex: number;
+  lastOperations: Document[];
 
   /**
    * Create the command builder.
@@ -51,6 +53,8 @@ export class ClientBulkWriteCommandBuilder {
     this.models = models;
     this.options = options;
     this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;
+    this.currentModelIndex = 0;
+    this.lastOperations = [];
   }
 
   /**
@@ -65,68 +69,55 @@ export class ClientBulkWriteCommandBuilder {
   }
 
   /**
-   * Build the bulk write commands from the models.
+   * Determines if there is another batch to process.
+   * @returns True if not all batches have been built.
    */
-  buildCommands(maxMessageSizeBytes: number, maxWriteBatchSize: number): ClientBulkWriteCommand[] {
-    // Iterate the models to build the ops and nsInfo fields.
-    // We need to do this in a loop which creates one command each up
-    // to the max bson size or max message size.
-    const commands: ClientBulkWriteCommand[] = [];
-    let currentCommandLength = 0;
+  hasNextBatch(): boolean {
+    return this.currentModelIndex < this.models.length;
+  }
+
+  /**
+   * Build a single batch of a client bulk write command.
+   * @param maxMessageSizeBytes - The max message size in bytes.
+   * @param maxWriteBatchSize - The max write batch size.
+   * @returns The client bulk write command.
+   */
+  buildBatch(maxMessageSizeBytes: number, maxWriteBatchSize: number): ClientBulkWriteCommand {
+    let commandLength = 0;
     let currentNamespaceIndex = 0;
-    let currentCommand: ClientBulkWriteCommand = this.baseCommand();
+    const command: ClientBulkWriteCommand = this.baseCommand();
     const namespaces = new Map<string, number>();
 
-    for (const model of this.models) {
+    while (this.currentModelIndex < this.models.length) {
+      const model = this.models[this.currentModelIndex];
       const ns = model.namespace;
-      const index = namespaces.get(ns);
-
-      /**
-       * Convenience function for resetting everything when a new batch
-       * is started.
-       */
-      const reset = () => {
-        commands.push(currentCommand);
-        namespaces.clear();
-        currentNamespaceIndex = 0;
-        currentCommand = this.baseCommand();
-        namespaces.set(ns, currentNamespaceIndex);
-      };
+      const nsIndex = namespaces.get(ns);
 
-      if (index != null) {
-        // Pushing to the ops document sequence returns the bytes length added.
-        const operation = buildOperation(model, index, this.pkFactory);
+      if (nsIndex != null) {
+        // Build the operation and serialize it to get the bytes buffer.
+        const operation = buildOperation(model, nsIndex, this.pkFactory);
         const operationBuffer = BSON.serialize(operation);
 
-        // Check if the operation buffer can fit in the current command. If it can,
+        // Check if the operation buffer can fit in the command. If it can,
         // then add the operation to the document sequence and increment the
         // current length as long as the ops don't exceed the maxWriteBatchSize.
         if (
-          currentCommandLength + operationBuffer.length < maxMessageSizeBytes &&
-          currentCommand.ops.documents.length < maxWriteBatchSize
+          commandLength + operationBuffer.length < maxMessageSizeBytes &&
+          command.ops.documents.length < maxWriteBatchSize
         ) {
           // Pushing to the ops document sequence returns the total byte length of the document sequence.
-          currentCommandLength =
-            MESSAGE_OVERHEAD_BYTES + this.addOperation(currentCommand, operation, operationBuffer);
+          commandLength = MESSAGE_OVERHEAD_BYTES + command.ops.push(operation, operationBuffer);
+          // Increment the builder's current model index.
+          this.currentModelIndex++;
         } else {
-          // We need to batch. Push the current command to the commands
-          // array and create a new current command. We aslo need to clear the namespaces
-          // map for the new command.
-          reset();
-
-          const nsInfo = { ns: ns };
-          const nsInfoBuffer = BSON.serialize(nsInfo);
-          currentCommandLength =
-            MESSAGE_OVERHEAD_BYTES +
-            this.addOperationAndNsInfo(
-              currentCommand,
-              operation,
-              operationBuffer,
-              nsInfo,
-              nsInfoBuffer
-            );
+          // The operation cannot fit in the current command and will need to
+          // go in the next batch. Exit the loop and set the last ops.
+          this.lastOperations = command.ops.documents;
+          break;
         }
       } else {
+        // The namespace is not already in the nsInfo so we will set it in the map, and
+        // construct our nsInfo and ops documents and buffers.
         namespaces.set(ns, currentNamespaceIndex);
         const nsInfo = { ns: ns };
         const nsInfoBuffer = BSON.serialize(nsInfo);
@@ -138,68 +129,27 @@ export class ClientBulkWriteCommandBuilder {
         // sequences and increment the current length as long as the ops don't exceed
         // the maxWriteBatchSize.
         if (
-          currentCommandLength + nsInfoBuffer.length + operationBuffer.length <
-            maxMessageSizeBytes &&
-          currentCommand.ops.documents.length < maxWriteBatchSize
+          commandLength + nsInfoBuffer.length + operationBuffer.length < maxMessageSizeBytes &&
+          command.ops.documents.length < maxWriteBatchSize
         ) {
-          currentCommandLength =
+          // Pushing to the ops document sequence returns the total byte length of the document sequence.
+          commandLength =
             MESSAGE_OVERHEAD_BYTES +
-            this.addOperationAndNsInfo(
-              currentCommand,
-              operation,
-              operationBuffer,
-              nsInfo,
-              nsInfoBuffer
-            );
+            command.nsInfo.push(nsInfo, nsInfoBuffer) +
+            command.ops.push(operation, operationBuffer);
+          // We've added a new namespace, increment the namespace index.
+          currentNamespaceIndex++;
+          // Increment the builder's current model index.
+          this.currentModelIndex++;
         } else {
-          // We need to batch. Push the current command to the commands
-          // array and create a new current command. Aslo clear the namespaces map.
-          reset();
-
-          currentCommandLength =
-            MESSAGE_OVERHEAD_BYTES +
-            this.addOperationAndNsInfo(
-              currentCommand,
-              operation,
-              operationBuffer,
-              nsInfo,
-              nsInfoBuffer
-            );
+          // The operation cannot fit in the current command and will need to
+          // go in the next batch. Exit the loop and set the last ops.
+          this.lastOperations = command.ops.documents;
+          break;
         }
-        // We've added a new namespace, increment the namespace index.
-        currentNamespaceIndex++;
       }
     }
-
-    // After we've finisihed iterating all the models put the last current command
-    // only if there are operations in it.
-    if (currentCommand.ops.documents.length > 0) {
-      commands.push(currentCommand);
-    }
-
-    return commands;
-  }
-
-  private addOperation(
-    command: ClientBulkWriteCommand,
-    operation: Document,
-    operationBuffer: Uint8Array
-  ): number {
-    // Pushing to the ops document sequence returns the total byte length of the document sequence.
-    return command.ops.push(operation, operationBuffer);
-  }
-
-  private addOperationAndNsInfo(
-    command: ClientBulkWriteCommand,
-    operation: Document,
-    operationBuffer: Uint8Array,
-    nsInfo: Document,
-    nsInfoBuffer: Uint8Array
-  ): number {
-    // Pushing to the nsInfo document sequence returns the total byte length of the document sequence.
-    const nsInfoLength = command.nsInfo.push(nsInfo, nsInfoBuffer);
-    const opsLength = this.addOperation(command, operation, operationBuffer);
-    return nsInfoLength + opsLength;
+    return command;
   }
 
   private baseCommand(): ClientBulkWriteCommand {