Add csot support for bulk write

Author: baileympearson

src/cursor/abstract_cursor.ts

@@ -220,7 +220,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');
     }
     this.cursorOptions.omitMaxTimeMS =

src/cursor/client_bulk_write_cursor.ts

@@ -36,7 +36,7 @@ export class ClientBulkWriteCursor extends AbstractCursor {
   constructor(
     client: MongoClient,
     commandBuilder: ClientBulkWriteCommandBuilder,
-    options: ClientBulkWriteOptions = {}
+    options: ClientBulkWriteCursorOptions = {}
   ) {
     super(client, new MongoDBNamespace('admin', '$cmd'), options);
 
@@ -78,7 +78,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,5 +1,8 @@
+import { CursorTimeoutContext, CursorTimeoutMode } from '../../cursor/abstract_cursor';
 import { ClientBulkWriteCursor } from '../../cursor/client_bulk_write_cursor';
 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';
@@ -56,18 +59,27 @@ 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 { ok: 1 };
     } else {
       const resultsMerger = new ClientBulkWriteResultsMerger(this.options);
       // For each command will will create and exhaust a cursor for the results.
       let currentBatchOffset = 0;
       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);
         const docs = await cursor.toArray();
         const operations = cursor.operations;
         resultsMerger.merge(currentBatchOffset, operations, cursor.response, docs);

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/timeout.ts

@@ -171,7 +171,7 @@ function isCSOTTimeoutContextOptions(v: unknown): v is CSOTTimeoutContextOptions
 
 /** @internal */
 export abstract class TimeoutContext {
-  static create(options: TimeoutContextOptions): TimeoutContext {
+  static create(options: Partial<TimeoutContextOptions>): TimeoutContext {
     if (options.session?.timeoutContext != null) return options.session?.timeoutContext;
     if (isCSOTTimeoutContextOptions(options)) return new CSOTTimeoutContext(options);
     else if (isLegacyTimeoutContextOptions(options)) return new LegacyTimeoutContext(options);

src/utils.ts

@@ -35,6 +35,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';
 
 /**
@@ -514,6 +515,21 @@ 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
+): Pick<
+  MongoClient['s']['options'],
+  'serverSelectionTimeoutMS' | 'socketTimeoutMS' | 'waitQueueTimeoutMS' | 'timeoutMS'
+> &
+  T {
+  const { socketTimeoutMS, serverSelectionTimeoutMS, waitQueueTimeoutMS, timeoutMS } =
+    client.s.options;
+  return Object.assign(
+    { 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

@@ -1104,7 +1104,7 @@ describe('CSOT spec prose tests', function () {
     });
   });
 
-  describe.skip(
+  describe(
     '11. Multi-batch bulkWrites',
     { requires: { mongodb: '>=8.0', serverless: 'forbid' } },
     function () {
@@ -1184,7 +1184,7 @@ describe('CSOT spec prose tests', function () {
         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));
 
@@ -1199,7 +1199,7 @@ describe('CSOT spec prose tests', function () {
 
         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';
+      });
     }
   );
 });