Use timeoutcontext for state machine execute()

Author: baileympearson

src/client-side-encryption/state_machine.ts

@@ -3,6 +3,7 @@ import { type MongoCryptContext, type MongoCryptKMSRequest } from 'mongodb-clien
 import * as net from 'net';
 import * as tls from 'tls';
 
+import { type FindOptions } from '../beta';
 import {
   type BSONSerializeOptions,
   deserialize,
@@ -11,6 +12,7 @@ import {
   serialize
 } from '../bson';
 import { type ProxyOptions } from '../cmap/connection';
+import { CursorTimeoutContext } from '../cursor/abstract_cursor';
 import { getSocks, type SocksLib } from '../deps';
 import { MongoOperationTimeoutError } from '../error';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
@@ -524,9 +526,7 @@ export class StateMachine {
       .listCollections(filter, {
         promoteLongs: false,
         promoteValues: false,
-        ...(timeoutContext?.csotEnabled()
-          ? { timeoutMS: timeoutContext?.remainingTimeMS, timeoutMode: 'cursorLifetime' }
-          : {})
+        timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
       })
       .toArray();
 
@@ -582,12 +582,9 @@ export class StateMachine {
     return client
       .db(dbName)
       .collection<DataKey>(collectionName, { readConcern: { level: 'majority' } })
-      .find(
-        deserialize(filter),
-        timeoutContext?.csotEnabled()
-          ? { timeoutMS: timeoutContext?.remainingTimeMS, timeoutMode: 'cursorLifetime' }
-          : {}
-      )
+      .find(deserialize(filter), {
+        timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
+      })
       .toArray();
   }
 }

src/operations/list_collections.ts

@@ -1,6 +1,6 @@
 import type { Binary, Document } from '../bson';
 import { CursorResponse } from '../cmap/wire_protocol/responses';
-import { type CursorTimeoutMode } from '../cursor/abstract_cursor';
+import { type AbstractCursorOptions, type CursorTimeoutMode } from '../cursor/abstract_cursor';
 import type { Db } from '../db';
 import type { Server } from '../sdam/server';
 import type { ClientSession } from '../sessions';
@@ -10,7 +10,9 @@ import { CommandOperation, type CommandOperationOptions } from './command';
 import { Aspect, defineAspects } from './operation';
 
 /** @public */
-export interface ListCollectionsOptions extends Omit<CommandOperationOptions, 'writeConcern'> {
+export interface ListCollectionsOptions
+  extends Omit<CommandOperationOptions, 'writeConcern'>,
+    Pick<AbstractCursorOptions, 'timeoutContext'> {
   /** Since 4.0: If true, will only return the collection name in the response, and will omit additional info */
   nameOnly?: boolean;
   /** Since 4.0: If true and nameOnly is true, allows a user without the required privilege (i.e. listCollections action on the database) to run the command when access control is enforced. */

test/unit/client-side-encryption/state_machine.test.ts

@@ -16,7 +16,10 @@ import {
   BSON,
   Collection,
   CSOTTimeoutContext,
+  CursorTimeoutContext,
+  type FindOptions,
   Int32,
+  type ListCollectionsOptions,
   Long,
   MongoClient,
   serialize,
@@ -484,26 +487,29 @@ describe('StateMachine', function () {
       });
 
       context('when StateMachine.fetchKeys() is passed a `CSOTimeoutContext`', function () {
-        it('collection.find runs with its timeoutMS property set to remainingTimeMS', async function () {
-          const timeoutContext = new CSOTTimeoutContext({
+        it('collection.find uses the provided timeout context', async function () {
+          const context = new CSOTTimeoutContext({
             timeoutMS: 500,
             serverSelectionTimeoutMS: 30000
           });
-          await sleep(300);
+
           await stateMachine
-            .fetchKeys(client, 'keyVault', BSON.serialize({ a: 1 }), timeoutContext)
+            .fetchKeys(client, 'keyVault', BSON.serialize({ a: 1 }), context)
             .catch(e => squashError(e));
-          expect(findSpy.getCalls()[0].args[1].timeoutMS).to.not.be.undefined;
-          expect(findSpy.getCalls()[0].args[1].timeoutMS).to.be.lessThanOrEqual(205);
+
+          const { timeoutContext } = findSpy.getCalls()[0].args[1] as FindOptions;
+          expect(timeoutContext).to.be.instanceOf(CursorTimeoutContext);
+          expect(timeoutContext.timeoutContext).to.equal(context);
         });
       });
 
       context('when StateMachine.fetchKeys() is not passed a `CSOTimeoutContext`', function () {
-        it('collection.find runs with an undefined timeoutMS property', async function () {
+        it('a timeoutContext is not provided to the find cursor', async function () {
           await stateMachine
             .fetchKeys(client, 'keyVault', BSON.serialize({ a: 1 }))
             .catch(e => squashError(e));
-          expect(findSpy.getCalls()[0].args[1].timeoutMS).to.be.undefined;
+          const { timeoutContext } = findSpy.getCalls()[0].args[1] as FindOptions;
+          expect(timeoutContext).to.be.undefined;
         });
       });
     });
@@ -564,29 +570,35 @@ describe('StateMachine', function () {
       context(
         'when StateMachine.fetchCollectionInfo() is passed a `CSOTimeoutContext`',
         function () {
-          it('listCollections runs with its timeoutMS property set to remainingTimeMS', async function () {
-            const timeoutContext = new CSOTTimeoutContext({
+          it('listCollections uses the provided timeoutContext', async function () {
+            const context = new CSOTTimeoutContext({
               timeoutMS: 500,
               serverSelectionTimeoutMS: 30000
             });
-            await sleep(300);
+
             await stateMachine
-              .fetchCollectionInfo(client, 'keyVault', BSON.serialize({ a: 1 }), timeoutContext)
+              .fetchCollectionInfo(client, 'keyVault', BSON.serialize({ a: 1 }), context)
               .catch(e => squashError(e));
-            expect(listCollectionsSpy.getCalls()[0].args[1].timeoutMS).to.not.be.undefined;
-            expect(listCollectionsSpy.getCalls()[0].args[1].timeoutMS).to.be.lessThanOrEqual(205);
+
+            const { timeoutContext } = listCollectionsSpy.getCalls()[0]
+              .args[1] as ListCollectionsOptions;
+            expect(timeoutContext).to.be.instanceOf(CursorTimeoutContext);
+            expect(timeoutContext.timeoutContext).to.equal(context);
           });
         }
       );
 
       context(
         'when StateMachine.fetchCollectionInfo() is not passed a `CSOTimeoutContext`',
         function () {
-          it('listCollections runs with an undefined timeoutMS property', async function () {
+          it('no timeoutContext is provided to listCollections', async function () {
             await stateMachine
               .fetchCollectionInfo(client, 'keyVault', BSON.serialize({ a: 1 }))
               .catch(e => squashError(e));
-            expect(listCollectionsSpy.getCalls()[0].args[1].timeoutMS).to.be.undefined;
+
+            const { timeoutContext } = listCollectionsSpy.getCalls()[0]
+              .args[1] as ListCollectionsOptions;
+            expect(timeoutContext).to.be.undefined;
           });
         }
       );