mongodb 6.12.0-dev.20250125.sha.c1bcf0de → 6.12.0-dev.20250128.sha.654069fc

package/src/client-side-encryption/state_machine.ts

@@ -15,8 +15,15 @@ import { CursorTimeoutContext } from '../cursor/abstract_cursor';
 import { getSocks, type SocksLib } from '../deps';
 import { MongoOperationTimeoutError } from '../error';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
+import { type Abortable } from '../mongo_types';
 import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';
-import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
+import {
+  addAbortListener,
+  BufferPool,
+  kDispose,
+  MongoDBCollectionNamespace,
+  promiseWithResolvers
+} from '../utils';
 import { autoSelectSocketOptions, type DataKey } from './client_encryption';
 import { MongoCryptError } from './errors';
 import { type MongocryptdManager } from './mongocryptd_manager';
@@ -189,7 +196,7 @@ export class StateMachine {
   async execute(
     executor: StateMachineExecutable,
     context: MongoCryptContext,
-    timeoutContext?: TimeoutContext
+    options: { timeoutContext?: TimeoutContext } & Abortable
   ): Promise<Uint8Array> {
     const keyVaultNamespace = executor._keyVaultNamespace;
     const keyVaultClient = executor._keyVaultClient;
@@ -199,6 +206,7 @@ export class StateMachine {
     let result: Uint8Array | null = null;
 
     while (context.state !== MONGOCRYPT_CTX_DONE && context.state !== MONGOCRYPT_CTX_ERROR) {
+      options.signal?.throwIfAborted();
       debug(`[context#${context.id}] ${stateToString.get(context.state) || context.state}`);
 
       switch (context.state) {
@@ -214,7 +222,7 @@ export class StateMachine {
             metaDataClient,
             context.ns,
             filter,
-            timeoutContext
+            options
           );
           if (collInfo) {
             context.addMongoOperationResponse(collInfo);
@@ -235,9 +243,9 @@ export class StateMachine {
           // When we are using the shared library, we don't have a mongocryptd manager.
           const markedCommand: Uint8Array = mongocryptdManager
             ? await mongocryptdManager.withRespawn(
-                this.markCommand.bind(this, mongocryptdClient, context.ns, command, timeoutContext)
+                this.markCommand.bind(this, mongocryptdClient, context.ns, command, options)
               )
-            : await this.markCommand(mongocryptdClient, context.ns, command, timeoutContext);
+            : await this.markCommand(mongocryptdClient, context.ns, command, options);
 
           context.addMongoOperationResponse(markedCommand);
           context.finishMongoOperation();
@@ -246,12 +254,7 @@ export class StateMachine {
 
         case MONGOCRYPT_CTX_NEED_MONGO_KEYS: {
           const filter = context.nextMongoOperation();
-          const keys = await this.fetchKeys(
-            keyVaultClient,
-            keyVaultNamespace,
-            filter,
-            timeoutContext
-          );
+          const keys = await this.fetchKeys(keyVaultClient, keyVaultNamespace, filter, options);
 
           if (keys.length === 0) {
             // See docs on EMPTY_V
@@ -273,7 +276,7 @@ export class StateMachine {
         }
 
         case MONGOCRYPT_CTX_NEED_KMS: {
-          await Promise.all(this.requests(context, timeoutContext));
+          await Promise.all(this.requests(context, options));
           context.finishKMSRequests();
           break;
         }
@@ -315,11 +318,13 @@ export class StateMachine {
    * @param kmsContext - A C++ KMS context returned from the bindings
    * @returns A promise that resolves when the KMS reply has be fully parsed
    */
-  async kmsRequest(request: MongoCryptKMSRequest, timeoutContext?: TimeoutContext): Promise<void> {
+  async kmsRequest(
+    request: MongoCryptKMSRequest,
+    options?: { timeoutContext?: TimeoutContext } & Abortable
+  ): Promise<void> {
     const parsedUrl = request.endpoint.split(':');
     const port = parsedUrl[1] != null ? Number.parseInt(parsedUrl[1], 10) : HTTPS_PORT;
-    const socketOptions = autoSelectSocketOptions(this.options.socketOptions || {});
-    const options: tls.ConnectionOptions & {
+    const socketOptions: tls.ConnectionOptions & {
       host: string;
       port: number;
       autoSelectFamily?: boolean;
@@ -328,7 +333,7 @@ export class StateMachine {
       host: parsedUrl[0],
       servername: parsedUrl[0],
       port,
-      ...socketOptions
+      ...autoSelectSocketOptions(this.options.socketOptions || {})
     };
     const message = request.message;
     const buffer = new BufferPool();
@@ -363,7 +368,7 @@ export class StateMachine {
           throw error;
         }
         try {
-          await this.setTlsOptions(providerTlsOptions, options);
+          await this.setTlsOptions(providerTlsOptions, socketOptions);
         } catch (err) {
           throw onerror(err);
         }
@@ -380,23 +385,25 @@ export class StateMachine {
       .once('close', () => rejectOnNetSocketError(onclose()))
       .once('connect', () => resolveOnNetSocketConnect());
 
+    let abortListener;
+
     try {
       if (this.options.proxyOptions && this.options.proxyOptions.proxyHost) {
         const netSocketOptions = {
+          ...socketOptions,
           host: this.options.proxyOptions.proxyHost,
-          port: this.options.proxyOptions.proxyPort || 1080,
-          ...socketOptions
+          port: this.options.proxyOptions.proxyPort || 1080
         };
         netSocket.connect(netSocketOptions);
         await willConnect;
 
         try {
           socks ??= loadSocks();
-          options.socket = (
+          socketOptions.socket = (
             await socks.SocksClient.createConnection({
               existing_socket: netSocket,
               command: 'connect',
-              destination: { host: options.host, port: options.port },
+              destination: { host: socketOptions.host, port: socketOptions.port },
               proxy: {
                 // host and port are ignored because we pass existing_socket
                 host: 'iLoveJavaScript',
@@ -412,7 +419,7 @@ export class StateMachine {
         }
       }
 
-      socket = tls.connect(options, () => {
+      socket = tls.connect(socketOptions, () => {
         socket.write(message);
       });
 
@@ -422,6 +429,11 @@ export class StateMachine {
         resolve
       } = promiseWithResolvers<void>();
 
+      abortListener = addAbortListener(options?.signal, function () {
+        destroySockets();
+        rejectOnTlsSocketError(this.reason);
+      });
+
       socket
         .once('error', err => rejectOnTlsSocketError(onerror(err)))
         .once('close', () => rejectOnTlsSocketError(onclose()))
@@ -436,8 +448,11 @@ export class StateMachine {
             resolve();
           }
         });
-      await (timeoutContext?.csotEnabled()
-        ? Promise.all([willResolveKmsRequest, Timeout.expires(timeoutContext?.remainingTimeMS)])
+      await (options?.timeoutContext?.csotEnabled()
+        ? Promise.all([
+            willResolveKmsRequest,
+            Timeout.expires(options.timeoutContext?.remainingTimeMS)
+          ])
         : willResolveKmsRequest);
     } catch (error) {
       if (error instanceof TimeoutError)
@@ -446,16 +461,17 @@ export class StateMachine {
     } finally {
       // There's no need for any more activity on this socket at this point.
       destroySockets();
+      abortListener?.[kDispose]();
     }
   }
 
-  *requests(context: MongoCryptContext, timeoutContext?: TimeoutContext) {
+  *requests(context: MongoCryptContext, options?: { timeoutContext?: TimeoutContext } & Abortable) {
     for (
       let request = context.nextKMSRequest();
       request != null;
       request = context.nextKMSRequest()
     ) {
-      yield this.kmsRequest(request, timeoutContext);
+      yield this.kmsRequest(request, options);
     }
   }
 
@@ -516,14 +532,16 @@ export class StateMachine {
     client: MongoClient,
     ns: string,
     filter: Document,
-    timeoutContext?: TimeoutContext
+    options?: { timeoutContext?: TimeoutContext } & Abortable
   ): Promise<Uint8Array | null> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
 
     const cursor = client.db(db).listCollections(filter, {
       promoteLongs: false,
       promoteValues: false,
-      timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
+      timeoutContext:
+        options?.timeoutContext && new CursorTimeoutContext(options?.timeoutContext, Symbol()),
+      signal: options?.signal
     });
 
     // There is always exactly zero or one matching documents, so this should always exhaust the cursor
@@ -547,17 +565,30 @@ export class StateMachine {
     client: MongoClient,
     ns: string,
     command: Uint8Array,
-    timeoutContext?: TimeoutContext
+    options?: { timeoutContext?: TimeoutContext } & Abortable
   ): Promise<Uint8Array> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
     const bsonOptions = { promoteLongs: false, promoteValues: false };
     const rawCommand = deserialize(command, bsonOptions);
 
+    const commandOptions: {
+      timeoutMS?: number;
+      signal?: AbortSignal;
+    } = {
+      timeoutMS: undefined,
+      signal: undefined
+    };
+
+    if (options?.timeoutContext?.csotEnabled()) {
+      commandOptions.timeoutMS = options.timeoutContext.remainingTimeMS;
+    }
+    if (options?.signal) {
+      commandOptions.signal = options.signal;
+    }
+
     const response = await client.db(db).command(rawCommand, {
       ...bsonOptions,
-      ...(timeoutContext?.csotEnabled()
-        ? { timeoutMS: timeoutContext?.remainingTimeMS }
-        : undefined)
+      ...commandOptions
     });
 
     return serialize(response, this.bsonOptions);
@@ -575,17 +606,30 @@ export class StateMachine {
     client: MongoClient,
     keyVaultNamespace: string,
     filter: Uint8Array,
-    timeoutContext?: TimeoutContext
+    options?: { timeoutContext?: TimeoutContext } & Abortable
   ): Promise<Array<DataKey>> {
     const { db: dbName, collection: collectionName } =
       MongoDBCollectionNamespace.fromString(keyVaultNamespace);
 
+    const commandOptions: {
+      timeoutContext?: CursorTimeoutContext;
+      signal?: AbortSignal;
+    } = {
+      timeoutContext: undefined,
+      signal: undefined
+    };
+
+    if (options?.timeoutContext != null) {
+      commandOptions.timeoutContext = new CursorTimeoutContext(options.timeoutContext, Symbol());
+    }
+    if (options?.signal != null) {
+      commandOptions.signal = options.signal;
+    }
+
     return client
       .db(dbName)
       .collection<DataKey>(collectionName, { readConcern: { level: 'majority' } })
-      .find(deserialize(filter), {
-        timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
-      })
+      .find(deserialize(filter), commandOptions)
       .toArray();
   }
 }

package/src/cmap/connection.ts

@@ -33,7 +33,7 @@ import {
 import type { ServerApi, SupportedNodeConnectionOptions } from '../mongo_client';
 import { type MongoClientAuthProviders } from '../mongo_client_auth_providers';
 import { MongoLoggableComponent, type MongoLogger, SeverityLevel } from '../mongo_logger';
-import { type CancellationToken, TypedEventEmitter } from '../mongo_types';
+import { type Abortable, type CancellationToken, TypedEventEmitter } from '../mongo_types';
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import { ServerType } from '../sdam/common';
 import { applySession, type ClientSession, updateSessionFromResponse } from '../sessions';
@@ -438,7 +438,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
 
   private async *sendWire(
     message: WriteProtocolMessageType,
-    options: CommandOptions,
+    options: CommandOptions & Abortable,
     responseType?: MongoDBResponseConstructor
   ): AsyncGenerator<MongoDBResponse> {
     this.throwIfAborted();
@@ -453,7 +453,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       await this.writeCommand(message, {
         agreedCompressor: this.description.compressor ?? 'none',
         zlibCompressionLevel: this.description.zlibCompressionLevel,
-        timeoutContext: options.timeoutContext
+        timeoutContext: options.timeoutContext,
+        signal: options.signal
       });
 
       if (options.noResponse || message.moreToCome) {
@@ -473,7 +474,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
         );
       }
 
-      for await (const response of this.readMany({ timeoutContext: options.timeoutContext })) {
+      for await (const response of this.readMany(options)) {
         this.socket.setTimeout(0);
         const bson = response.parse();
 
@@ -492,9 +493,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   private async *sendCommand(
     ns: MongoDBNamespace,
     command: Document,
-    options: CommandOptions,
+    options: CommandOptions & Abortable,
     responseType?: MongoDBResponseConstructor
   ) {
+    options?.signal?.throwIfAborted();
+
     const message = this.prepareCommand(ns.db, command, options);
     let started = 0;
     if (this.shouldEmitAndLogCommand) {
@@ -610,10 +613,12 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   public async command(
     ns: MongoDBNamespace,
     command: Document,
-    options: CommandOptions = {},
+    options: CommandOptions & Abortable = {},
     responseType?: MongoDBResponseConstructor
   ): Promise<Document> {
     this.throwIfAborted();
+    options.signal?.throwIfAborted();
+
     for await (const document of this.sendCommand(ns, command, options, responseType)) {
       if (options.timeoutContext?.csotEnabled()) {
         if (MongoDBResponse.is(document)) {
@@ -676,7 +681,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       agreedCompressor?: CompressorName;
       zlibCompressionLevel?: number;
       timeoutContext?: TimeoutContext;
-    }
+    } & Abortable
   ): Promise<void> {
     const finalCommand =
       options.agreedCompressor === 'none' || !OpCompressedRequest.canCompress(command)
@@ -701,23 +706,23 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
 
     if (this.socket.write(buffer)) return;
 
-    const drainEvent = once<void>(this.socket, 'drain');
+    const drainEvent = once<void>(this.socket, 'drain', options);
     const timeout = options?.timeoutContext?.timeoutForSocketWrite;
-    if (timeout) {
-      try {
-        return await Promise.race([drainEvent, timeout]);
-      } catch (error) {
-        let err = error;
-        if (TimeoutError.is(error)) {
-          err = new MongoOperationTimeoutError('Timed out at socket write');
-          this.cleanup(err);
-        }
-        throw error;
-      } finally {
-        timeout.clear();
+    const drained = timeout ? Promise.race([drainEvent, timeout]) : drainEvent;
+    try {
+      return await drained;
+    } catch (writeError) {
+      if (TimeoutError.is(writeError)) {
+        const timeoutError = new MongoOperationTimeoutError('Timed out at socket write');
+        this.onError(timeoutError);
+        throw timeoutError;
+      } else if (writeError === options.signal?.reason) {
+        this.onError(writeError);
       }
+      throw writeError;
+    } finally {
+      timeout?.clear();
     }
-    return await drainEvent;
   }
 
   /**
@@ -729,9 +734,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
    *
    * Note that `for-await` loops call `return` automatically when the loop is exited.
    */
-  private async *readMany(options: {
-    timeoutContext?: TimeoutContext;
-  }): AsyncGenerator<OpMsgResponse | OpReply> {
+  private async *readMany(
+    options: {
+      timeoutContext?: TimeoutContext;
+    } & Abortable
+  ): AsyncGenerator<OpMsgResponse | OpReply> {
     try {
       this.dataEvents = onData(this.messageStream, options);
       this.messageStream.resume();
@@ -745,16 +752,17 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
         }
       }
     } catch (readError) {
-      const err = readError;
       if (TimeoutError.is(readError)) {
-        const error = new MongoOperationTimeoutError(
+        const timeoutError = new MongoOperationTimeoutError(
           `Timed out during socket read (${readError.duration}ms)`
         );
         this.dataEvents = null;
-        this.onError(error);
-        throw error;
+        this.onError(timeoutError);
+        throw timeoutError;
+      } else if (readError === options.signal?.reason) {
+        this.onError(readError);
       }
-      throw err;
+      throw readError;
     } finally {
       this.dataEvents = null;
       this.messageStream.pause();

package/src/cmap/connection_pool.ts

@@ -25,10 +25,18 @@ import {
   MongoRuntimeError,
   MongoServerError
 } from '../error';
-import { CancellationToken, TypedEventEmitter } from '../mongo_types';
+import { type Abortable, CancellationToken, TypedEventEmitter } from '../mongo_types';
 import type { Server } from '../sdam/server';
 import { type TimeoutContext, TimeoutError } from '../timeout';
-import { type Callback, List, makeCounter, now, promiseWithResolvers } from '../utils';
+import {
+  addAbortListener,
+  type Callback,
+  kDispose,
+  List,
+  makeCounter,
+  now,
+  promiseWithResolvers
+} from '../utils';
 import { connect } from './connect';
 import { Connection, type ConnectionEvents, type ConnectionOptions } from './connection';
 import {
@@ -316,7 +324,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
    * will be held by the pool. This means that if a connection is checked out it MUST be checked back in or
    * explicitly destroyed by the new owner.
    */
-  async checkOut(options: { timeoutContext: TimeoutContext }): Promise<Connection> {
+  async checkOut(options: { timeoutContext: TimeoutContext } & Abortable): Promise<Connection> {
     const checkoutTime = now();
     this.emitAndLog(
       ConnectionPool.CONNECTION_CHECK_OUT_STARTED,
@@ -334,6 +342,11 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       checkoutTime
     };
 
+    const abortListener = addAbortListener(options.signal, function () {
+      waitQueueMember.cancelled = true;
+      reject(this.reason);
+    });
+
     this.waitQueue.push(waitQueueMember);
     process.nextTick(() => this.processWaitQueue());
 
@@ -364,6 +377,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       }
       throw error;
     } finally {
+      abortListener?.[kDispose]();
       timeout?.clear();
     }
   }

package/src/cmap/wire_protocol/on_data.ts

@@ -1,7 +1,8 @@
 import { type EventEmitter } from 'events';
 
+import { type Abortable } from '../../mongo_types';
 import { type TimeoutContext } from '../../timeout';
-import { List, promiseWithResolvers } from '../../utils';
+import { addAbortListener, kDispose, List, promiseWithResolvers } from '../../utils';
 
 /**
  * @internal
@@ -21,8 +22,10 @@ type PendingPromises = Omit<
  */
 export function onData(
   emitter: EventEmitter,
-  { timeoutContext }: { timeoutContext?: TimeoutContext }
+  { timeoutContext, signal }: { timeoutContext?: TimeoutContext } & Abortable
 ) {
+  signal?.throwIfAborted();
+
   // Setup pending events and pending promise lists
   /**
    * When the caller has not yet called .next(), we store the
@@ -90,6 +93,9 @@ export function onData(
   // Adding event handlers
   emitter.on('data', eventHandler);
   emitter.on('error', errorHandler);
+  const abortListener = addAbortListener(signal, function () {
+    errorHandler(this.reason);
+  });
 
   const timeoutForSocketRead = timeoutContext?.timeoutForSocketRead;
   timeoutForSocketRead?.throwIfExpired();
@@ -115,6 +121,7 @@ export function onData(
     // Adding event handlers
     emitter.off('data', eventHandler);
     emitter.off('error', errorHandler);
+    abortListener?.[kDispose]();
     finished = true;
     timeoutForSocketRead?.clear();
     const doneResult = { value: undefined, done: finished } as const;

package/src/collection.ts

@@ -14,6 +14,7 @@ import type { Db } from './db';
 import { MongoInvalidArgumentError, MongoOperationTimeoutError } from './error';
 import type { MongoClient, PkFactory } from './mongo_client';
 import type {
+  Abortable,
   Filter,
   Flatten,
   OptionalUnlessRequiredId,
@@ -505,7 +506,7 @@ export class Collection<TSchema extends Document = Document> {
   async findOne(filter: Filter<TSchema>): Promise<WithId<TSchema> | null>;
   async findOne(
     filter: Filter<TSchema>,
-    options: Omit<FindOptions, 'timeoutMode'>
+    options: Omit<FindOptions, 'timeoutMode'> & Abortable
   ): Promise<WithId<TSchema> | null>;
 
   // allow an override of the schema.
@@ -513,12 +514,12 @@ export class Collection<TSchema extends Document = Document> {
   async findOne<T = TSchema>(filter: Filter<TSchema>): Promise<T | null>;
   async findOne<T = TSchema>(
     filter: Filter<TSchema>,
-    options?: Omit<FindOptions, 'timeoutMode'>
+    options?: Omit<FindOptions, 'timeoutMode'> & Abortable
   ): Promise<T | null>;
 
   async findOne(
     filter: Filter<TSchema> = {},
-    options: FindOptions = {}
+    options: FindOptions & Abortable = {}
   ): Promise<WithId<TSchema> | null> {
     const cursor = this.find(filter, options).limit(-1).batchSize(1);
     const res = await cursor.next();
@@ -532,9 +533,15 @@ export class Collection<TSchema extends Document = Document> {
    * @param filter - The filter predicate. If unspecified, then all documents in the collection will match the predicate
    */
   find(): FindCursor<WithId<TSchema>>;
-  find(filter: Filter<TSchema>, options?: FindOptions): FindCursor<WithId<TSchema>>;
-  find<T extends Document>(filter: Filter<TSchema>, options?: FindOptions): FindCursor<T>;
-  find(filter: Filter<TSchema> = {}, options: FindOptions = {}): FindCursor<WithId<TSchema>> {
+  find(filter: Filter<TSchema>, options?: FindOptions & Abortable): FindCursor<WithId<TSchema>>;
+  find<T extends Document>(
+    filter: Filter<TSchema>,
+    options?: FindOptions & Abortable
+  ): FindCursor<T>;
+  find(
+    filter: Filter<TSchema> = {},
+    options: FindOptions & Abortable = {}
+  ): FindCursor<WithId<TSchema>> {
     return new FindCursor<WithId<TSchema>>(
       this.client,
       this.s.namespace,
@@ -792,7 +799,7 @@ export class Collection<TSchema extends Document = Document> {
    */
   async countDocuments(
     filter: Filter<TSchema> = {},
-    options: CountDocumentsOptions = {}
+    options: CountDocumentsOptions & Abortable = {}
   ): Promise<number> {
     const pipeline = [];
     pipeline.push({ $match: filter });
@@ -1006,7 +1013,7 @@ export class Collection<TSchema extends Document = Document> {
    */
   aggregate<T extends Document = Document>(
     pipeline: Document[] = [],
-    options?: AggregateOptions
+    options?: AggregateOptions & Abortable
   ): AggregationCursor<T> {
     if (!Array.isArray(pipeline)) {
       throw new MongoInvalidArgumentError(