refactor(NODE-6187): refactor to use TimeoutContext abstraction

src/cmap/connection.ts

@@ -29,7 +29,7 @@ import { type CancellationToken, TypedEventEmitter } from '../mongo_types';
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import { ServerType } from '../sdam/common';
 import { applySession, type ClientSession, updateSessionFromResponse } from '../sessions';
-import { type Timeout } from '../timeout';
+import { type TimeoutContext } from '../timeout';
 import {
   BufferPool,
   calculateDurationInMs,
@@ -95,7 +95,7 @@ export interface CommandOptions extends BSONSerializeOptions {
   directConnection?: boolean;
 
   /** @internal */
-  timeout?: Timeout;
+  timeoutContext?: TimeoutContext;
 }
 
 /** @public */

src/cmap/connection_pool.ts

@@ -27,8 +27,8 @@ import {
 } from '../error';
 import { CancellationToken, TypedEventEmitter } from '../mongo_types';
 import type { Server } from '../sdam/server';
-import { Timeout, TimeoutError } from '../timeout';
-import { type Callback, csotMin, List, makeCounter, promiseWithResolvers } from '../utils';
+import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';
+import { type Callback, List, makeCounter, promiseWithResolvers } from '../utils';
 import { connect } from './connect';
 import { Connection, type ConnectionEvents, type ConnectionOptions } from './connection';
 import {
@@ -354,41 +354,17 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
    * will be held by the pool. This means that if a connection is checked out it MUST be checked back in or
    * explicitly destroyed by the new owner.
    */
-  async checkOut(options?: { timeout?: Timeout }): Promise<Connection> {
+  async checkOut(options?: { timeoutContext?: TimeoutContext }): Promise<Connection> {
     this.emitAndLog(
       ConnectionPool.CONNECTION_CHECK_OUT_STARTED,
       new ConnectionCheckOutStartedEvent(this)
     );
 
-    const waitQueueTimeoutMS = this.options.waitQueueTimeoutMS;
-    const serverSelectionTimeoutMS = this[kServer].topology.s.serverSelectionTimeoutMS;
-
     const { promise, resolve, reject } = promiseWithResolvers<Connection>();
 
-    let timeout: Timeout | null = null;
-    if (options?.timeout) {
-      // CSOT enabled
-      // Determine if we're using the timeout passed in or a new timeout
-      if (options.timeout.duration > 0 || serverSelectionTimeoutMS > 0) {
-        // This check determines whether or not Topology.selectServer used the configured
-        // `timeoutMS` or `serverSelectionTimeoutMS` value for its timeout
-        if (
-          options.timeout.duration === serverSelectionTimeoutMS ||
-          csotMin(options.timeout.duration, serverSelectionTimeoutMS) < serverSelectionTimeoutMS
-        ) {
-          // server selection used `timeoutMS`, so we should use the existing timeout as the timeout
-          // here
-          timeout = options.timeout;
-        } else {
-          // server selection used `serverSelectionTimeoutMS`, so we construct a new timeout with
-          // the time remaining to ensure that Topology.selectServer and ConnectionPool.checkOut
-          // cumulatively don't spend more than `serverSelectionTimeoutMS` blocking
-          timeout = Timeout.expires(serverSelectionTimeoutMS - options.timeout.timeElapsed);
-        }
-      }
-    } else {
-      timeout = Timeout.expires(waitQueueTimeoutMS);
-    }
+    let timeout;
+    if (options?.timeoutContext) timeout = options.timeoutContext.connectionCheckoutTimeout;
+    else timeout = Timeout.expires(this.options.waitQueueTimeoutMS);
 
     const waitQueueMember: WaitQueueMember = {
       resolve,
@@ -403,6 +379,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       return await (timeout ? Promise.race([promise, timeout]) : promise);
     } catch (error) {
       if (TimeoutError.is(error)) {
+        timeout?.clear();
         waitQueueMember[kCancelled] = true;
 
         this.emitAndLog(
@@ -415,7 +392,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
             : 'Timed out while checking out a connection from connection pool',
           this.address
         );
-        if (options?.timeout) {
+        if (options?.timeoutContext?.csotEnabled()) {
           throw new MongoOperationTimeoutError('Timed out during connection checkout', {
             cause: timeoutError
           });
@@ -424,7 +401,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       }
       throw error;
     } finally {
-      if (timeout !== options?.timeout) timeout?.clear();
+      if (options?.timeoutContext?.clearConnectionCheckoutTimeout) timeout?.clear();
     }
   }
 

src/index.ts

@@ -557,7 +557,13 @@ export type {
   WithTransactionCallback
 } from './sessions';
 export type { Sort, SortDirection, SortDirectionForCmd, SortForCmd } from './sort';
-export type { Timeout } from './timeout';
+export type {
+  CSOTTimeoutContext,
+  LegacyTimeoutContext,
+  Timeout,
+  TimeoutContext,
+  TimeoutContextOptions
+} from './timeout';
 export type { Transaction, TransactionOptions, TxnState } from './transactions';
 export type {
   BufferPool,

src/operations/command.ts

@@ -118,7 +118,7 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
     const options = {
       ...this.options,
       ...this.bsonOptions,
-      timeout: this.timeout,
+      timeoutContext: this.timeoutContext,
       readPreference: this.readPreference,
       session
     };

src/operations/execute_operation.ts

@@ -27,7 +27,7 @@ import {
 } from '../sdam/server_selection';
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
-import { Timeout } from '../timeout';
+import { TimeoutContext } from '../timeout';
 import { squashError, supportsRetryableWrites } from '../utils';
 import { AbstractOperation, Aspect } from './operation';
 
@@ -118,6 +118,14 @@ export async function executeOperation<
     );
   }
 
+  const timeoutContext = TimeoutContext.create({
+    serverSelectionTimeoutMS: client.options.serverSelectionTimeoutMS,
+    waitQueueTimeoutMS: client.options.waitQueueTimeoutMS,
+    timeoutMS: operation.options.timeoutMS
+  });
+
+  operation.timeoutContext = timeoutContext;
+
   const readPreference = operation.readPreference ?? ReadPreference.primary;
   const inTransaction = !!session?.inTransaction();
 
@@ -153,13 +161,10 @@ export async function executeOperation<
     selector = readPreference;
   }
 
-  const timeout = operation.timeoutMS != null ? Timeout.expires(operation.timeoutMS) : undefined;
-  operation.timeout = timeout;
-
   const server = await topology.selectServer(selector, {
     session,
     operationName: operation.commandName,
-    timeout
+    timeoutContext: operation.timeoutContext
   });
 
   if (session == null) {
@@ -270,9 +275,9 @@ async function retryOperation<
   // select a new server, and attempt to retry the operation
   const server = await topology.selectServer(selector, {
     session,
-    timeout: operation.timeout,
     operationName: operation.commandName,
-    previousServer
+    previousServer,
+    timeoutContext: operation.timeoutContext
   });
 
   if (isWriteOperation && !supportsRetryableWrites(server)) {

src/operations/find.ts

@@ -113,7 +113,7 @@ export class FindOperation extends CommandOperation<Document> {
         ...this.bsonOptions,
         documentsReturnedIn: 'firstBatch',
         session,
-        timeout: this.timeout
+        timeoutContext: this.timeoutContext
       },
       undefined
     );

src/operations/get_more.ts

@@ -96,6 +96,7 @@ export class GetMoreOperation extends AbstractOperation {
     const commandOptions = {
       returnFieldSelector: null,
       documentsReturnedIn: 'nextBatch',
+      timeoutContext: this.timeoutContext,
       ...this.options
     };
 

src/operations/kill_cursors.ts

@@ -46,7 +46,10 @@ export class KillCursorsOperation extends AbstractOperation {
       cursors: [this.cursorId]
     };
     try {
-      await server.command(this.ns, killCursorsCommand, { session });
+      await server.command(this.ns, killCursorsCommand, {
+        session,
+        timeoutContext: this.timeoutContext
+      });
     } catch (error) {
       // The driver should never emit errors from killCursors, this is spec-ed behavior
       squashError(error);

src/operations/operation.ts

@@ -2,7 +2,7 @@ import { type BSONSerializeOptions, type Document, resolveBSONOptions } from '..
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import type { Server } from '../sdam/server';
 import type { ClientSession } from '../sessions';
-import { type Timeout } from '../timeout';
+import { type Timeout, type TimeoutContext } from '../timeout';
 import type { MongoDBNamespace } from '../utils';
 
 export const Aspect = {
@@ -67,6 +67,9 @@ export abstract class AbstractOperation<TResult = any> {
   /** @internal */
   timeoutMS?: number;
 
+  /** @internal */
+  timeoutContext!: TimeoutContext;
+
   [kSession]: ClientSession | undefined;
 
   constructor(options: OperationOptions = {}) {
@@ -82,8 +85,6 @@ export abstract class AbstractOperation<TResult = any> {
     this.options = options;
     this.bypassPinningCheck = !!options.bypassPinningCheck;
     this.trySecondaryWrite = false;
-
-    this.timeoutMS = options.timeoutMS;
   }
 
   /** Must match the first key of the command object sent to the server.

src/operations/run_command.ts

@@ -33,7 +33,7 @@ export class RunCommandOperation<T = Document> extends AbstractOperation<T> {
     const res: TODO_NODE_3286 = await server.command(this.ns, this.command, {
       ...this.options,
       readPreference: this.readPreference,
-      timeout: this.timeout,
+      timeoutContext: this.timeoutContext,
       session
     });
     return res;
@@ -62,7 +62,7 @@ export class RunAdminCommandOperation<T = Document> extends AbstractOperation<T>
       ...this.options,
       readPreference: this.readPreference,
       session,
-      timeout: this.timeout
+      timeoutContext: this.timeoutContext
     });
     return res;
   }

src/operations/search_indexes/create.ts

@@ -39,7 +39,10 @@ export class CreateSearchIndexesOperation extends AbstractOperation<string[]> {
       indexes: this.descriptions
     };
 
-    const res = await server.command(namespace, command, { session });
+    const res = await server.command(namespace, command, {
+      session,
+      timeoutContext: this.timeoutContext
+    });
 
     const indexesCreated: Array<{ name: string }> = res?.indexesCreated ?? [];
     return indexesCreated.map(({ name }) => name);

src/operations/search_indexes/drop.ts

@@ -28,7 +28,7 @@ export class DropSearchIndexOperation extends AbstractOperation<void> {
     }
 
     try {
-      await server.command(namespace, command, { session });
+      await server.command(namespace, command, { session, timeoutContext: this.timeoutContext });
     } catch (error) {
       const isNamespaceNotFoundError =
         error instanceof MongoServerError && error.code === MONGODB_ERROR_CODES.NamespaceNotFound;

src/operations/search_indexes/update.ts

@@ -27,7 +27,7 @@ export class UpdateSearchIndexOperation extends AbstractOperation<void> {
       definition: this.definition
     };
 
-    await server.command(namespace, command, { session });
+    await server.command(namespace, command, { session, timeoutContext: this.timeoutContext });
     return;
   }
 }

src/sdam/server.ts

@@ -273,7 +273,7 @@ export class Server extends TypedEventEmitter<ServerEvents> {
   public async command(
     ns: MongoDBNamespace,
     command: Document,
-    options?: CommandOptions
+    options: CommandOptions
   ): Promise<Document>;
 
   public async command(