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/connect.ts

@@ -13,9 +13,11 @@ import {
   MongoInvalidArgumentError,
   MongoNetworkError,
   MongoNetworkTimeoutError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   needsRetryableWriteLabel
 } from '../error';
+import { Timeout, TimeoutError } from '../timeout';
 import { HostAddress, ns, promiseWithResolvers } from '../utils';
 import { AuthContext } from './auth/auth_provider';
 import { AuthMechanism } from './auth/providers';
@@ -37,6 +39,7 @@ export type Stream = Socket | TLSSocket;
 
 export async function connect(options: ConnectionOptions): Promise<Connection> {
   let connection: Connection | null = null;
+  console.log(options.timeout);
   try {
     const socket = await makeSocket(options);
     connection = makeConnection(options, socket);
@@ -110,7 +113,25 @@ export async function performInitialHandshake(
   }
 
   const start = new Date().getTime();
-  const response = await conn.command(ns('admin.$cmd'), handshakeDoc, handshakeOptions);
+
+  let response: Document;
+  if (options.timeout) {
+    console.log(options.timeout);
+    try {
+      response = await Promise.race([
+        options.timeout,
+        conn.command(ns('admin.$cmd'), handshakeDoc, handshakeOptions)
+      ]);
+    } catch (error) {
+      if (TimeoutError.is(error)) {
+        throw new MongoOperationTimeoutError('Timed out during handshake');
+      } else {
+        throw error;
+      }
+    }
+  } else {
+    response = await conn.command(ns('admin.$cmd'), handshakeDoc, handshakeOptions);
+  }
 
   if (!('isWritablePrimary' in response)) {
     // Provide hello-style response document.
@@ -321,14 +342,40 @@ export async function makeSocket(options: MakeConnectionOptions): Promise<Stream
   const connectTimeoutMS = options.connectTimeoutMS ?? 30000;
   const existingSocket = options.existingSocket;
 
+  let timeout: Timeout | null;
+  if (options.timeout) {
+    options.timeout.throwIfExpired();
+    timeout = Timeout.expires(Timeout.min(connectTimeoutMS, options.timeout.remainingTime));
+  } else {
+    timeout = null;
+  }
+
   let socket: Stream;
 
   if (options.proxyHost != null) {
     // Currently, only Socks5 is supported.
-    return await makeSocks5Connection({
-      ...options,
-      connectTimeoutMS // Should always be present for Socks5
-    });
+    if (timeout) {
+      try {
+        timeout.throwIfExpired();
+        return await Promise.race([
+          timeout,
+          makeSocks5Connection({
+            ...options,
+            connectTimeoutMS // Should always be present for Socks5
+          })
+        ]);
+      } catch (error) {
+        if (TimeoutError.is(error)) {
+          throw new MongoOperationTimeoutError('Failed during socket creation');
+        } else {
+          throw error;
+        }
+      }
+    } else
+      return await makeSocks5Connection({
+        ...options,
+        connectTimeoutMS // Should always be present for Socks5
+      });
   }
 
   if (useTLS) {
@@ -370,11 +417,22 @@ export async function makeSocket(options: MakeConnectionOptions): Promise<Stream
   }
 
   try {
-    socket = await connectedSocket;
-    return socket;
+    // TODO: race Promises here
+    if (timeout) {
+      timeout.throwIfExpired();
+      socket = await Promise.race([timeout, connectedSocket]);
+      return socket;
+    } else {
+      socket = await connectedSocket;
+      return socket;
+    }
   } catch (error) {
     socket.destroy();
-    throw error;
+    if (TimeoutError.is(error)) {
+      throw new MongoOperationTimeoutError('Timed out during socket creation');
+    } else {
+      throw error;
+    }
   } finally {
     socket.setTimeout(0);
     socket.removeAllListeners();
@@ -397,6 +455,7 @@ function loadSocks() {
 }
 
 async function makeSocks5Connection(options: MakeConnectionOptions): Promise<Stream> {
+  // TODO: ADD CSOT support here
   const hostAddress = HostAddress.fromHostPort(
     options.proxyHost ?? '', // proxyHost is guaranteed to set here
     options.proxyPort ?? 1080

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,
@@ -88,6 +89,11 @@ export interface CommandOptions extends BSONSerializeOptions {
   writeConcern?: WriteConcern;
 
   directConnection?: boolean;
+
+  /** @internal */
+  timeout?: Timeout;
+  /** @internal */
+  serverSelectionTimeout?: Timeout;
 }
 
 /** @public */
@@ -126,6 +132,8 @@ export interface ConnectionOptions
   extendedMetadata: Promise<Document>;
   /** @internal */
   mongoLogger?: MongoLogger | undefined;
+  /** @internal */
+  timeout?: Timeout;
 }
 
 /** @public */

src/cmap/connection_pool.ts

@@ -21,6 +21,7 @@ import {
   MongoInvalidArgumentError,
   MongoMissingCredentialsError,
   MongoNetworkError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError
 } from '../error';
@@ -354,7 +355,7 @@ 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)
@@ -364,7 +365,13 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
 
     const { promise, resolve, reject } = promiseWithResolvers<Connection>();
 
-    const timeout = Timeout.expires(waitQueueTimeoutMS);
+    let timeout: Timeout;
+    if (options?.timeout) {
+      options.timeout.throwIfExpired();
+      timeout = options.timeout;
+    } else {
+      timeout = Timeout.expires(waitQueueTimeoutMS);
+    }
 
     const waitQueueMember: WaitQueueMember = {
       resolve,
@@ -373,7 +380,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     };
 
     this[kWaitQueue].push(waitQueueMember);
-    process.nextTick(() => this.processWaitQueue());
+    process.nextTick(() => this.processWaitQueue({ timeout }));
 
     try {
       return await Promise.race([promise, waitQueueMember.timeout]);
@@ -393,6 +400,9 @@ 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', { cause: timeoutError });
+        }
         throw timeoutError;
       }
       throw error;
@@ -616,14 +626,15 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     return true;
   }
 
-  private createConnection(callback: Callback<Connection>) {
+  private createConnection(callback: Callback<Connection>, options?: { timeout?: Timeout }) {
     const connectOptions: ConnectionOptions = {
       ...this.options,
       id: this[kConnectionCounter].next().value,
       generation: this[kGeneration],
       cancellationToken: this[kCancellationToken],
       mongoLogger: this.mongoLogger,
-      authProviders: this[kServer].topology.client.s.authProviders
+      authProviders: this[kServer].topology.client.s.authProviders,
+      timeout: options?.timeout
     };
 
     this[kPending]++;
@@ -735,7 +746,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     }
   }
 
-  private processWaitQueue() {
+  private processWaitQueue(options?: { timeout?: Timeout }) {
     if (this[kProcessingWaitQueue]) {
       return;
     }
@@ -823,7 +834,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
           waitQueueMember.timeout.clear();
         }
         process.nextTick(() => this.processWaitQueue());
-      });
+      }, options);
     }
     this[kProcessingWaitQueue] = false;
   }

src/collection.ts

@@ -254,6 +254,10 @@ export class Collection<TSchema extends Document = Document> {
     this.s.collectionHint = normalizeHintField(v);
   }
 
+  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,10 @@ export class Db {
     return this.s.namespace.toString();
   }
 
+  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 +276,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,12 @@ export class MongoUnexpectedServerResponseError extends MongoRuntimeError {
   }
 }
 
+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,8 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
     const options = {
       ...this.options,
       ...this.bsonOptions,
+      timeout: this.timeout,
+      serverSelectionTimeout: this.serverSelectionTimeout,
       readPreference: this.readPreference,
       session
     };

src/operations/execute_operation.ts

@@ -26,6 +26,7 @@ import {
 } from '../sdam/server_selection';
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
+import { Timeout } from '../timeout';
 import { squashError, supportsRetryableWrites } from '../utils';
 import { AbstractOperation, Aspect } from './operation';
 
@@ -151,9 +152,21 @@ export async function executeOperation<
     selector = readPreference;
   }
 
+  const timeout = operation.timeout;
+  timeout?.throwIfExpired();
+  // TODO: construct serverSelection timeout here, pass it into topology.selectServer and store on
+  // operation for use in operation.execute when we have to perform connection checkout
+  operation.serverSelectionTimeout =
+    operation.timeout != null
+      ? Timeout.expires(
+          Timeout.min(operation.timeout.remainingTime, topology.s.serverSelectionTimeoutMS)
+        )
+      : undefined;
+
   const server = await topology.selectServer(selector, {
     session,
-    operationName: operation.commandName
+    operationName: operation.commandName,
+    timeout: operation.serverSelectionTimeout
   });
 
   if (session == null) {
@@ -264,6 +277,7 @@ 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
   });

src/operations/find.ts

@@ -114,6 +114,7 @@ export class FindOperation extends CommandOperation<Document> {
     return await server.command(this.ns, findCommand, {
       ...this.options,
       ...this.bsonOptions,
+      timeout: this.timeout,
       documentsReturnedIn: 'firstBatch',
       session
     });

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,12 @@ export abstract class AbstractOperation<TResult = any> {
 
   options: OperationOptions;
 
+  /** @internal */
+  timeout?: Timeout;
+
+  /** @internal */
+  serverSelectionTimeout?: Timeout;
+
   [kSession]: ClientSession | undefined;
 
   constructor(options: OperationOptions = {}) {
@@ -76,6 +83,10 @@ export abstract class AbstractOperation<TResult = any> {
     this.options = options;
     this.bypassPinningCheck = !!options.bypassPinningCheck;
     this.trySecondaryWrite = false;
+
+    if (options.timeoutMS != null) {
+      this.timeout = Timeout.expires(options.timeoutMS);
+    }
   }
 
   /** Must match the first key of the command object sent to the server.