refactor(NODE-6524): consolidate socket timeout calculation (#4320)

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/cmap/connection_pool.ts

@@ -402,7 +402,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       }
       throw error;
     } finally {
-      if (options.timeoutContext.clearConnectionCheckoutTimeout) timeout?.clear();
+      timeout?.clear();
     }
   }
 

src/cursor/abstract_cursor.ts

@@ -1169,9 +1169,6 @@ export class CursorTimeoutContext extends TimeoutContext {
   override get clearServerSelectionTimeout(): boolean {
     return this.timeoutContext.clearServerSelectionTimeout;
   }
-  override get clearConnectionCheckoutTimeout(): boolean {
-    return this.timeoutContext.clearConnectionCheckoutTimeout;
-  }
   override get timeoutForSocketWrite(): Timeout | null {
     return this.timeoutContext.timeoutForSocketWrite;
   }
@@ -1190,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,5 +1,6 @@
 import { clearTimeout, setTimeout } from 'timers';
 
+import { type Document } from './bson';
 import { MongoInvalidArgumentError, MongoOperationTimeoutError, MongoRuntimeError } from './error';
 import { type ClientSession } from './sessions';
 import { csotMin, noop } from './utils';
@@ -171,8 +172,6 @@ export abstract class TimeoutContext {
 
   abstract get clearServerSelectionTimeout(): boolean;
 
-  abstract get clearConnectionCheckoutTimeout(): boolean;
-
   abstract get timeoutForSocketWrite(): Timeout | null;
 
   abstract get timeoutForSocketRead(): Timeout | null;
@@ -185,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 */
@@ -193,7 +196,6 @@ export class CSOTTimeoutContext extends TimeoutContext {
   serverSelectionTimeoutMS: number;
   socketTimeoutMS?: number;
 
-  clearConnectionCheckoutTimeout: boolean;
   clearServerSelectionTimeout: boolean;
 
   private _serverSelectionTimeout?: Timeout | null;
@@ -212,7 +214,6 @@ export class CSOTTimeoutContext extends TimeoutContext {
     this.socketTimeoutMS = options.socketTimeoutMS;
 
     this.clearServerSelectionTimeout = false;
-    this.clearConnectionCheckoutTimeout = true;
   }
 
   get maxTimeMS(): number {
@@ -325,19 +326,27 @@ 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 */
 export class LegacyTimeoutContext extends TimeoutContext {
   options: LegacyTimeoutContextOptions;
   clearServerSelectionTimeout: boolean;
-  clearConnectionCheckoutTimeout: boolean;
 
   constructor(options: LegacyTimeoutContextOptions) {
     super();
     this.options = options;
     this.clearServerSelectionTimeout = true;
-    this.clearConnectionCheckoutTimeout = true;
   }
 
   csotEnabled(): this is CSOTTimeoutContext {
@@ -379,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;
+  }
 }