Consolidate logic

Author: baileympearson

src/cmap/connection.ts

@@ -427,10 +427,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       ...options
     };
 
-    if (!options.omitMaxTimeMS) {
-      const maxTimeMS = options.timeoutContext?.maxTimeMS;
-      if (maxTimeMS && maxTimeMS > 0 && Number.isFinite(maxTimeMS)) cmd.maxTimeMS = maxTimeMS;
-    }
+    options.timeoutContext?.addMaxTimeMSToCommand(cmd, options);
 
     const message = this.supportsOpMsg
       ? new OpMsgRequest(db, cmd, commandOptions)
@@ -446,13 +443,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   ): AsyncGenerator<MongoDBResponse> {
     this.throwIfAborted();
 
-    if (options.timeoutContext?.csotEnabled()) {
-      this.socket.setTimeout(0);
-    } else if (typeof options.socketTimeoutMS === 'number') {
-      this.socket.setTimeout(options.socketTimeoutMS);
-    } else if (this.socketTimeoutMS !== 0) {
-      this.socket.setTimeout(this.socketTimeoutMS);
-    }
+    const timeout =
+      options.socketTimeoutMS ??
+      options?.timeoutContext?.getSocketTimeoutMS() ??
+      this.socketTimeoutMS;
+    this.socket.setTimeout(timeout);
 
     try {
       await this.writeCommand(message, {
@@ -487,11 +482,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
         yield document;
         this.throwIfAborted();
 
-        if (typeof options.socketTimeoutMS === 'number') {
-          this.socket.setTimeout(options.socketTimeoutMS);
-        } else if (this.socketTimeoutMS !== 0) {
-          this.socket.setTimeout(this.socketTimeoutMS);
-        }
+        this.socket.setTimeout(timeout);
       }
     } finally {
       this.socket.setTimeout(0);

src/cursor/abstract_cursor.ts

@@ -1187,12 +1187,16 @@ export class CursorTimeoutContext extends TimeoutContext {
   override get maxTimeMS(): number | null {
     return this.timeoutContext.maxTimeMS;
   }
-
   get timeoutMS(): number | null {
     return this.timeoutContext.csotEnabled() ? this.timeoutContext.timeoutMS : null;
   }
-
   override refreshed(): CursorTimeoutContext {
     return new CursorTimeoutContext(this.timeoutContext.refreshed(), this.owner);
   }
+  override addMaxTimeMSToCommand(command: Document, options: { omitMaxTimeMS?: boolean }): void {
+    this.timeoutContext.addMaxTimeMSToCommand(command, options);
+  }
+  override getSocketTimeoutMS(): number | undefined {
+    return this.timeoutContext.getSocketTimeoutMS();
+  }
 }

src/timeout.ts

@@ -1,3 +1,4 @@
+import { type Document } from 'bson';
 import { clearTimeout, setTimeout } from 'timers';
 
 import { MongoInvalidArgumentError, MongoOperationTimeoutError, MongoRuntimeError } from './error';
@@ -183,6 +184,10 @@ export abstract class TimeoutContext {
 
   /** Returns a new instance of the TimeoutContext, with all timeouts refreshed and restarted. */
   abstract refreshed(): TimeoutContext;
+
+  abstract addMaxTimeMSToCommand(command: Document, options: { omitMaxTimeMS?: boolean }): void;
+
+  abstract getSocketTimeoutMS(): number | undefined;
 }
 
 /** @internal */
@@ -321,6 +326,16 @@ export class CSOTTimeoutContext extends TimeoutContext {
   override refreshed(): CSOTTimeoutContext {
     return new CSOTTimeoutContext(this);
   }
+
+  override addMaxTimeMSToCommand(command: Document, options: { omitMaxTimeMS?: boolean }): void {
+    if (options.omitMaxTimeMS) return;
+    const maxTimeMS = this.remainingTimeMS - this.minRoundTripTime;
+    if (maxTimeMS > 0 && Number.isFinite(maxTimeMS)) command.maxTimeMS = maxTimeMS;
+  }
+
+  override getSocketTimeoutMS(): number | undefined {
+    return 0;
+  }
 }
 
 /** @internal */
@@ -373,4 +388,12 @@ export class LegacyTimeoutContext extends TimeoutContext {
   override refreshed(): LegacyTimeoutContext {
     return new LegacyTimeoutContext(this.options);
   }
+
+  override addMaxTimeMSToCommand(_command: Document, _options: { omitMaxTimeMS?: boolean }): void {
+    // No max timeMS is added to commands in legacy timeout mode.
+  }
+
+  override getSocketTimeoutMS(): number | undefined {
+    return this.options.socketTimeoutMS;
+  }
 }