Start CSOT impl

Author: W-A-James

src/bulk/common.ts

@@ -20,6 +20,7 @@ import { makeUpdateStatement, UpdateOperation, type UpdateStatement } from '../o
 import type { Server } from '../sdam/server';
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
+import { type Timeout } from '../timeout';
 import { maybeAddIdToDocuments } from '../utils';
 import {
   applyRetryableWrites,
@@ -874,6 +875,8 @@ export interface BulkWriteOptions extends CommandOperationOptions {
   forceServerObjectId?: boolean;
   /** Map of parameter names and values that can be accessed using $$var (requires MongoDB 5.0). */
   let?: Document;
+  /** @internal */
+  timeout?: Timeout | null;
 }
 
 const executeCommandsAsync = promisify(executeCommands);

src/cmap/connection.ts

@@ -29,7 +29,7 @@ import { 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';
-import { type Timeout } from '../timeout';
+import { Timeout } from '../timeout';
 import {
   BufferPool,
   calculateDurationInMs,
@@ -59,6 +59,7 @@ import {
   type WriteProtocolMessageType
 } from './commands';
 import type { Stream } from './connect';
+import { type ConnectionPool } from './connection_pool';
 import type { ClientMetadata } from './handshake/client_metadata';
 import { StreamDescription, type StreamDescriptionOptions } from './stream_description';
 import { type CompressorName, decompressResponse } from './wire_protocol/compression';
@@ -183,6 +184,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
    * Once connection is established, command logging can log events (if enabled)
    */
   public established: boolean;
+  public pool?: ConnectionPool;
   /** Indicates that the connection (including underlying TCP socket) has been closed. */
   public closed = false;
 
@@ -279,6 +281,10 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     );
   }
 
+  private get minRoundTripTime(): number {
+    return this.pool?.server.description.minRoundTripTime ?? 0;
+  }
+
   public markAvailable(): void {
     this.lastUseTime = now();
   }
@@ -343,6 +349,10 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
 
     let clusterTime = this.clusterTime;
 
+    if (Timeout.is(options.timeout) && options.timeout.duration > 0) {
+      cmd.maxTimeMS = options.timeout.getMaxTimeMS(this.minRoundTripTime);
+    }
+
     if (this.serverApi) {
       const { version, strict, deprecationErrors } = this.serverApi;
       cmd.apiVersion = version;
@@ -432,7 +442,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     try {
       await this.writeCommand(message, {
         agreedCompressor: this.description.compressor ?? 'none',
-        zlibCompressionLevel: this.description.zlibCompressionLevel
+        zlibCompressionLevel: this.description.zlibCompressionLevel,
+        timeout: options.timeout
       });
 
       if (options.noResponse) {
@@ -442,7 +453,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
 
       this.throwIfAborted();
 
-      for await (const response of this.readMany()) {
+      for await (const response of this.readMany({ timeout: options.timeout })) {
         this.socket.setTimeout(0);
         const bson = response.parse();
 
@@ -635,7 +646,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
    */
   private async writeCommand(
     command: WriteProtocolMessageType,
-    options: { agreedCompressor?: CompressorName; zlibCompressionLevel?: number }
+    options: {
+      agreedCompressor?: CompressorName;
+      zlibCompressionLevel?: number;
+      timeout?: Timeout | null;
+    }
   ): Promise<void> {
     const finalCommand =
       options.agreedCompressor === 'none' || !OpCompressedRequest.canCompress(command)
@@ -647,8 +662,15 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
 
     const buffer = Buffer.concat(await finalCommand.toBin());
 
-    if (this.socket.write(buffer)) return;
-    return await once(this.socket, 'drain');
+    if (this.socket.write(buffer)) {
+      return;
+    }
+    const drain = once(this.socket, 'drain');
+
+    if (options.timeout) {
+      await Promise.race([drain, options.timeout]);
+    }
+    await drain;
   }
 
   /**
@@ -660,9 +682,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
    *
    * Note that `for-await` loops call `return` automatically when the loop is exited.
    */
-  private async *readMany(): AsyncGenerator<OpMsgResponse | OpReply> {
+  private async *readMany(options: {
+    timeout?: Timeout | null;
+  }): AsyncGenerator<OpMsgResponse | OpReply> {
     try {
-      this.dataEvents = onData(this.messageStream);
+      this.dataEvents = onData(this.messageStream, options);
       for await (const message of this.dataEvents) {
         const response = await decompressResponse(message);
         yield response;

src/cmap/connection_pool.ts

@@ -168,6 +168,8 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
   [kMetrics]: ConnectionPoolMetrics;
   [kProcessingWaitQueue]: boolean;
 
+  server: Server;
+
   /**
    * Emitted when the connection pool is created.
    * @event
@@ -247,6 +249,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
 
     this[kPoolState] = PoolState.paused;
     this[kServer] = server;
+    this.server = server;
     this[kConnections] = new List();
     this[kPending] = 0;
     this[kCheckedOut] = new Set();

src/operations/execute_operation.ts

@@ -187,6 +187,9 @@ export async function executeOperation<
   }
 
   try {
+    if (operation.timeout) {
+      return await Promise.race([operation.execute(server, session), operation.timeout]);
+    }
     return await operation.execute(server, session);
   } catch (operationError) {
     if (willRetry && operationError instanceof MongoError) {
@@ -271,6 +274,9 @@ async function retryOperation<
   }
 
   try {
+    if (operation.timeout) {
+      return await Promise.race([operation.execute(server, session), operation.timeout]);
+    }
     return await operation.execute(server, session);
   } catch (retryError) {
     if (

src/timeout.ts

@@ -24,7 +24,10 @@ export class CSOTError extends MongoError {
   }
 }
 
-/** @internal */
+/** @internal
+ * This class is an abstraction over CSOT timeouts, implementing the specification outlined in
+ * https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/client-side-operations-timeout.md
+ * */
 export class Timeout extends Promise<never> {
   get [Symbol.toStringTag](): 'MongoDBTimeout' {
     return 'MongoDBTimeout';
@@ -39,6 +42,9 @@ export class Timeout extends Promise<never> {
   public duration: number;
   public timedOut = false;
 
+  /**
+   * Return the amount of time remaining until a CSOTError is thrown
+   * */
   public get remainingTime(): number {
     if (this.duration === 0) return Infinity;
     if (this.timedOut) return 0;
@@ -59,13 +65,13 @@ export class Timeout extends Promise<never> {
       executor(noop, promiseReject);
     });
 
-    // Construct timeout error at point of Timeout instantiation to preserve stack traces
+    // NOTE: Construct timeout error at point of Timeout instantiation to preserve stack traces
     this.timeoutError = new CSOTError('Timeout!');
 
     this.expireTimeout = () => {
       this.ended = Math.trunc(performance.now());
       this.timedOut = true;
-      // Wrap error here: Why?
+      // NOTE: Wrap error here: Why?
       reject(CSOTError.from(this.timeoutError));
     };
 
@@ -108,7 +114,12 @@ export class Timeout extends Promise<never> {
     }
   }
 
+  /**
+   * Implement maxTimeMS calculation detailed in https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/client-side-operations-timeout.md#command-execution
+   * */
   public getMaxTimeMS(minRoundTripTime: number): any {
+    console.log(`remaining time: ${this.remainingTime}, minRTT: ${minRoundTripTime}`);
+    if (!Number.isFinite(this.remainingTime)) return 0;
     if (minRoundTripTime < this.remainingTime) return this.remainingTime - minRoundTripTime;
     throw CSOTError.from(this.timeoutError);
   }
@@ -118,6 +129,7 @@ export class Timeout extends Promise<never> {
     return Timeout.expires(this.duration);
   }
 
+  /** Create a new timeout that expires in `duration` ms */
   public static expires(duration: number): Timeout {
     return new Timeout(undefined, duration);
   }