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 { type 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,6 +457,8 @@ 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> {
     if (this.transaction.state === TxnState.NO_TRANSACTION) {
@@ -538,6 +553,8 @@ export class ClientSession
 
   /**
    * Aborts the currently active transaction in this session.
+   *
+   * @param options - Optional options, can be used to override `defaultTimeoutMS`.
    */
   async abortTransaction(): Promise<void> {
     if (this.transaction.state === TxnState.NO_TRANSACTION) {

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;