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