From 20564f7a2cddf071a638afc78df9c5437e0fe265 Mon Sep 17 00:00:00 2001 From: Bailey Pearson Date: Tue, 12 Nov 2024 12:46:12 -0700 Subject: [PATCH] refactor(NODE-6524): consolidate socket timeout calculation (#4320) --- src/cmap/connection.ts | 23 +++++++---------------- src/cmap/connection_pool.ts | 2 +- src/cursor/abstract_cursor.ts | 11 ++++++----- src/timeout.ts | 29 +++++++++++++++++++++++------ 4 files changed, 37 insertions(+), 28 deletions(-) diff --git a/src/cmap/connection.ts b/src/cmap/connection.ts index ca7c86a0bad..b92adb50511 100644 --- a/src/cmap/connection.ts +++ b/src/cmap/connection.ts @@ -427,10 +427,7 @@ export class Connection extends TypedEventEmitter { ...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 { ): AsyncGenerator { 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 { 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); diff --git a/src/cmap/connection_pool.ts b/src/cmap/connection_pool.ts index 2cd2bcc2c19..83449a91ef9 100644 --- a/src/cmap/connection_pool.ts +++ b/src/cmap/connection_pool.ts @@ -402,7 +402,7 @@ export class ConnectionPool extends TypedEventEmitter { } throw error; } finally { - if (options.timeoutContext.clearConnectionCheckoutTimeout) timeout?.clear(); + timeout?.clear(); } } diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 66bfbed0078..8eccdfcf630 100644 --- a/src/cursor/abstract_cursor.ts +++ b/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(); + } } diff --git a/src/timeout.ts b/src/timeout.ts index 47c27c7b90e..3b1dbcb2346 100644 --- a/src/timeout.ts +++ b/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; + } }