refactor(NODE-6411): AbstractCursor accepts an external timeout context (#4264)

Author: baileympearson

src/cmap/connection.ts

@@ -427,9 +427,9 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       ...options
     };
 
-    if (!options.omitMaxTimeMS && options.timeoutContext?.csotEnabled()) {
-      const { maxTimeMS } = options.timeoutContext;
-      if (maxTimeMS > 0 && Number.isFinite(maxTimeMS)) cmd.maxTimeMS = maxTimeMS;
+    if (!options.omitMaxTimeMS) {
+      const maxTimeMS = options.timeoutContext?.maxTimeMS;
+      if (maxTimeMS && maxTimeMS > 0 && Number.isFinite(maxTimeMS)) cmd.maxTimeMS = maxTimeMS;
     }
 
     const message = this.supportsOpMsg

src/cursor/abstract_cursor.ts

@@ -21,7 +21,7 @@ import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import { type AsyncDisposable, configureResourceManagement } from '../resource_management';
 import type { Server } from '../sdam/server';
 import { ClientSession, maybeClearPinnedConnection } from '../sessions';
-import { TimeoutContext } from '../timeout';
+import { type CSOTTimeoutContext, type Timeout, TimeoutContext } from '../timeout';
 import { type MongoDBNamespace, squashError } from '../utils';
 
 /**
@@ -119,6 +119,14 @@ export interface AbstractCursorOptions extends BSONSerializeOptions {
   timeoutMS?: number;
   /** @internal TODO(NODE-5688): make this public */
   timeoutMode?: CursorTimeoutMode;
+
+  /**
+   * @internal
+   *
+   * A timeout context to govern the total time the cursor can live.  If provided, the cursor
+   * cannot be used in ITERATION mode.
+   */
+  timeoutContext?: CursorTimeoutContext;
 }
 
 /** @internal */
@@ -171,7 +179,7 @@ export abstract class AbstractCursor<
   /** @internal */
   protected readonly cursorOptions: InternalAbstractCursorOptions;
   /** @internal */
-  protected timeoutContext?: TimeoutContext;
+  protected timeoutContext?: CursorTimeoutContext;
 
   /** @event */
   static readonly CLOSE = 'close' as const;
@@ -205,20 +213,12 @@ export abstract class AbstractCursor<
     };
     this.cursorOptions.timeoutMS = options.timeoutMS;
     if (this.cursorOptions.timeoutMS != null) {
-      if (options.timeoutMode == null) {
-        if (options.tailable) {
-          this.cursorOptions.timeoutMode = CursorTimeoutMode.ITERATION;
-        } else {
-          this.cursorOptions.timeoutMode = CursorTimeoutMode.LIFETIME;
-        }
-      } else {
-        if (options.tailable && this.cursorOptions.timeoutMode === CursorTimeoutMode.LIFETIME) {
-          throw new MongoInvalidArgumentError(
-            "Cannot set tailable cursor's timeoutMode to LIFETIME"
-          );
-        }
-        this.cursorOptions.timeoutMode = options.timeoutMode;
+      if (options.tailable && this.cursorOptions.timeoutMode === CursorTimeoutMode.LIFETIME) {
+        throw new MongoInvalidArgumentError("Cannot set tailable cursor's timeoutMode to LIFETIME");
       }
+      this.cursorOptions.timeoutMode =
+        options.timeoutMode ??
+        (options.tailable ? CursorTimeoutMode.ITERATION : CursorTimeoutMode.LIFETIME);
     } else {
       if (options.timeoutMode != null)
         throw new MongoInvalidArgumentError('Cannot set timeoutMode without setting timeoutMS');
@@ -264,6 +264,17 @@ export abstract class AbstractCursor<
         utf8: options?.enableUtf8Validation === false ? false : true
       }
     };
+
+    if (
+      options.timeoutContext != null &&
+      options.timeoutMS != null &&
+      this.cursorOptions.timeoutMode !== CursorTimeoutMode.LIFETIME
+    ) {
+      throw new MongoAPIError(
+        `cannot create a cursor with an externally provided timeout context that doesn't use timeoutMode=CURSOR_LIFETIME.`
+      );
+    }
+    this.timeoutContext = options.timeoutContext;
   }
 
   /**
@@ -721,6 +732,9 @@ export abstract class AbstractCursor<
    * if the resultant data has already been retrieved by this cursor.
    */
   rewind(): void {
+    if (this.timeoutContext && this.timeoutContext.owner !== this) {
+      throw new MongoAPIError(`Cannot rewind cursor that does not own its timeout context.`);
+    }
     if (!this.initialized) {
       return;
     }
@@ -790,10 +804,13 @@ export abstract class AbstractCursor<
    */
   private async cursorInit(): Promise<void> {
     if (this.cursorOptions.timeoutMS != null) {
-      this.timeoutContext = TimeoutContext.create({
-        serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
-        timeoutMS: this.cursorOptions.timeoutMS
-      });
+      this.timeoutContext ??= new CursorTimeoutContext(
+        TimeoutContext.create({
+          serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
+          timeoutMS: this.cursorOptions.timeoutMS
+        }),
+        this
+      );
     }
     try {
       const state = await this._initialize(this.cursorSession);
@@ -872,6 +889,20 @@ export abstract class AbstractCursor<
   private async cleanup(timeoutMS?: number, error?: Error) {
     this.isClosed = true;
     const session = this.cursorSession;
+    const timeoutContextForKillCursors = (): CursorTimeoutContext | undefined => {
+      if (timeoutMS != null) {
+        this.timeoutContext?.clear();
+        return new CursorTimeoutContext(
+          TimeoutContext.create({
+            serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
+            timeoutMS
+          }),
+          this
+        );
+      } else {
+        return this.timeoutContext?.refreshed();
+      }
+    };
     try {
       if (
         !this.isKilled &&
@@ -884,23 +915,13 @@ export abstract class AbstractCursor<
         this.isKilled = true;
         const cursorId = this.cursorId;
         this.cursorId = Long.ZERO;
-        let timeoutContext: TimeoutContext | undefined;
-        if (timeoutMS != null) {
-          this.timeoutContext?.clear();
-          timeoutContext = TimeoutContext.create({
-            serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
-            timeoutMS
-          });
-        } else {
-          this.timeoutContext?.refresh();
-          timeoutContext = this.timeoutContext;
-        }
+
         await executeOperation(
           this.cursorClient,
           new KillCursorsOperation(cursorId, this.cursorNamespace, this.selectedServer, {
             session
           }),
-          timeoutContext
+          timeoutContextForKillCursors()
         );
       }
     } catch (error) {
@@ -1042,3 +1063,54 @@ class ReadableCursorStream extends Readable {
 }
 
 configureResourceManagement(AbstractCursor.prototype);
+
+/**
+ * @internal
+ * The cursor timeout context is a wrapper around a timeout context
+ * that keeps track of the "owner" of the cursor.  For timeout contexts
+ * instantiated inside a cursor, the owner will be the cursor.
+ *
+ * All timeout behavior is exactly the same as the wrapped timeout context's.
+ */
+export class CursorTimeoutContext extends TimeoutContext {
+  constructor(
+    public timeoutContext: TimeoutContext,
+    public owner: symbol | AbstractCursor
+  ) {
+    super();
+  }
+  override get serverSelectionTimeout(): Timeout | null {
+    return this.timeoutContext.serverSelectionTimeout;
+  }
+  override get connectionCheckoutTimeout(): Timeout | null {
+    return this.timeoutContext.connectionCheckoutTimeout;
+  }
+  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;
+  }
+  override get timeoutForSocketRead(): Timeout | null {
+    return this.timeoutContext.timeoutForSocketRead;
+  }
+  override csotEnabled(): this is CSOTTimeoutContext {
+    return this.timeoutContext.csotEnabled();
+  }
+  override refresh(): void {
+    return this.timeoutContext.refresh();
+  }
+  override clear(): void {
+    return this.timeoutContext.clear();
+  }
+  override get maxTimeMS(): number | null {
+    return this.timeoutContext.maxTimeMS;
+  }
+
+  override refreshed(): CursorTimeoutContext {
+    return new CursorTimeoutContext(this.timeoutContext.refreshed(), this.owner);
+  }
+}

src/index.ts

@@ -359,6 +359,7 @@ export type {
   CursorStreamOptions
 } from './cursor/abstract_cursor';
 export type {
+  CursorTimeoutContext,
   InitialCursorResponse,
   InternalAbstractCursorOptions
 } from './cursor/abstract_cursor';

src/operations/find.ts

@@ -1,6 +1,6 @@
 import type { Document } from '../bson';
 import { CursorResponse, ExplainedCursorResponse } from '../cmap/wire_protocol/responses';
-import { type CursorTimeoutMode } from '../cursor/abstract_cursor';
+import { type AbstractCursorOptions, type CursorTimeoutMode } from '../cursor/abstract_cursor';
 import { MongoInvalidArgumentError } from '../error';
 import { type ExplainOptions } from '../explain';
 import { ReadConcern } from '../read_concern';
@@ -18,7 +18,8 @@ import { Aspect, defineAspects, type Hint } from './operation';
  */
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 export interface FindOptions<TSchema extends Document = Document>
-  extends Omit<CommandOperationOptions, 'writeConcern' | 'explain'> {
+  extends Omit<CommandOperationOptions, 'writeConcern' | 'explain'>,
+    AbstractCursorOptions {
   /** Sets the limit of documents returned in the query. */
   limit?: number;
   /** Set to sort the documents coming back from the query. Array of indexes, `[['a', 1]]` etc. */

src/timeout.ts

@@ -178,6 +178,8 @@ export abstract class TimeoutContext {
     else throw new MongoRuntimeError('Unrecognized options');
   }
 
+  abstract get maxTimeMS(): number | null;
+
   abstract get serverSelectionTimeout(): Timeout | null;
 
   abstract get connectionCheckoutTimeout(): Timeout | null;
@@ -195,6 +197,9 @@ export abstract class TimeoutContext {
   abstract refresh(): void;
 
   abstract clear(): void;
+
+  /** Returns a new instance of the TimeoutContext, with all timeouts refreshed and restarted. */
+  abstract refreshed(): TimeoutContext;
 }
 
 /** @internal */
@@ -317,6 +322,10 @@ export class CSOTTimeoutContext extends TimeoutContext {
       throw new MongoOperationTimeoutError(message ?? `Expired after ${this.timeoutMS}ms`);
     return remainingTimeMS;
   }
+
+  override refreshed(): CSOTTimeoutContext {
+    return new CSOTTimeoutContext(this);
+  }
 }
 
 /** @internal */
@@ -363,4 +372,12 @@ export class LegacyTimeoutContext extends TimeoutContext {
   clear(): void {
     return;
   }
+
+  get maxTimeMS() {
+    return null;
+  }
+
+  override refreshed(): LegacyTimeoutContext {
+    return new LegacyTimeoutContext(this.options);
+  }
 }

test/integration/client-side-operations-timeout/node_csot.test.ts

@@ -26,7 +26,7 @@ import {
   MongoServerError,
   ObjectId
 } from '../../mongodb';
-import { type FailPoint } from '../../tools/utils';
+import { type FailPoint, waitUntilPoolsFilled } from '../../tools/utils';
 
 const metadata = { requires: { mongodb: '>=4.4' } };
 
@@ -362,7 +362,7 @@ describe('CSOT driver tests', metadata, () => {
     };
 
     beforeEach(async function () {
-      internalClient = this.configuration.newClient();
+      internalClient = this.configuration.newClient({});
       await internalClient
         .db('db')
         .dropCollection('coll')
@@ -378,7 +378,11 @@ describe('CSOT driver tests', metadata, () => {
 
       await internalClient.db().admin().command(failpoint);
 
-      client = this.configuration.newClient(undefined, { monitorCommands: true });
+      client = this.configuration.newClient(undefined, { monitorCommands: true, minPoolSize: 10 });
+
+      // wait for a handful of connections to have been established
+      await waitUntilPoolsFilled(client, AbortSignal.timeout(30_000), 5);
+
       commandStarted = [];
       commandSucceeded = [];
       client.on('commandStarted', ev => commandStarted.push(ev));
@@ -492,7 +496,13 @@ describe('CSOT driver tests', metadata, () => {
 
         await internalClient.db().admin().command(failpoint);
 
-        client = this.configuration.newClient(undefined, { monitorCommands: true });
+        client = this.configuration.newClient(undefined, {
+          monitorCommands: true,
+          minPoolSize: 10
+        });
+        // wait for a handful of connections to have been established
+        await waitUntilPoolsFilled(client, AbortSignal.timeout(30_000), 5);
+
         commandStarted = [];
         commandSucceeded = [];
         client.on('commandStarted', ev => commandStarted.push(ev));