feat(NODE-6389): add support for timeoutMS in StateMachine.execute()

etc/notes/errors.md

@@ -67,7 +67,7 @@ Children of `MongoError` include:
 ### `MongoDriverError`
 
 This class represents errors which originate in the driver itself or when the user incorrectly uses the driver. This class should **never** be directly instantiated.
-Its children are the main classes of errors that most users will interact with: [**`MongoAPIError`**](#MongoAPIError) and [**`MongoRuntimeError`**](#MongoRuntimeError).
+Its children are the main classes of errors that most users will interact with: [**`MongoAPIError`**](#MongoAPIError), [**`MongoRuntimeError`**](#MongoRuntimeError) and [**`MongoOperationTimeoutError`**](#MongoOperationTimeoutError).
 
 ### `MongoAPIError`
 
@@ -109,6 +109,10 @@ This class should **never** be directly instantiated.
 | **MongoGridFSChunkError**   | Thrown when a malformed or invalid chunk is encountered when reading from a GridFS Stream. |
 | **MongoUnexpectedServerResponseError**   | Thrown when the driver receives a **parsable** response it did not expect from the server. |
 
+### `MongoOperationTimeoutError`
+
+- TODO(NODE-5688): Add MongoOperationTimeoutError documentation
+
 ### MongoUnexpectedServerResponseError
 
 Intended for the scenario where the MongoDB returns an unexpected response in relation to some state the driver is in.

package.json

@@ -97,7 +97,7 @@
     "mocha": "^10.4.0",
     "mocha-sinon": "^2.1.2",
     "mongodb-client-encryption": "^6.1.0",
-    "mongodb-legacy": "^6.1.0",
+    "mongodb-legacy": "^6.1.1",
     "nyc": "^15.1.0",
     "prettier": "^3.3.3",
     "semver": "^7.6.3",

src/admin.ts

@@ -78,7 +78,8 @@ export class Admin {
       new RunAdminCommandOperation(command, {
         ...resolveBSONOptions(options),
         session: options?.session,
-        readPreference: options?.readPreference
+        readPreference: options?.readPreference,
+        timeoutMS: options?.timeoutMS ?? this.s.db.timeoutMS
       })
     );
   }
@@ -154,7 +155,10 @@ export class Admin {
    * @param options - Optional settings for the command
    */
   async listDatabases(options?: ListDatabasesOptions): Promise<ListDatabasesResult> {
-    return await executeOperation(this.s.db.client, new ListDatabasesOperation(this.s.db, options));
+    return await executeOperation(
+      this.s.db.client,
+      new ListDatabasesOperation(this.s.db, { timeoutMS: this.s.db.timeoutMS, ...options })
+    );
   }
 
   /**

src/bulk/common.ts

@@ -20,6 +20,7 @@ import { makeUpdateStatement, UpdateOperation, type UpdateStatement } from '../o
 import type { Server } from '../sdam/server';
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
+import { type TimeoutContext } from '../timeout';
 import {
   applyRetryableWrites,
   type Callback,
@@ -873,6 +874,9 @@ export interface BulkWriteOptions extends CommandOperationOptions {
   forceServerObjectId?: boolean;
   /** Map of parameter names and values that can be accessed using $$var (requires MongoDB 5.0). */
   let?: Document;
+
+  /** @internal */
+  timeoutContext?: TimeoutContext;
 }
 
 const executeCommandsAsync = promisify(executeCommands);

src/client-side-encryption/state_machine.ts

@@ -12,7 +12,9 @@ import {
 } from '../bson';
 import { type ProxyOptions } from '../cmap/connection';
 import { getSocks, type SocksLib } from '../deps';
+import { MongoOperationTimeoutError } from '../error';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
+import { type CSOTTimeoutContext, Timeout, TimeoutError } from '../timeout';
 import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
 import { autoSelectSocketOptions, type DataKey } from './client_encryption';
 import { MongoCryptError } from './errors';
@@ -182,7 +184,11 @@ export class StateMachine {
   /**
    * Executes the state machine according to the specification
    */
-  async execute(executor: StateMachineExecutable, context: MongoCryptContext): Promise<Uint8Array> {
+  async execute(
+    executor: StateMachineExecutable,
+    context: MongoCryptContext,
+    timeoutContext?: CSOTTimeoutContext
+  ): Promise<Uint8Array> {
     const keyVaultNamespace = executor._keyVaultNamespace;
     const keyVaultClient = executor._keyVaultClient;
     const metaDataClient = executor._metaDataClient;
@@ -201,8 +207,13 @@ export class StateMachine {
               'unreachable state machine state: entered MONGOCRYPT_CTX_NEED_MONGO_COLLINFO but metadata client is undefined'
             );
           }
-          const collInfo = await this.fetchCollectionInfo(metaDataClient, context.ns, filter);
 
+          const collInfo = await this.fetchCollectionInfo(
+            metaDataClient,
+            context.ns,
+            filter,
+            timeoutContext?.csotEnabled() ? timeoutContext.remainingTimeMS : null
+          );
           if (collInfo) {
             context.addMongoOperationResponse(collInfo);
           }
@@ -222,9 +233,20 @@ export class StateMachine {
           // When we are using the shared library, we don't have a mongocryptd manager.
           const markedCommand: Uint8Array = mongocryptdManager
             ? await mongocryptdManager.withRespawn(
-                this.markCommand.bind(this, mongocryptdClient, context.ns, command)
+                this.markCommand.bind(
+                  this,
+                  mongocryptdClient,
+                  context.ns,
+                  command,
+                  timeoutContext?.csotEnabled() ? timeoutContext.remainingTimeMS : null
+                )
               )
-            : await this.markCommand(mongocryptdClient, context.ns, command);
+            : await this.markCommand(
+                mongocryptdClient,
+                context.ns,
+                command,
+                timeoutContext?.csotEnabled() ? timeoutContext.remainingTimeMS : null
+              );
 
           context.addMongoOperationResponse(markedCommand);
           context.finishMongoOperation();
@@ -233,7 +255,12 @@ export class StateMachine {
 
         case MONGOCRYPT_CTX_NEED_MONGO_KEYS: {
           const filter = context.nextMongoOperation();
-          const keys = await this.fetchKeys(keyVaultClient, keyVaultNamespace, filter);
+          const keys = await this.fetchKeys(
+            keyVaultClient,
+            keyVaultNamespace,
+            filter,
+            timeoutContext?.csotEnabled() ? timeoutContext.remainingTimeMS : null
+          );
 
           if (keys.length === 0) {
             // See docs on EMPTY_V
@@ -255,9 +282,12 @@ export class StateMachine {
         }
 
         case MONGOCRYPT_CTX_NEED_KMS: {
-          const requests = Array.from(this.requests(context));
-          await Promise.all(requests);
-
+          await Promise.all(
+            this.requests(
+              context,
+              timeoutContext?.csotEnabled() ? timeoutContext.remainingTimeMS : null
+            )
+          );
           context.finishKMSRequests();
           break;
         }
@@ -299,7 +329,7 @@ export class StateMachine {
    * @param kmsContext - A C++ KMS context returned from the bindings
    * @returns A promise that resolves when the KMS reply has be fully parsed
    */
-  async kmsRequest(request: MongoCryptKMSRequest): Promise<void> {
+  async kmsRequest(request: MongoCryptKMSRequest, timeoutMS?: number | null): Promise<void> {
     const parsedUrl = request.endpoint.split(':');
     const port = parsedUrl[1] != null ? Number.parseInt(parsedUrl[1], 10) : HTTPS_PORT;
     const socketOptions = autoSelectSocketOptions(this.options.socketOptions || {});
@@ -329,10 +359,6 @@ export class StateMachine {
       }
     }
 
-    function ontimeout() {
-      return new MongoCryptError('KMS request timed out');
-    }
-
     function onerror(cause: Error) {
       return new MongoCryptError('KMS request failed', { cause });
     }
@@ -364,7 +390,6 @@ export class StateMachine {
       resolve: resolveOnNetSocketConnect
     } = promiseWithResolvers<void>();
     netSocket
-      .once('timeout', () => rejectOnNetSocketError(ontimeout()))
       .once('error', err => rejectOnNetSocketError(onerror(err)))
       .once('close', () => rejectOnNetSocketError(onclose()))
       .once('connect', () => resolveOnNetSocketConnect());
@@ -410,8 +435,8 @@ export class StateMachine {
         reject: rejectOnTlsSocketError,
         resolve
       } = promiseWithResolvers<void>();
+
       socket
-        .once('timeout', () => rejectOnTlsSocketError(ontimeout()))
         .once('error', err => rejectOnTlsSocketError(onerror(err)))
         .once('close', () => rejectOnTlsSocketError(onclose()))
         .on('data', data => {
@@ -425,20 +450,27 @@ export class StateMachine {
             resolve();
           }
         });
-      await willResolveKmsRequest;
+
+      await (typeof timeoutMS === 'number'
+        ? Promise.all([willResolveKmsRequest, Timeout.expires(timeoutMS)])
+        : willResolveKmsRequest);
+    } catch (error) {
+      if (error instanceof TimeoutError)
+        throw new MongoOperationTimeoutError('KMS request timed out');
+      throw error;
     } finally {
       // There's no need for any more activity on this socket at this point.
       destroySockets();
     }
   }
 
-  *requests(context: MongoCryptContext) {
+  *requests(context: MongoCryptContext, timeoutMS?: number | null) {
     for (
       let request = context.nextKMSRequest();
       request != null;
       request = context.nextKMSRequest()
     ) {
-      yield this.kmsRequest(request);
+      yield this.kmsRequest(request, timeoutMS);
     }
   }
 
@@ -498,15 +530,19 @@ export class StateMachine {
   async fetchCollectionInfo(
     client: MongoClient,
     ns: string,
-    filter: Document
+    filter: Document,
+    timeoutMS?: number | null
   ): Promise<Uint8Array | null> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
 
     const collections = await client
       .db(db)
       .listCollections(filter, {
         promoteLongs: false,
-        promoteValues: false
+        promoteValues: false,
+        ...(typeof timeoutMS === 'number'
+          ? { timeoutMS, timeoutMode: 'cursorLifetime' }
+          : undefined)
       })
       .toArray();
 
@@ -522,12 +558,20 @@ export class StateMachine {
    * @param command - The command to execute.
    * @param callback - Invoked with the serialized and marked bson command, or with an error
    */
-  async markCommand(client: MongoClient, ns: string, command: Uint8Array): Promise<Uint8Array> {
-    const options = { promoteLongs: false, promoteValues: false };
+  async markCommand(
+    client: MongoClient,
+    ns: string,
+    command: Uint8Array,
+    timeoutMS?: number | null
+  ): Promise<Uint8Array> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
-    const rawCommand = deserialize(command, options);
+    const bsonOptions = { promoteLongs: false, promoteValues: false };
+    const rawCommand = deserialize(command, bsonOptions);
 
-    const response = await client.db(db).command(rawCommand, options);
+    const response = await client.db(db).command(rawCommand, {
+      ...bsonOptions,
+      ...(typeof timeoutMS === 'number' ? { timeoutMS, omitMaxTimeMS: true } : undefined)
+    });
 
     return serialize(response, this.bsonOptions);
   }
@@ -543,15 +587,16 @@ export class StateMachine {
   fetchKeys(
     client: MongoClient,
     keyVaultNamespace: string,
-    filter: Uint8Array
+    filter: Uint8Array,
+    timeoutMS?: number | null
   ): Promise<Array<DataKey>> {
     const { db: dbName, collection: collectionName } =
       MongoDBCollectionNamespace.fromString(keyVaultNamespace);
 
     return client
       .db(dbName)
       .collection<DataKey>(collectionName, { readConcern: { level: 'majority' } })
-      .find(deserialize(filter))
+      .find(deserialize(filter), { timeoutMS: timeoutMS != null ? timeoutMS : undefined })
       .toArray();
   }
 }