start csot cursor implementation

Author: W-A-James

src/cmap/connection.ts

@@ -85,6 +85,7 @@ export interface CommandOptions extends BSONSerializeOptions {
   documentsReturnedIn?: string;
   noResponse?: boolean;
   omitReadPreference?: boolean;
+  omitMaxTimeMS?: boolean;
 
   // TODO(NODE-2802): Currently the CommandOptions take a property willRetryWrite which is a hint
   // from executeOperation that the txnNum should be applied to this command.
@@ -418,7 +419,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       ...options
     };
 
-    if (options.timeoutContext?.csotEnabled()) {
+    if (!options.omitMaxTimeMS && options.timeoutContext?.csotEnabled()) {
       const { maxTimeMS } = options.timeoutContext;
       if (maxTimeMS > 0 && Number.isFinite(maxTimeMS)) cmd.maxTimeMS = maxTimeMS;
     }

src/cursor/abstract_cursor.ts

@@ -20,6 +20,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 MongoDBNamespace, squashError } from '../utils';
 
 /**
@@ -59,11 +60,15 @@ export interface CursorStreamOptions {
 /** @public */
 export type CursorFlag = (typeof CURSOR_FLAGS)[number];
 
+/** @internal */
 export const CursorTimeoutMode = Object.freeze({
   ITERATION: 'iteration',
   LIFETIME: 'lifetime'
 } as const);
 
+/** @internal
+ * TODO(NODE-5688): Document and release
+ * */
 export type CursorTimeoutMode = (typeof CursorTimeoutMode)[keyof typeof CursorTimeoutMode];
 
 /** @public */
@@ -111,6 +116,7 @@ export interface AbstractCursorOptions extends BSONSerializeOptions {
   noCursorTimeout?: boolean;
   /** @internal TODO(NODE-5688): make this public */
   timeoutMS?: number;
+  /** @internal TODO(NODE-5688): make this public */
   timeoutMode?: CursorTimeoutMode;
 }
 
@@ -131,6 +137,12 @@ export type AbstractCursorEvents = {
   [AbstractCursor.CLOSE](): void;
 };
 
+/** @internal */
+export type CursorInitializeOptions = {
+  omitMaxTimeMS?: boolean;
+  timeoutContext?: TimeoutContext;
+};
+
 /** @public */
 export abstract class AbstractCursor<
     TSchema = any,
@@ -161,6 +173,8 @@ export abstract class AbstractCursor<
   private isKilled: boolean;
   /** @internal */
   protected readonly cursorOptions: InternalAbstractCursorOptions;
+  /** @internal */
+  protected timeoutContext?: TimeoutContext;
 
   /** @event */
   static readonly CLOSE = 'close' as const;
@@ -191,11 +205,7 @@ export abstract class AbstractCursor<
     };
     this.cursorOptions.timeoutMS = options.timeoutMS;
     if (this.cursorOptions.timeoutMS != null) {
-      this.cursorOptions.timeoutMode = options.timeoutMode;
-    }
-
-    if (this.cursorOptions.timeoutMS != null) {
-      if (this.cursorOptions.timeoutMode == null) {
+      if (options.timeoutMode == null) {
         if (this.cursorOptions.tailable) {
           this.cursorOptions.timeoutMode = CursorTimeoutMode.ITERATION;
         } else {
@@ -442,6 +452,9 @@ export abstract class AbstractCursor<
       await this.fetchBatch();
     } while (!this.isDead || (this.documents?.length ?? 0) !== 0);
 
+    if (this.cursorOptions.timeoutMode === 'iteration') {
+      this.timeoutContext?.refresh();
+    }
     return null;
   }
 
@@ -493,8 +506,8 @@ export abstract class AbstractCursor<
   /**
    * Frees any client-side resources used by the cursor.
    */
-  async close(): Promise<void> {
-    await this.cleanup();
+  async close(timeoutMS?: number): Promise<void> {
+    await this.cleanup(timeoutMS);
   }
 
   /**
@@ -699,7 +712,8 @@ export abstract class AbstractCursor<
 
   /** @internal */
   protected abstract _initialize(
-    session: ClientSession | undefined
+    session: ClientSession | undefined,
+    options?: CursorInitializeOptions
   ): Promise<InitialCursorResponse>;
 
   /** @internal */
@@ -721,11 +735,12 @@ export abstract class AbstractCursor<
       {
         ...this.cursorOptions,
         session: this.cursorSession,
-        batchSize
+        batchSize,
+        omitMaxTimeMS: this.cursorOptions.timeoutMode != null
       }
     );
 
-    return await executeOperation(this.cursorClient, getMoreOperation);
+    return await executeOperation(this.cursorClient, getMoreOperation, this.timeoutContext);
   }
 
   /**
@@ -736,8 +751,22 @@ export abstract class AbstractCursor<
    * a significant refactor.
    */
   private async cursorInit(): Promise<void> {
+    if (this.cursorOptions.timeoutMS != null) {
+      this.timeoutContext = TimeoutContext.create({
+        serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
+        timeoutMS: this.cursorOptions.timeoutMS,
+        cursorTimeoutMode: this.cursorOptions.timeoutMode
+      });
+    }
+    const omitMaxTimeMS =
+      this.cursorOptions.timeoutMS != null &&
+      ((this.cursorOptions.timeoutMode === 'iteration' && !this.cursorOptions.tailable) ||
+        (this.cursorOptions.tailable && !this.cursorOptions.awaitData));
     try {
-      const state = await this._initialize(this.cursorSession);
+      const state = await this._initialize(this.cursorSession, {
+        timeoutContext: this.timeoutContext,
+        omitMaxTimeMS
+      });
       const response = state.response;
       this.selectedServer = state.server;
       this.cursorId = response.id;
@@ -747,7 +776,7 @@ export abstract class AbstractCursor<
     } catch (error) {
       // the cursor is now initialized, even if an error occurred
       this.initialized = true;
-      await this.cleanup(error);
+      await this.cleanup(undefined, error);
       throw error;
     }
 
@@ -788,7 +817,7 @@ export abstract class AbstractCursor<
       this.documents = response;
     } catch (error) {
       try {
-        await this.cleanup(error);
+        await this.cleanup(undefined, error);
       } catch (error) {
         // `cleanupCursor` should never throw, squash and throw the original error
         squashError(error);
@@ -809,7 +838,7 @@ export abstract class AbstractCursor<
   }
 
   /** @internal */
-  private async cleanup(error?: Error) {
+  private async cleanup(timeoutMS?: number, error?: Error) {
     this.isClosed = true;
     const session = this.cursorSession;
     try {
@@ -824,11 +853,22 @@ export abstract class AbstractCursor<
         this.isKilled = true;
         const cursorId = this.cursorId;
         this.cursorId = Long.ZERO;
+        let timeoutContext: TimeoutContext | undefined;
+        if (timeoutMS != null) {
+          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
         );
       }
     } catch (error) {

src/cursor/aggregation_cursor.ts

@@ -8,7 +8,11 @@ import type { ClientSession } from '../sessions';
 import type { Sort } from '../sort';
 import type { MongoDBNamespace } from '../utils';
 import { mergeOptions } from '../utils';
-import type { AbstractCursorOptions, InitialCursorResponse } from './abstract_cursor';
+import type {
+  AbstractCursorOptions,
+  CursorInitializeOptions,
+  InitialCursorResponse
+} from './abstract_cursor';
 import { AbstractCursor, CursorTimeoutMode } from './abstract_cursor';
 
 /** @public */
@@ -61,14 +65,22 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
   }
 
   /** @internal */
-  async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
+  async _initialize(
+    session: ClientSession,
+    options?: CursorInitializeOptions
+  ): Promise<InitialCursorResponse> {
     const aggregateOperation = new AggregateOperation(this.namespace, this.pipeline, {
       ...this.aggregateOptions,
       ...this.cursorOptions,
+      omitMaxTimeMS: options?.omitMaxTimeMS,
       session
     });
 
-    const response = await executeOperation(this.client, aggregateOperation);
+    const response = await executeOperation(
+      this.client,
+      aggregateOperation,
+      options?.timeoutContext
+    );
 
     return { server: aggregateOperation.server, session, response };
   }

src/cursor/change_stream_cursor.ts

@@ -17,6 +17,7 @@ import { maxWireVersion, type MongoDBNamespace } from '../utils';
 import {
   AbstractCursor,
   type AbstractCursorOptions,
+  type CursorInitializeOptions,
   type InitialCursorResponse
 } from './abstract_cursor';
 
@@ -126,14 +127,22 @@ export class ChangeStreamCursor<
     });
   }
 
-  async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
+  async _initialize(
+    session: ClientSession,
+    options?: CursorInitializeOptions
+  ): Promise<InitialCursorResponse> {
     const aggregateOperation = new AggregateOperation(this.namespace, this.pipeline, {
       ...this.cursorOptions,
       ...this.changeStreamCursorOptions,
+      omitMaxTimeMS: options?.omitMaxTimeMS,
       session
     });
 
-    const response = await executeOperation(session.client, aggregateOperation);
+    const response = await executeOperation(
+      session.client,
+      aggregateOperation,
+      options?.timeoutContext
+    );
 
     const server = aggregateOperation.server;
     this.maxWireVersion = maxWireVersion(server);

src/cursor/find_cursor.ts

@@ -11,7 +11,11 @@ import type { Hint } from '../operations/operation';
 import type { ClientSession } from '../sessions';
 import { formatSort, type Sort, type SortDirection } from '../sort';
 import { emitWarningOnce, mergeOptions, type MongoDBNamespace, squashError } from '../utils';
-import { AbstractCursor, type InitialCursorResponse } from './abstract_cursor';
+import {
+  AbstractCursor,
+  type CursorInitializeOptions,
+  type InitialCursorResponse
+} from './abstract_cursor';
 
 /** @public Flags allowed for cursor */
 export const FLAGS = [
@@ -62,14 +66,18 @@ export class FindCursor<TSchema = any> extends AbstractCursor<TSchema> {
   }
 
   /** @internal */
-  async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
+  async _initialize(
+    session: ClientSession,
+    options?: CursorInitializeOptions
+  ): Promise<InitialCursorResponse> {
     const findOperation = new FindOperation(this.namespace, this.cursorFilter, {
       ...this.findOptions, // NOTE: order matters here, we may need to refine this
       ...this.cursorOptions,
+      omitMaxTimeMS: options?.omitMaxTimeMS,
       session
     });
 
-    const response = await executeOperation(this.client, findOperation);
+    const response = await executeOperation(this.client, findOperation, options?.timeoutContext);
 
     // the response is not a cursor when `explain` is enabled
     this.numReturned = response.batchSize;

src/cursor/list_collections_cursor.ts

@@ -7,7 +7,11 @@ import {
   type ListCollectionsOptions
 } from '../operations/list_collections';
 import type { ClientSession } from '../sessions';
-import { AbstractCursor, type InitialCursorResponse } from './abstract_cursor';
+import {
+  AbstractCursor,
+  type CursorInitializeOptions,
+  type InitialCursorResponse
+} from './abstract_cursor';
 
 /** @public */
 export class ListCollectionsCursor<
@@ -34,14 +38,18 @@ export class ListCollectionsCursor<
   }
 
   /** @internal */
-  async _initialize(session: ClientSession | undefined): Promise<InitialCursorResponse> {
+  async _initialize(
+    session: ClientSession | undefined,
+    options?: CursorInitializeOptions
+  ): Promise<InitialCursorResponse> {
     const operation = new ListCollectionsOperation(this.parent, this.filter, {
       ...this.cursorOptions,
       ...this.options,
+      omitMaxTimeMS: options?.omitMaxTimeMS,
       session
     });
 
-    const response = await executeOperation(this.parent.client, operation);
+    const response = await executeOperation(this.parent.client, operation, options?.timeoutContext);
 
     return { server: operation.server, session, response };
   }

src/cursor/list_indexes_cursor.ts

@@ -2,7 +2,11 @@ import type { Collection } from '../collection';
 import { executeOperation } from '../operations/execute_operation';
 import { ListIndexesOperation, type ListIndexesOptions } from '../operations/indexes';
 import type { ClientSession } from '../sessions';
-import { AbstractCursor, type InitialCursorResponse } from './abstract_cursor';
+import {
+  AbstractCursor,
+  type CursorInitializeOptions,
+  type InitialCursorResponse
+} from './abstract_cursor';
 
 /** @public */
 export class ListIndexesCursor extends AbstractCursor {
@@ -23,10 +27,14 @@ export class ListIndexesCursor extends AbstractCursor {
   }
 
   /** @internal */
-  async _initialize(session: ClientSession | undefined): Promise<InitialCursorResponse> {
+  async _initialize(
+    session: ClientSession | undefined,
+    options?: CursorInitializeOptions
+  ): Promise<InitialCursorResponse> {
     const operation = new ListIndexesOperation(this.parent, {
       ...this.cursorOptions,
       ...this.options,
+      omitMaxTimeMS: options?.omitMaxTimeMS,
       session
     });