add CSOT logic to server selection

Author: W-A-James

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,8 @@ export interface CommandOptions extends BSONSerializeOptions {
   writeConcern?: WriteConcern;
 
   directConnection?: boolean;
+
+  timeout?: Timeout | null;
 }
 
 /** @public */

src/cmap/wire_protocol/on_data.ts

@@ -1,5 +1,6 @@
 import { type EventEmitter } from 'events';
 
+import { type Timeout } from '../../timeout';
 import { List, promiseWithResolvers } from '../../utils';
 
 /**
@@ -18,7 +19,7 @@ type PendingPromises = Omit<
  * Returns an AsyncIterator that iterates each 'data' event emitted from emitter.
  * It will reject upon an error event.
  */
-export function onData(emitter: EventEmitter) {
+export function onData(emitter: EventEmitter, options?: { timeout?: Timeout | null }) {
   // Setup pending events and pending promise lists
   /**
    * When the caller has not yet called .next(), we store the
@@ -86,6 +87,8 @@ export function onData(emitter: EventEmitter) {
   // Adding event handlers
   emitter.on('data', eventHandler);
   emitter.on('error', errorHandler);
+  // eslint-disable-next-line github/no-then
+  options?.timeout?.then(() => null, errorHandler);
 
   return iterator;
 

src/operations/bulk_write.ts

@@ -36,7 +36,12 @@ export class BulkWriteOperation extends AbstractOperation<BulkWriteResult> {
   ): Promise<BulkWriteResult> {
     const coll = this.collection;
     const operations = this.operations;
-    const options = { ...this.options, ...this.bsonOptions, readPreference: this.readPreference };
+    const options = {
+      ...this.options,
+      ...this.bsonOptions,
+      readPreference: this.readPreference,
+      timeout: this.timeout
+    };
 
     // Create the bulk operation
     const bulk: BulkOperationBase =

src/operations/command.ts

@@ -118,7 +118,8 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
       ...this.options,
       ...this.bsonOptions,
       readPreference: this.readPreference,
-      session
+      session,
+      timeout: this.timeout
     };
 
     const serverWireVersion = maxWireVersion(server);

src/operations/execute_operation.ts

@@ -260,7 +260,8 @@ async function retryOperation<
   const server = await topology.selectServer(selector, {
     session,
     operationName: operation.commandName,
-    previousServer
+    previousServer,
+    timeout: operation.timeout
   });
 
   if (isWriteOperation && !supportsRetryableWrites(server)) {

src/operations/find.ts

@@ -115,7 +115,8 @@ export class FindOperation extends CommandOperation<Document> {
       ...this.options,
       ...this.bsonOptions,
       documentsReturnedIn: 'firstBatch',
-      session
+      session,
+      timeout: this.timeout
     });
   }
 }

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 = {
@@ -63,7 +64,10 @@ export abstract class AbstractOperation<TResult = any> {
 
   [kSession]: ClientSession | undefined;
 
+  timeout?: Timeout | null;
+
   constructor(options: OperationOptions = {}) {
+    this.timeout = options.timeoutMS != null ? Timeout.expires(options.timeoutMS) : null;
     this.readPreference = this.hasAspect(Aspect.WRITE_OPERATION)
       ? ReadPreference.primary
       : ReadPreference.fromOptions(options) ?? ReadPreference.primary;

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,7 +33,8 @@ 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,
-      session
+      session,
+      timeout: this.timeout
     });
     return res;
   }

src/sdam/topology.ts

@@ -180,7 +180,7 @@ export interface SelectServerOptions {
   operationName: string;
   previousServer?: ServerDescription;
   /** @internal*/
-  timeout?: Timeout;
+  timeout?: Timeout | null;
 }
 
 /** @public */
@@ -285,8 +285,8 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
 
     const selectedHosts =
       options.srvMaxHosts == null ||
-        options.srvMaxHosts === 0 ||
-        options.srvMaxHosts >= seedlist.length
+      options.srvMaxHosts === 0 ||
+      options.srvMaxHosts >= seedlist.length
         ? seedlist
         : shuffle(seedlist, options.srvMaxHosts);
 
@@ -932,10 +932,10 @@ function processWaitQueue(topology: Topology) {
       const previousServer = waitQueueMember.previousServer;
       selectedDescriptions = serverSelector
         ? serverSelector(
-          topology.description,
-          serverDescriptions,
-          previousServer ? [previousServer] : []
-        )
+            topology.description,
+            serverDescriptions,
+            previousServer ? [previousServer] : []
+          )
         : serverDescriptions;
     } catch (selectorError) {
       waitQueueMember.timeoutController.clear();