add tests

Author: W-A-James

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; //TODO(NODE-6187): Make this required
 }
 
 /** @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

@@ -556,7 +556,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) {
@@ -272,7 +277,8 @@ async function retryOperation<
     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/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/sdam/topology.ts

@@ -34,11 +34,10 @@ import { MongoLoggableComponent, type MongoLogger, SeverityLevel } from '../mong
 import { TypedEventEmitter } from '../mongo_types';
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import type { ClientSession } from '../sessions';
-import { Timeout, TimeoutError } from '../timeout';
+import { Timeout, TimeoutContext, TimeoutError } from '../timeout';
 import type { Transaction } from '../transactions';
 import {
   type Callback,
-  csotMin,
   type EventEmitterWithState,
   HostAddress,
   List,
@@ -181,6 +180,8 @@ export interface SelectServerOptions {
   previousServer?: ServerDescription;
   /** @internal*/
   timeout?: Timeout;
+  /** @internal */
+  timeoutContext?: TimeoutContext;
 }
 
 /** @public */
@@ -459,12 +460,21 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
     }
 
     const timeoutMS = this.client.options.timeoutMS;
-    const timeout = timeoutMS != null ? Timeout.expires(timeoutMS) : undefined;
+    const serverSelectionTimeoutMS = this.client.options.serverSelectionTimeoutMS;
     const readPreference = options.readPreference ?? ReadPreference.primary;
+
+    // TODO: figure out if there is a way we can remove this and use the command that initiated the
+    // autoconnect's timeoutContext
+    const timeoutContext = TimeoutContext.create({
+      timeoutMS,
+      serverSelectionTimeoutMS,
+      waitQueueTimeoutMS: this.client.options.waitQueueTimeoutMS
+    });
+
     const selectServerOptions = {
       operationName: 'ping',
-      timeout,
-      ...options
+      ...options,
+      timeoutContext
     };
     try {
       const server = await this.selectServer(
@@ -474,7 +484,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
 
       const skipPingOnConnect = this.s.options[Symbol.for('@@mdb.skipPingOnConnect')] === true;
       if (!skipPingOnConnect && server && this.s.credentials) {
-        await server.command(ns('admin.$cmd'), { ping: 1 }, { timeout });
+        await server.command(ns('admin.$cmd'), { ping: 1 }, { timeoutContext });
         stateTransition(this, STATE_CONNECTED);
         this.emit(Topology.OPEN, this);
         this.emit(Topology.CONNECT, this);
@@ -563,24 +573,10 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
         new ServerSelectionStartedEvent(selector, this.description, options.operationName)
       );
     }
-    const serverSelectionTimeoutMS = options.serverSelectionTimeoutMS ?? 0;
-    let timeout: Timeout | null;
-    if (options.timeout) {
-      // CSOT Enabled
-      if (options.timeout.duration > 0 || serverSelectionTimeoutMS > 0) {
-        if (
-          options.timeout.duration === serverSelectionTimeoutMS ||
-          csotMin(options.timeout.duration, serverSelectionTimeoutMS) < serverSelectionTimeoutMS
-        ) {
-          timeout = options.timeout;
-        } else {
-          timeout = Timeout.expires(serverSelectionTimeoutMS);
-        }
-      } else {
-        timeout = null;
-      }
-    } else {
-      timeout = Timeout.expires(serverSelectionTimeoutMS);
+    let timeout;
+    if (options.timeoutContext) timeout = options.timeoutContext.serverSelectionTimeout;
+    else {
+      timeout = Timeout.expires(options.serverSelectionTimeoutMS ?? 0);
     }
 
     const isSharded = this.description.type === TopologyType.Sharded;
@@ -654,7 +650,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
           );
         }
 
-        if (options.timeout) {
+        if (options.timeoutContext?.csotEnabled()) {
           throw new MongoOperationTimeoutError('Timed out during server selection', {
             cause: timeoutError
           });
@@ -664,7 +660,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
       // Other server selection error
       throw error;
     } finally {
-      if (timeout !== options.timeout) timeout?.clear();
+      if (options.timeoutContext?.clearServerSelectionTimeout) timeout?.clear();
     }
   }
   /**