feat(NODE-6090): Implement CSOT logic for connection checkout and server selection

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 TimeoutContext } from '../timeout';
+import { type Timeout } from '../timeout';
 import {
   BufferPool,
   calculateDurationInMs,
@@ -96,7 +96,7 @@ export interface CommandOptions extends BSONSerializeOptions {
   directConnection?: boolean;
 
   /** @internal */
-  timeoutContext?: TimeoutContext;
+  timeout?: Timeout;
 }
 
 /** @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 { type TimeoutContext, TimeoutError } from '../timeout';
-import { type Callback, List, makeCounter, promiseWithResolvers } from '../utils';
+import { Timeout, TimeoutError } from '../timeout';
+import { type Callback, csotMin, List, makeCounter, promiseWithResolvers } from '../utils';
 import { connect } from './connect';
 import { Connection, type ConnectionEvents, type ConnectionOptions } from './connection';
 import {
@@ -354,15 +354,41 @@ 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: { timeoutContext: TimeoutContext }): 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 = options.timeoutContext.connectionCheckoutTimeout;
+    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);
+    }
 
     const waitQueueMember: WaitQueueMember = {
       resolve,
@@ -390,7 +416,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
             : 'Timed out while checking out a connection from connection pool',
           this.address
         );
-        if (options.timeoutContext.csotEnabled()) {
+        if (options?.timeout) {
           throw new MongoOperationTimeoutError('Timed out during connection checkout', {
             cause: timeoutError
           });
@@ -399,7 +425,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       }
       throw error;
     } finally {
-      if (options.timeoutContext.clearConnectionCheckoutTimeout) timeout?.clear();
+      if (timeout !== options?.timeout) timeout?.clear();
     }
   }
 

src/operations/command.ts

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

src/operations/execute_operation.ts

@@ -24,8 +24,8 @@ import {
 } from '../sdam/server_selection';
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
-import { TimeoutContext } from '../timeout';
-import { supportsRetryableWrites } from '../utils';
+import { Timeout } from '../timeout';
+import { squashError, supportsRetryableWrites } from '../utils';
 import { AbstractOperation, Aspect } from './operation';
 
 const MMAPv1_RETRY_WRITES_ERROR_CODE = MONGODB_ERROR_CODES.IllegalOperation;
@@ -200,10 +200,13 @@ async function tryOperation<
     selector = readPreference;
   }
 
-  let server = await topology.selectServer(selector, {
+  const timeout = operation.timeoutMS != null ? Timeout.expires(operation.timeoutMS) : undefined;
+  operation.timeout = timeout;
+
+  const server = await topology.selectServer(selector, {
     session,
     operationName: operation.commandName,
-    timeoutContext
+    timeout
   });
 
   const hasReadAspect = operation.hasAspect(Aspect.READ_OPERATION);
@@ -278,14 +281,67 @@ async function tryOperation<
     } catch (operationError) {
       if (!(operationError instanceof MongoError)) throw operationError;
 
-      if (
-        previousOperationError != null &&
-        operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)
-      ) {
-        throw previousOperationError;
-      }
-      previousServer = server.description;
-      previousOperationError = operationError;
+async function retryOperation<
+  T extends AbstractOperation<TResult>,
+  TResult = ResultTypeFromOperation<T>
+>(
+  operation: T,
+  originalError: MongoError,
+  { session, topology, selector, previousServer }: RetryOptions
+): Promise<TResult> {
+  const isWriteOperation = operation.hasAspect(Aspect.WRITE_OPERATION);
+  const isReadOperation = operation.hasAspect(Aspect.READ_OPERATION);
+
+  if (isWriteOperation && originalError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
+    throw new MongoServerError({
+      message: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,
+      errmsg: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,
+      originalError
+    });
+  }
+
+  if (isWriteOperation && !isRetryableWriteError(originalError)) {
+    throw originalError;
+  }
+
+  if (isReadOperation && !isRetryableReadError(originalError)) {
+    throw originalError;
+  }
+
+  if (
+    originalError instanceof MongoNetworkError &&
+    session.isPinned &&
+    !session.inTransaction() &&
+    operation.hasAspect(Aspect.CURSOR_CREATING)
+  ) {
+    // If we have a cursor and the initial command fails with a network error,
+    // we can retry it on another connection. So we need to check it back in, clear the
+    // pool for the service id, and retry again.
+    session.unpin({ force: true, forceClear: true });
+  }
+
+  // select a new server, and attempt to retry the operation
+  const server = await topology.selectServer(selector, {
+    session,
+    timeout: operation.timeout,
+    operationName: operation.commandName,
+    previousServer
+  });
+
+  if (isWriteOperation && !supportsRetryableWrites(server)) {
+    throw new MongoUnexpectedServerResponseError(
+      'Selected server does not support retryable writes'
+    );
+  }
+
+  try {
+    return await operation.execute(server, session);
+  } catch (retryError) {
+    if (
+      retryError instanceof MongoError &&
+      retryError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)
+    ) {
+      throw originalError;
     }
   }
 

src/operations/find.ts

@@ -119,7 +119,7 @@ export class FindOperation extends CommandOperation<CursorResponse> {
         ...this.bsonOptions,
         documentsReturnedIn: 'firstBatch',
         session,
-        timeoutContext
+        timeout: this.timeout
       },
       this.explain ? ExplainedCursorResponse : CursorResponse
     );

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, type TimeoutContext } from '../timeout';
+import { type Timeout } from '../timeout';
 import type { MongoDBNamespace } from '../utils';
 
 export const Aspect = {
@@ -82,6 +82,8 @@ 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

@@ -47,7 +47,7 @@ export class RunCommandOperation<T = Document> extends AbstractOperation<T> {
         ...this.options,
         readPreference: this.readPreference,
         session,
-        timeoutContext
+        timeout: this.timeout
       },
       this.options.responseType
     );
@@ -82,7 +82,7 @@ export class RunAdminCommandOperation<T = Document> extends AbstractOperation<T>
       ...this.options,
       readPreference: this.readPreference,
       session,
-      timeoutContext
+      timeout: this.timeout
     });
     return res;
   }

src/sdam/topology.ts

@@ -38,6 +38,7 @@ import { Timeout, TimeoutContext, TimeoutError } from '../timeout';
 import type { Transaction } from '../transactions';
 import {
   type Callback,
+  csotMin,
   type EventEmitterWithState,
   HostAddress,
   List,
@@ -460,20 +461,13 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
       }
     }
 
-    const timeoutMS = this.client.s.options.timeoutMS;
-    const serverSelectionTimeoutMS = this.client.s.options.serverSelectionTimeoutMS;
+    const timeoutMS = this.client.options.timeoutMS;
+    const timeout = timeoutMS != null ? Timeout.expires(timeoutMS) : undefined;
     const readPreference = options.readPreference ?? ReadPreference.primary;
-
-    const timeoutContext = TimeoutContext.create({
-      timeoutMS,
-      serverSelectionTimeoutMS,
-      waitQueueTimeoutMS: this.client.s.options.waitQueueTimeoutMS
-    });
-
     const selectServerOptions = {
       operationName: 'ping',
-      ...options,
-      timeoutContext
+      timeout,
+      ...options
     };
     try {
       const server = await this.selectServer(
@@ -483,7 +477,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 }, { timeoutContext });
+        await server.command(ns('admin.$cmd'), { ping: 1 }, { timeout });
         stateTransition(this, STATE_CONNECTED);
         this.emit(Topology.OPEN, this);
         this.emit(Topology.CONNECT, this);
@@ -572,10 +566,24 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
         new ServerSelectionStartedEvent(selector, this.description, options.operationName)
       );
     }
-    let timeout;
-    if (options.timeoutContext) timeout = options.timeoutContext.serverSelectionTimeout;
-    else {
-      timeout = Timeout.expires(options.serverSelectionTimeoutMS ?? 0);
+    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);
     }
 
     const isSharded = this.description.type === TopologyType.Sharded;
@@ -599,7 +607,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
           )
         );
       }
-      if (options.timeoutContext?.clearServerSelectionTimeout) timeout?.clear();
+      if (timeout !== options.timeout) timeout?.clear();
       return transaction.server;
     }
 
@@ -649,7 +657,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
           );
         }
 
-        if (options.timeoutContext?.csotEnabled()) {
+        if (options.timeout) {
           throw new MongoOperationTimeoutError('Timed out during server selection', {
             cause: timeoutError
           });
@@ -659,7 +667,7 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
       // Other server selection error
       throw error;
     } finally {
-      if (options.timeoutContext?.clearServerSelectionTimeout) timeout?.clear();
+      if (timeout !== options.timeout) timeout?.clear();
     }
   }
   /**

test/integration/client-side-operations-timeout/client_side_operations_timeout.unit.test.ts

@@ -33,20 +33,16 @@ describe('CSOT spec unit tests', function () {
       client = this.configuration.newClient({ timeoutMS: 1000 });
       // Spy on connection checkout and pull options argument
       const checkoutSpy = sinon.spy(ConnectionPool.prototype, 'checkOut');
+      const selectServerSpy = sinon.spy(Topology.prototype, 'selectServer');
       const expiresSpy = sinon.spy(Timeout, 'expires');
 
       await client.db('db').collection('collection').insertOne({ x: 1 });
 
       expect(checkoutSpy).to.have.been.calledOnce;
-      const timeoutContext = checkoutSpy.lastCall.args[0].timeoutContext;
-      expect(timeoutContext).to.exist;
+      expect(checkoutSpy.firstCall.args[0].timeout).to.exist;
       // Check that we passed through the timeout
-      //  @ts-expect-error accessing private properties
-      expect(timeoutContext._serverSelectionTimeout).to.be.instanceOf(Timeout);
-      //  @ts-expect-error accessing private properties
-      expect(timeoutContext._serverSelectionTimeout).to.equal(
-        //  @ts-expect-error accessing private properties
-        timeoutContext._connectionCheckoutTimeout
+      expect(checkoutSpy.firstCall.args[0].timeout).to.equal(
+        selectServerSpy.lastCall.lastArg.timeout
       );
 
       // Check that no more Timeouts are constructed after we enter checkout

test/integration/client-side-operations-timeout/node_csot.test.ts

@@ -143,7 +143,7 @@ describe('CSOT driver tests', () => {
       });
 
       it('throws a MongoOperationTimeoutError', {
-        metadata: { requires: { mongodb: '>=4.4', topology: '!load-balanced' } },
+        metadata: { requires: { mongodb: '>=4.4' } },
         test: async function () {
           const commandsStarted = [];
           client = this.configuration.newClient(undefined, { timeoutMS: 1, monitorCommands: true });

test/unit/cmap/connection_pool.test.js

@@ -129,8 +129,8 @@ describe('Connection Pool', function () {
 
     pool.ready();
 
-    const conn = await pool.checkOut({ timeoutContext });
-    const err = await pool.checkOut({ timeoutContext }).catch(e => e);
+    const conn = await pool.checkOut();
+    const err = await pool.checkOut().catch(e => e);
     expect(err).to.exist.and.be.instanceOf(WaitQueueTimeoutError);
     sinon.stub(pool, 'availableConnectionCount').get(() => 0);
     pool.checkIn(conn);