feat(CSOT) - feature branch

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
       })
     );
   }

src/cmap/connection.ts

@@ -29,6 +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 {
   BufferPool,
   calculateDurationInMs,
@@ -92,6 +93,9 @@ export interface CommandOptions extends BSONSerializeOptions {
   writeConcern?: WriteConcern;
 
   directConnection?: boolean;
+
+  /** @internal */
+  operationTimeout?: Timeout;
 }
 
 /** @public */

src/cmap/connection_pool.ts

@@ -21,13 +21,14 @@ import {
   MongoInvalidArgumentError,
   MongoMissingCredentialsError,
   MongoNetworkError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError
 } from '../error';
 import { CancellationToken, TypedEventEmitter } from '../mongo_types';
 import type { Server } from '../sdam/server';
 import { Timeout, TimeoutError } from '../timeout';
-import { type Callback, List, makeCounter, promiseWithResolvers } from '../utils';
+import { type Callback, csotMin, List, makeCounter, promiseWithResolvers } from '../utils';
 import { connect } from './connect';
 import { Connection, type ConnectionEvents, type ConnectionOptions } from './connection';
 import {
@@ -354,17 +355,33 @@ 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(): Promise<Connection> {
+  async checkOut(options?: { timeout?: Timeout }): 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>();
 
-    const timeout = Timeout.expires(waitQueueTimeoutMS);
+    let timeout: Timeout;
+    if (options?.timeout) {
+      // CSOT enabled
+      // Determine if we're using the timeout passed in or a new timeout
+      if (
+        csotMin(options.timeout.duration, serverSelectionTimeoutMS) === serverSelectionTimeoutMS
+      ) {
+        timeout = Timeout.expires(serverSelectionTimeoutMS - options.timeout.timeElapsed);
+      } else {
+        timeout = options.timeout;
+      }
+    } else {
+      timeout = Timeout.expires(waitQueueTimeoutMS);
+    }
+
+    timeout.throwIfExpired();
 
     const waitQueueMember: WaitQueueMember = {
       resolve,
@@ -393,6 +410,11 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
             : 'Timed out while checking out a connection from connection pool',
           this.address
         );
+        if (options?.timeout) {
+          throw new MongoOperationTimeoutError('Timed out during connection checkout', {
+            cause: timeoutError
+          });
+        }
         throw timeoutError;
       }
       throw error;

src/collection.ts

@@ -254,6 +254,11 @@ export class Collection<TSchema extends Document = Document> {
     this.s.collectionHint = normalizeHintField(v);
   }
 
+  /** @internal */
+  get timeoutMS(): number | undefined {
+    return this.s.options.timeoutMS;
+  }
+
   /**
    * Inserts a single document into MongoDB. If documents passed in do not contain the **_id** field,
    * one will be added to each of the documents missing it by the driver, mutating the document. This behavior

src/db.ts

@@ -222,6 +222,11 @@ export class Db {
     return this.s.namespace.toString();
   }
 
+  /** @internal */
+  get timeoutMS(): number | undefined {
+    return this.s.options?.timeoutMS;
+  }
+
   /**
    * Create a new collection on a server with the specified options. Use this to create capped collections.
    * More information about command options available at https://www.mongodb.com/docs/manual/reference/command/create/
@@ -272,6 +277,7 @@ export class Db {
       this.client,
       new RunCommandOperation(this, command, {
         ...resolveBSONOptions(options),
+        timeoutMS: options?.timeoutMS,
         session: options?.session,
         readPreference: options?.readPreference
       })

src/error.ts

@@ -702,6 +702,15 @@ export class MongoUnexpectedServerResponseError extends MongoRuntimeError {
   }
 }
 
+/**
+ * @internal
+ */
+export class MongoOperationTimeoutError extends MongoRuntimeError {
+  override get name(): string {
+    return 'MongoOperationTimeoutError';
+  }
+}
+
 /**
  * An error thrown when the user attempts to add options to a cursor that has already been
  * initialized

src/operations/command.ts

@@ -64,6 +64,7 @@ export interface OperationParent {
   writeConcern?: WriteConcern;
   readPreference?: ReadPreference;
   bsonOptions?: BSONSerializeOptions;
+  timeoutMS?: number;
 }
 
 /** @internal */
@@ -117,6 +118,7 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
     const options = {
       ...this.options,
       ...this.bsonOptions,
+      operationTimeout: this.operationTimeout,
       readPreference: this.readPreference,
       session
     };

src/operations/execute_operation.ts

@@ -152,9 +152,13 @@ export async function executeOperation<
     selector = readPreference;
   }
 
+  const timeout = operation.operationTimeout;
+  timeout?.throwIfExpired();
+
   const server = await topology.selectServer(selector, {
     session,
-    operationName: operation.commandName
+    operationName: operation.commandName,
+    timeout
   });
 
   if (session == null) {
@@ -265,6 +269,7 @@ async function retryOperation<
   // select a new server, and attempt to retry the operation
   const server = await topology.selectServer(selector, {
     session,
+    timeout: operation.operationTimeout,
     operationName: operation.commandName,
     previousServer
   });

src/operations/operation.ts

@@ -2,6 +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 { Timeout } from '../timeout';
 import type { MongoDBNamespace } from '../utils';
 
 export const Aspect = {
@@ -61,6 +62,9 @@ export abstract class AbstractOperation<TResult = any> {
 
   options: OperationOptions;
 
+  /** @internal */
+  operationTimeout?: Timeout;
+
   [kSession]: ClientSession | undefined;
 
   constructor(options: OperationOptions = {}) {
@@ -76,6 +80,10 @@ export abstract class AbstractOperation<TResult = any> {
     this.options = options;
     this.bypassPinningCheck = !!options.bypassPinningCheck;
     this.trySecondaryWrite = false;
+
+    if (options.timeoutMS != null) {
+      this.operationTimeout = Timeout.expires(options.timeoutMS);
+    }
   }
 
   /** Must match the first key of the command object sent to the server.

src/operations/run_command.ts

@@ -13,6 +13,8 @@ export type RunCommandOptions = {
   session?: ClientSession;
   /** The read preference */
   readPreference?: ReadPreferenceLike;
+  /** @internal */
+  timeoutMS?: number;
 } & BSONSerializeOptions;
 
 /** @internal */
@@ -31,6 +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,
+      operationTimeout: this.operationTimeout,
       session
     });
     return res;
@@ -58,7 +61,8 @@ export class RunAdminCommandOperation<T = Document> extends AbstractOperation<T>
     const res: TODO_NODE_3286 = await server.command(this.ns, this.command, {
       ...this.options,
       readPreference: this.readPreference,
-      session
+      session,
+      operationTimeout: this.operationTimeout
     });
     return res;
   }

src/sdam/server.ts

@@ -310,7 +310,11 @@ export class Server extends TypedEventEmitter<ServerEvents> {
     this.incrementOperationCount();
     if (conn == null) {
       try {
-        conn = await this.pool.checkOut();
+        if (options.operationTimeout) {
+          conn = await this.pool.checkOut({ timeout: options.operationTimeout });
+        } else {
+          conn = await this.pool.checkOut();
+        }
         if (this.loadBalanced && isPinnableCommand(cmd, session)) {
           session?.pin(conn);
         }
@@ -333,6 +337,7 @@ export class Server extends TypedEventEmitter<ServerEvents> {
         operationError.code === MONGODB_ERROR_CODES.Reauthenticate
       ) {
         await this.pool.reauthenticate(conn);
+        // TODO(NODE-5682): Implement CSOT support for socket read/write at the connection layer
         try {
           return await conn.command(ns, cmd, finalOptions, responseType);
         } catch (commandError) {

src/sdam/topology.ts

@@ -24,6 +24,7 @@ import {
   type MongoDriverError,
   MongoError,
   MongoErrorLabel,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerSelectionError,
   MongoTopologyClosedError
@@ -457,8 +458,14 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
       }
     }
 
+    const timeoutMS = this.client.options.timeoutMS;
+    const operationTimeout = timeoutMS != null ? Timeout.expires(timeoutMS) : undefined;
     const readPreference = options.readPreference ?? ReadPreference.primary;
-    const selectServerOptions = { operationName: 'ping', ...options };
+    const selectServerOptions = {
+      operationName: 'ping',
+      timeout: operationTimeout,
+      ...options
+    };
     try {
       const server = await this.selectServer(
         readPreferenceServerSelector(readPreference),
@@ -467,7 +474,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 }, {});
+        await server.command(ns('admin.$cmd'), { ping: 1 }, { operationTimeout });
         stateTransition(this, STATE_CONNECTED);
         this.emit(Topology.OPEN, this);
         this.emit(Topology.CONNECT, this);
@@ -556,6 +563,18 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
         new ServerSelectionStartedEvent(selector, this.description, options.operationName)
       );
     }
+    const serverSelectionTimeoutMS = options.serverSelectionTimeoutMS ?? 0;
+    let timeout: Timeout;
+    if (options.timeout) {
+      // CSOT Enabled
+      timeout =
+        Math.min(options.timeout.remainingTime, serverSelectionTimeoutMS) ===
+        serverSelectionTimeoutMS
+          ? Timeout.expires(serverSelectionTimeoutMS)
+          : options.timeout;
+    } else {
+      timeout = Timeout.expires(serverSelectionTimeoutMS);
+    }
 
     const isSharded = this.description.type === TopologyType.Sharded;
     const session = options.session;
@@ -582,7 +601,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
     }
 
     const { promise: serverPromise, resolve, reject } = promiseWithResolvers<Server>();
-    const timeout = Timeout.expires(options.serverSelectionTimeoutMS ?? 0);
+
     const waitQueueMember: ServerSelectionRequest = {
       serverSelector,
       topologyDescription: this.description,
@@ -601,12 +620,12 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
     processWaitQueue(this);
 
     try {
+      waitQueueMember.timeout.throwIfExpired();
       return await Promise.race([serverPromise, waitQueueMember.timeout]);
     } catch (error) {
       if (TimeoutError.is(error)) {
         // Timeout
         waitQueueMember[kCancelled] = true;
-        timeout.clear();
         const timeoutError = new MongoServerSelectionError(
           `Server selection timed out after ${options.serverSelectionTimeoutMS} ms`,
           this.description
@@ -628,6 +647,11 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
           );
         }
 
+        if (options.timeout) {
+          throw new MongoOperationTimeoutError('Timed out during server selection', {
+            cause: timeoutError
+          });
+        }
         throw timeoutError;
       }
       // Other server selection error

src/timeout.ts

@@ -40,6 +40,16 @@ export class Timeout extends Promise<never> {
   public duration: number;
   public timedOut = false;
 
+  get remainingTime(): number {
+    if (this.timedOut) return 0;
+    if (this.duration === 0) return Infinity;
+    return this.start + this.duration - Math.trunc(performance.now());
+  }
+
+  get timeElapsed(): number {
+    return Math.trunc(performance.now()) - this.start;
+  }
+
   /** Create a new timeout that expires in `duration` ms */
   private constructor(executor: Executor = () => null, duration: number, unref = false) {
     let reject!: Reject;
@@ -78,6 +88,10 @@ export class Timeout extends Promise<never> {
     this.id = undefined;
   }
 
+  throwIfExpired(): void {
+    if (this.timedOut) throw new TimeoutError('Timed out');
+  }
+
   public static expires(durationMS: number, unref?: boolean): Timeout {
     return new Timeout(undefined, durationMS, unref);
   }

src/utils.ts

@@ -527,6 +527,10 @@ export function resolveOptions<T extends CommandOperationOptions>(
     result.readPreference = readPreference;
   }
 
+  const timeoutMS = options?.timeoutMS;
+
+  result.timeoutMS = timeoutMS ?? parent?.timeoutMS;
+
   return result;
 }
 
@@ -1330,6 +1334,12 @@ export async function fileIsAccessible(fileName: string, mode?: number) {
   }
 }
 
+export function csotMin(duration1: number, duration2: number): number {
+  if (duration1 === 0) return duration2;
+  if (duration2 === 0) return duration1;
+  return Math.min(duration1, duration2);
+}
+
 export function noop() {
   return;
 }