feat(NODE-6403): add CSOT support to client bulk write (#4261)

Co-authored-by: Warren James <[email protected]>

src/cmap/connection.ts

@@ -721,6 +721,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
           throw new MongoOperationTimeoutError('Timed out at socket write');
         }
         throw error;
+      } finally {
+        timeout.clear();
       }
     }
     return await drainEvent;

src/cmap/wire_protocol/on_data.ts

@@ -116,6 +116,7 @@ export function onData(
     emitter.off('data', eventHandler);
     emitter.off('error', errorHandler);
     finished = true;
+    timeoutForSocketRead?.clear();
     const doneResult = { value: undefined, done: finished } as const;
 
     for (const promise of unconsumedPromises) {

src/cursor/abstract_cursor.ts

@@ -243,7 +243,7 @@ export abstract class AbstractCursor<
         options.timeoutMode ??
         (options.tailable ? CursorTimeoutMode.ITERATION : CursorTimeoutMode.LIFETIME);
     } else {
-      if (options.timeoutMode != null)
+      if (options.timeoutMode != null && options.timeoutContext == null)
         throw new MongoInvalidArgumentError('Cannot set timeoutMode without setting timeoutMS');
     }
 

src/cursor/client_bulk_write_cursor.ts

@@ -34,7 +34,7 @@ export class ClientBulkWriteCursor extends AbstractCursor {
   constructor(
     client: MongoClient,
     commandBuilder: ClientBulkWriteCommandBuilder,
-    options: ClientBulkWriteOptions = {}
+    options: ClientBulkWriteCursorOptions = {}
   ) {
     super(client, new MongoDBNamespace('admin', '$cmd'), options);
 
@@ -71,7 +71,11 @@ export class ClientBulkWriteCursor extends AbstractCursor {
       session
     });
 
-    const response = await executeOperation(this.client, clientBulkWriteOperation);
+    const response = await executeOperation(
+      this.client,
+      clientBulkWriteOperation,
+      this.timeoutContext
+    );
     this.cursorResponse = response;
 
     return { server: clientBulkWriteOperation.server, session, response };

src/operations/client_bulk_write/executor.ts

@@ -1,4 +1,5 @@
 import { type Document } from '../../bson';
+import { CursorTimeoutContext, CursorTimeoutMode } from '../../cursor/abstract_cursor';
 import { ClientBulkWriteCursor } from '../../cursor/client_bulk_write_cursor';
 import {
   MongoClientBulkWriteError,
@@ -7,6 +8,8 @@ import {
   MongoServerError
 } from '../../error';
 import { type MongoClient } from '../../mongo_client';
+import { TimeoutContext } from '../../timeout';
+import { resolveTimeoutOptions } from '../../utils';
 import { WriteConcern } from '../../write_concern';
 import { executeOperation } from '../execute_operation';
 import { ClientBulkWriteOperation } from './client_bulk_write';
@@ -86,17 +89,26 @@ export class ClientBulkWriteExecutor {
       pkFactory
     );
     // Unacknowledged writes need to execute all batches and return { ok: 1}
+    const resolvedOptions = resolveTimeoutOptions(this.client, this.options);
+    const context = TimeoutContext.create(resolvedOptions);
+
     if (this.options.writeConcern?.w === 0) {
       while (commandBuilder.hasNextBatch()) {
         const operation = new ClientBulkWriteOperation(commandBuilder, this.options);
-        await executeOperation(this.client, operation);
+        await executeOperation(this.client, operation, context);
       }
       return ClientBulkWriteResultsMerger.unacknowledged();
     } else {
       const resultsMerger = new ClientBulkWriteResultsMerger(this.options);
       // For each command will will create and exhaust a cursor for the results.
       while (commandBuilder.hasNextBatch()) {
-        const cursor = new ClientBulkWriteCursor(this.client, commandBuilder, this.options);
+        const cursorContext = new CursorTimeoutContext(context, Symbol());
+        const options = {
+          ...this.options,
+          timeoutContext: cursorContext,
+          ...(resolvedOptions.timeoutMS != null && { timeoutMode: CursorTimeoutMode.LIFETIME })
+        };
+        const cursor = new ClientBulkWriteCursor(this.client, commandBuilder, options);
         try {
           await resultsMerger.merge(cursor);
         } catch (error) {

src/sdam/server.ts

@@ -106,7 +106,7 @@ export type ServerEvents = {
   EventEmitterWithState;
 
 /** @internal */
-export type ServerCommandOptions = Omit<CommandOptions, 'timeoutContext'> & {
+export type ServerCommandOptions = Omit<CommandOptions, 'timeoutContext' | 'socketTimeoutMS'> & {
   timeoutContext: TimeoutContext;
 };
 

src/utils.ts

@@ -36,6 +36,7 @@ import { ServerType } from './sdam/common';
 import type { Server } from './sdam/server';
 import type { Topology } from './sdam/topology';
 import type { ClientSession } from './sessions';
+import { type TimeoutContextOptions } from './timeout';
 import { WriteConcern } from './write_concern';
 
 /**
@@ -515,6 +516,18 @@ export function hasAtomicOperators(doc: Document | Document[]): boolean {
   return keys.length > 0 && keys[0][0] === '$';
 }
 
+export function resolveTimeoutOptions<T extends Partial<TimeoutContextOptions>>(
+  client: MongoClient,
+  options: T
+): T &
+  Pick<
+    MongoClient['s']['options'],
+    'timeoutMS' | 'serverSelectionTimeoutMS' | 'waitQueueTimeoutMS' | 'socketTimeoutMS'
+  > {
+  const { socketTimeoutMS, serverSelectionTimeoutMS, waitQueueTimeoutMS, timeoutMS } =
+    client.s.options;
+  return { socketTimeoutMS, serverSelectionTimeoutMS, waitQueueTimeoutMS, timeoutMS, ...options };
+}
 /**
  * Merge inherited properties from parent into options, prioritizing values from options,
  * then values from parent.

test/integration/client-side-operations-timeout/client_side_operations_timeout.prose.test.ts

@@ -21,7 +21,8 @@ import {
   promiseWithResolvers,
   squashError
 } from '../../mongodb';
-import { type FailPoint } from '../../tools/utils';
+import { type FailPoint, makeMultiBatchWrite } from '../../tools/utils';
+import { filterForCommands } from '../shared';
 
 // TODO(NODE-5824): Implement CSOT prose tests
 describe('CSOT spec prose tests', function () {
@@ -1183,9 +1184,9 @@ describe('CSOT spec prose tests', function () {
     });
   });
 
-  describe.skip(
+  describe(
     '11. Multi-batch bulkWrites',
-    { requires: { mongodb: '>=8.0', serverless: 'forbid' } },
+    { requires: { mongodb: '>=8.0', serverless: 'forbid', topology: 'single' } },
     function () {
       /**
        * ### 11. Multi-batch bulkWrites
@@ -1245,9 +1246,6 @@ describe('CSOT spec prose tests', function () {
         }
       };
 
-      let maxBsonObjectSize: number;
-      let maxMessageSizeBytes: number;
-
       beforeEach(async function () {
         await internalClient
           .db('db')
@@ -1256,29 +1254,20 @@ describe('CSOT spec prose tests', function () {
           .catch(() => null);
         await internalClient.db('admin').command(failpoint);
 
-        const hello = await internalClient.db('admin').command({ hello: 1 });
-        maxBsonObjectSize = hello.maxBsonObjectSize;
-        maxMessageSizeBytes = hello.maxMessageSizeBytes;
-
         client = this.configuration.newClient({ timeoutMS: 2000, monitorCommands: true });
       });
 
-      it.skip('performs two bulkWrites which fail to complete before 2000 ms', async function () {
+      it('performs two bulkWrites which fail to complete before 2000 ms', async function () {
         const writes = [];
-        client.on('commandStarted', ev => writes.push(ev));
+        client.on('commandStarted', filterForCommands('bulkWrite', writes));
 
-        const length = maxMessageSizeBytes / maxBsonObjectSize + 1;
-        const models = Array.from({ length }, () => ({
-          namespace: 'db.coll',
-          name: 'insertOne' as const,
-          document: { a: 'b'.repeat(maxBsonObjectSize - 500) }
-        }));
+        const models = await makeMultiBatchWrite(this.configuration);
 
         const error = await client.bulkWrite(models).catch(error => error);
 
         expect(error, error.stack).to.be.instanceOf(MongoOperationTimeoutError);
-        expect(writes.map(ev => ev.commandName)).to.deep.equal(['bulkWrite', 'bulkWrite']);
-      }).skipReason = 'TODO(NODE-6403): client.bulkWrite is implemented in a follow up';
+        expect(writes).to.have.lengthOf(2);
+      });
     }
   );
 });

test/integration/client-side-operations-timeout/node_csot.test.ts

@@ -279,12 +279,16 @@ describe('CSOT driver tests', metadata, () => {
           .stub(Connection.prototype, 'readMany')
           .callsFake(async function* (...args) {
             const realIterator = readManyStub.wrappedMethod.call(this, ...args);
-            const cmd = commandSpy.lastCall.args.at(1);
-            if ('giveMeWriteErrors' in cmd) {
-              await realIterator.next().catch(() => null); // dismiss response
-              yield { parse: () => writeErrorsReply };
-            } else {
-              yield (await realIterator.next()).value;
+            try {
+              const cmd = commandSpy.lastCall.args.at(1);
+              if ('giveMeWriteErrors' in cmd) {
+                await realIterator.next().catch(() => null); // dismiss response
+                yield { parse: () => writeErrorsReply };
+              } else {
+                yield (await realIterator.next()).value;
+              }
+            } finally {
+              realIterator.return();
             }
           });
       });

test/integration/collection-management/collection_db_management.test.ts

@@ -1,6 +1,6 @@
 import { expect } from 'chai';
 
-import { Collection, type Db, type MongoClient } from '../../mongodb';
+import { Collection, type Db, type MongoClient, ObjectId } from '../../mongodb';
 
 describe('Collection Management and Db Management', function () {
   let client: MongoClient;
@@ -16,7 +16,7 @@ describe('Collection Management and Db Management', function () {
   });
 
   it('returns a collection object after calling createCollection', async function () {
-    const collection = await db.createCollection('collection');
+    const collection = await db.createCollection(new ObjectId().toHexString());
     expect(collection).to.be.instanceOf(Collection);
   });