chore: create and pass around timeout contexts

Author: nbbeeken

package-lock.json

@@ -99,7 +99,7 @@
     "mocha": "^10.4.0",
     "mocha-sinon": "^2.1.2",
     "mongodb-client-encryption": "^6.1.0-alpha.0",
-    "mongodb-legacy": "^6.0.1",
+    "mongodb-legacy": "^6.1.1",
     "nyc": "^15.1.0",
     "prettier": "^2.8.8",
     "semver": "^7.6.0",

package.json

@@ -765,9 +765,22 @@ export class MongoUnexpectedServerResponseError extends MongoRuntimeError {
  * @internal
  */
 export class MongoOperationTimeoutError extends MongoRuntimeError {
+  get [Symbol.toStringTag]() {
+    return 'MongoOperationTimeoutError';
+  }
+
   override get name(): string {
     return 'MongoOperationTimeoutError';
   }
+
+  static is(error: unknown): error is MongoOperationTimeoutError {
+    return (
+      error != null &&
+      typeof error === 'object' &&
+      Symbol.toStringTag in error &&
+      error[Symbol.toStringTag] === 'MongoOperationTimeoutError'
+    );
+  }
 }
 
 /**

src/error.ts

@@ -82,11 +82,6 @@ export async function executeOperation<
   } else if (session.client !== client) {
     throw new MongoInvalidArgumentError('ClientSession must be from the same MongoClient');
   }
-  if (session.explicit && session?.timeoutMS != null && operation.options.timeoutMS != null) {
-    throw new MongoInvalidArgumentError(
-      'Do not specify timeoutMS on operation if already specified on an explicit session'
-    );
-  }
 
   const readPreference = operation.readPreference ?? ReadPreference.primary;
   const inTransaction = !!session?.inTransaction();
@@ -107,11 +102,13 @@ export async function executeOperation<
     session.unpin();
   }
 
-  timeoutContext ??= TimeoutContext.create({
-    serverSelectionTimeoutMS: client.s.options.serverSelectionTimeoutMS,
-    waitQueueTimeoutMS: client.s.options.waitQueueTimeoutMS,
-    timeoutMS: operation.options.timeoutMS
-  });
+  timeoutContext ??=
+    session.timeoutContext ??
+    TimeoutContext.create({
+      serverSelectionTimeoutMS: client.s.options.serverSelectionTimeoutMS,
+      waitQueueTimeoutMS: client.s.options.waitQueueTimeoutMS,
+      timeoutMS: operation.options.timeoutMS ?? session.timeoutMS
+    });
 
   try {
     return await tryOperation(operation, {

src/operations/execute_operation.ts

@@ -16,6 +16,7 @@ import {
   MongoErrorLabel,
   MongoExpiredSessionError,
   MongoInvalidArgumentError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError,
   MongoTransactionError,
@@ -29,6 +30,7 @@ import { ReadConcernLevel } from './read_concern';
 import { ReadPreference } from './read_preference';
 import { type AsyncDisposable, configureResourceManagement } from './resource_management';
 import { _advanceClusterTime, type ClusterTime, TopologyType } from './sdam/common';
+import { TimeoutContext } from './timeout';
 import {
   isTransactionCommand,
   Transaction,
@@ -101,6 +103,9 @@ export interface EndSessionOptions {
   error?: AnyError;
   force?: boolean;
   forceClear?: boolean;
+
+  /** @internal */
+  timeoutMS?: number;
 }
 
 /**
@@ -118,7 +123,7 @@ export class ClientSession
   /** @internal */
   sessionPool: ServerSessionPool;
   hasEnded: boolean;
-  clientOptions?: MongoOptions;
+  clientOptions: MongoOptions;
   supports: { causalConsistency: boolean };
   clusterTime?: ClusterTime;
   operationTime?: Timestamp;
@@ -140,6 +145,9 @@ export class ClientSession
   /** @internal */
   timeoutMS?: number;
 
+  /** @internal */
+  public timeoutContext: TimeoutContext | null = null;
+
   /**
    * Create a client session.
    * @internal
@@ -152,7 +160,7 @@ export class ClientSession
     client: MongoClient,
     sessionPool: ServerSessionPool,
     options: ClientSessionOptions,
-    clientOptions?: MongoOptions
+    clientOptions: MongoOptions
   ) {
     super();
 
@@ -272,7 +280,11 @@ export class ClientSession
   async endSession(options?: EndSessionOptions): Promise<void> {
     try {
       if (this.inTransaction()) {
-        await this.abortTransaction();
+        if (typeof options?.timeoutMS === 'number') {
+          await this.abortTransaction({ timeoutMS: options.timeoutMS });
+        } else {
+          await this.abortTransaction();
+        }
       }
       if (!this.hasEnded) {
         const serverSession = this[kServerSession];
@@ -291,6 +303,7 @@ export class ClientSession
       }
     } catch (error) {
       // spec indicates that we should ignore all errors for `endSessions`
+      if (MongoOperationTimeoutError.is(error)) throw error;
       squashError(error);
     } finally {
       maybeClearPinnedConnection(this, { force: true, ...options });
@@ -444,16 +457,20 @@ export class ClientSession
 
   /**
    * Commits the currently active transaction in this session.
+   *
+   * @param options - Optional options, can be used to override `defaultTimeoutMS`.
    */
-  async commitTransaction(): Promise<void> {
-    return await endTransaction(this, 'commitTransaction');
+  async commitTransaction(options?: { timeoutMS: number }): Promise<void> {
+    return await endTransaction(this, 'commitTransaction', options);
   }
 
   /**
    * Aborts the currently active transaction in this session.
+   *
+   * @param options - Optional options, can be used to override `defaultTimeoutMS`.
    */
-  async abortTransaction(): Promise<void> {
-    return await endTransaction(this, 'abortTransaction');
+  async abortTransaction(options?: { timeoutMS: number }): Promise<void> {
+    return await endTransaction(this, 'abortTransaction', options);
   }
 
   /**
@@ -499,7 +516,15 @@ export class ClientSession
     fn: WithTransactionCallback<T>,
     options?: TransactionOptions
   ): Promise<T> {
-    const startTime = now();
+    if (typeof options?.timeoutMS === 'number')
+      this.timeoutContext = TimeoutContext.create({
+        timeoutMS: options.timeoutMS,
+        serverSelectionTimeoutMS: this.clientOptions.serverSelectionTimeoutMS,
+        socketTimeoutMS: this.clientOptions.socketTimeoutMS
+      });
+    const { timeoutContext } = this;
+
+    const startTime = timeoutContext?.csotEnabled() ? timeoutContext.start : now();
     return await attemptTransaction(this, startTime, fn, options);
   }
 }
@@ -677,7 +702,8 @@ async function attemptTransaction<T>(
 
 async function endTransaction(
   session: ClientSession,
-  commandName: 'abortTransaction' | 'commitTransaction'
+  commandName: 'abortTransaction' | 'commitTransaction',
+  options: { timeoutMS?: number } = {}
 ): Promise<void> {
   // handle any initial problematic cases
   const txnState = session.transaction.state;
@@ -749,6 +775,25 @@ async function endTransaction(
     command.recoveryToken = session.transaction.recoveryToken;
   }
 
+  const timeoutMS =
+    'timeoutMS' in options && typeof options.timeoutMS === 'number'
+      ? options.timeoutMS
+      : typeof session.timeoutMS === 'number'
+      ? session.timeoutMS
+      : session.timeoutContext?.csotEnabled()
+      ? session.timeoutContext.timeoutMS
+      : null;
+
+  const timeoutContext =
+    // override for this operation
+    TimeoutContext.create({
+      serverSelectionTimeoutMS: session.clientOptions.serverSelectionTimeoutMS,
+      socketTimeoutMS: session.clientOptions.socketTimeoutMS,
+      ...(timeoutMS != null
+        ? { timeoutMS }
+        : { waitQueueTimeoutMS: session.clientOptions.waitQueueTimeoutMS })
+    });
+
   try {
     // send the command
     await executeOperation(
@@ -757,7 +802,8 @@ async function endTransaction(
         session,
         readPreference: ReadPreference.primary,
         bypassPinningCheck: true
-      })
+      }),
+      timeoutContext
     );
     if (command.abortTransaction) {
       // always unpin on abort regardless of command outcome
@@ -794,7 +840,8 @@ async function endTransaction(
             session,
             readPreference: ReadPreference.primary,
             bypassPinningCheck: true
-          })
+          }),
+          timeoutContext
         );
         if (commandName !== 'commitTransaction') {
           session.transaction.transition(TxnState.TRANSACTION_ABORTED);
@@ -824,6 +871,7 @@ function handleEndTransactionError(
       maybeClearPinnedConnection(session, { force: false });
     }
     // The spec indicates that if the operation times out or fails with a non-retryable error, we should ignore all errors on `abortTransaction`
+    if (MongoOperationTimeoutError.is(error)) throw error; // But not if it is CSOT
     return;
   }
 

src/sessions.ts

@@ -184,7 +184,7 @@ export class CSOTTimeoutContext extends TimeoutContext {
   private _serverSelectionTimeout?: Timeout | null;
   private _connectionCheckoutTimeout?: Timeout | null;
   public minRoundTripTime = 0;
-  private start: number;
+  public start: number;
 
   constructor(options: CSOTTimeoutContextOptions) {
     super();

src/timeout.ts

@@ -528,6 +528,12 @@ export function resolveOptions<T extends CommandOperationOptions>(
     result.readPreference = readPreference;
   }
 
+  if (session?.explicit && session.timeoutMS != null && options?.timeoutMS != null) {
+    throw new MongoInvalidArgumentError(
+      'Do not specify timeoutMS on operation if already specified on an explicit session'
+    );
+  }
+
   const timeoutMS = options?.timeoutMS;
 
   result.timeoutMS = timeoutMS ?? parent?.timeoutMS;