fix: batch reset on retry

src/operations/client_bulk_write/client_bulk_write.ts

@@ -27,6 +27,10 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
     this.ns = new MongoDBNamespace('admin', '$cmd');
   }
 
+  override resetBatch(): boolean {
+    return this.commandBuilder.resetBatch();
+  }
+
   /**
    * Execute the command. Superclass will handle write concern, etc.
    * @param server - The server.
@@ -84,5 +88,6 @@ defineAspects(ClientBulkWriteOperation, [
   Aspect.WRITE_OPERATION,
   Aspect.SKIP_COLLATION,
   Aspect.CURSOR_CREATING,
-  Aspect.RETRYABLE
+  Aspect.RETRYABLE,
+  Aspect.COMMAND_BATCHING
 ]);

src/operations/client_bulk_write/command_builder.ts

@@ -40,6 +40,7 @@ export class ClientBulkWriteCommandBuilder {
   options: ClientBulkWriteOptions;
   pkFactory: PkFactory;
   currentModelIndex: number;
+  previousModelIndex: number;
   lastOperations: Document[];
 
   /**
@@ -55,6 +56,7 @@ export class ClientBulkWriteCommandBuilder {
     this.options = options;
     this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;
     this.currentModelIndex = 0;
+    this.previousModelIndex = 0;
     this.lastOperations = [];
   }
 
@@ -77,6 +79,15 @@ export class ClientBulkWriteCommandBuilder {
     return this.currentModelIndex < this.models.length;
   }
 
+  /**
+   * When we need to retry a command we need to set the current
+   * model index back to its previous value.
+   */
+  resetBatch(): boolean {
+    this.currentModelIndex = this.previousModelIndex;
+    return true;
+  }
+
   /**
    * Build a single batch of a client bulk write command.
    * @param maxMessageSizeBytes - The max message size in bytes.
@@ -92,6 +103,8 @@ export class ClientBulkWriteCommandBuilder {
     let currentNamespaceIndex = 0;
     const command: ClientBulkWriteCommand = this.baseCommand();
     const namespaces = new Map<string, number>();
+    // In the case of retries we need to mark where we started this batch.
+    this.previousModelIndex = this.currentModelIndex;
 
     while (this.currentModelIndex < this.models.length) {
       const model = this.models[this.currentModelIndex];

src/operations/client_bulk_write/common.ts

@@ -181,6 +181,7 @@ export interface ClientBulkWriteResult {
   deleteResults?: Map<number, ClientDeleteResult>;
 }
 
+/** @public */
 export interface ClientBulkWriteError {
   code: number;
   message: string;

src/operations/execute_operation.ts

@@ -260,6 +260,10 @@ async function tryOperation<
     }
 
     try {
+      // If tries > 0 and we are command batching we need to reset the batch.
+      if (tries > 0 && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
+        operation.resetBatch();
+      }
       return await operation.execute(server, session);
     } catch (operationError) {
       if (!(operationError instanceof MongoError)) throw operationError;

src/operations/operation.ts

@@ -11,7 +11,8 @@ export const Aspect = {
   EXPLAINABLE: Symbol('EXPLAINABLE'),
   SKIP_COLLATION: Symbol('SKIP_COLLATION'),
   CURSOR_CREATING: Symbol('CURSOR_CREATING'),
-  MUST_SELECT_SAME_SERVER: Symbol('MUST_SELECT_SAME_SERVER')
+  MUST_SELECT_SAME_SERVER: Symbol('MUST_SELECT_SAME_SERVER'),
+  COMMAND_BATCHING: Symbol('COMMAND_BATCHING')
 } as const;
 
 /** @public */
@@ -98,6 +99,10 @@ export abstract class AbstractOperation<TResult = any> {
     this[kSession] = undefined;
   }
 
+  resetBatch(): boolean {
+    return true;
+  }
+
   get canRetryRead(): boolean {
     return this.hasAspect(Aspect.RETRYABLE) && this.hasAspect(Aspect.READ_OPERATION);
   }

test/unit/index.test.ts

@@ -27,7 +27,6 @@ const EXPECTED_EXPORTS = [
   'CancellationToken',
   'ChangeStream',
   'ChangeStreamCursor',
-  'ClientBulkWriteError',
   'ClientEncryption',
   'ClientSession',
   'Code',