refactor(NODE-6057): implement CursorResponse for lazy document parsing

Author: nbbeeken

src/cmap/wire_protocol/on_demand/document.ts

@@ -197,6 +197,13 @@ export class OnDemandDocument {
     }
   }
 
+  /**
+   * Returns the number of elements in this BSON document
+   */
+  public size() {
+    return this.elements.length;
+  }
+
   /**
    * Checks for the existence of an element by name.
    *
@@ -303,12 +310,18 @@ export class OnDemandDocument {
     });
   }
 
+  /** Returns this document's bytes only */
+  toBytes() {
+    const size = getInt32LE(this.bson, this.offset);
+    return this.bson.subarray(this.offset, this.offset + size);
+  }
+
   /**
    * Iterates through the elements of a document reviving them using the `as` BSONType.
    *
    * @param as - The type to revive all elements as
    */
-  public *valuesAs<const T extends keyof JSTypeOf>(as: T): Generator<JSTypeOf[T]> {
+  public *valuesAs<const T extends keyof JSTypeOf>(as: T): Generator<JSTypeOf[T], void, void> {
     if (!this.isArray) {
       throw new BSONError('Unexpected conversion of non-array value to array');
     }

src/cmap/wire_protocol/responses.ts

@@ -1,5 +1,13 @@
-import { type BSONSerializeOptions, BSONType, type Document, type Timestamp } from '../../bson';
+import {
+  type BSONSerializeOptions,
+  BSONType,
+  type Document,
+  Long,
+  type Timestamp
+} from '../../bson';
+import { MongoUnexpectedServerResponseError } from '../../error';
 import { type ClusterTime } from '../../sdam/common';
+import { type MongoDBNamespace, ns } from '../../utils';
 import { OnDemandDocument } from './on_demand/document';
 
 /** @internal */
@@ -107,3 +115,80 @@ export class MongoDBResponse extends OnDemandDocument {
     return { utf8: { writeErrors: false } };
   }
 }
+
+function throwUnsupportedError() {
+  throw new Error('Unsupported method');
+}
+
+export class CursorResponse extends MongoDBResponse {
+  id: Long | null = null;
+  ns: MongoDBNamespace | null = null;
+
+  documents: any | null = null;
+  bufferForUnshift: any[] = [];
+
+  private batch: OnDemandDocument | null = null;
+  private values: Generator<OnDemandDocument, void, void> | null = null;
+  private batchSize = 0;
+  private iterated = 0;
+
+  constructor(b: Uint8Array, o?: number, a?: boolean) {
+    super(b, o, a);
+
+    if (this.isError) return;
+
+    const cursor = this.get('cursor', BSONType.object, true);
+
+    const id = cursor.get('id', BSONType.long, true);
+    this.id = new Long(Number(id & 0xffff_ffffn), Number((id >> 32n) & 0xffff_ffffn));
+
+    const namespace = cursor.get('ns', BSONType.string) ?? '';
+    if (namespace) this.ns = ns(namespace);
+
+    if (cursor.has('firstBatch')) this.batch = cursor.get('firstBatch', BSONType.array, true);
+    else if (cursor.has('nextBatch')) this.batch = cursor.get('nextBatch', BSONType.array, true);
+    else throw new MongoUnexpectedServerResponseError('Cursor document did not contain a batch');
+
+    this.values = this.batch.valuesAs(BSONType.object);
+    this.batchSize = this.batch.size();
+    this.iterated = 0;
+    this.documents = Object.defineProperties(Object.create(null), {
+      length: {
+        get: () => {
+          return this.batchSize - this.iterated;
+        }
+      },
+      shift: {
+        value: (options?: BSONSerializeOptions) => {
+          this.iterated += 1;
+          if (this.bufferForUnshift.length) return this.bufferForUnshift.pop();
+          const r = this.values?.next();
+          if (!r || r.done) return null;
+          if (options.raw) {
+            return r.value.toBytes();
+          } else {
+            return r.value.toObject(options);
+          }
+        }
+      },
+      unshift: {
+        value: (v: any) => {
+          this.iterated -= 1;
+          this.bufferForUnshift.push(v);
+        }
+      },
+      clear: {
+        value: () => {
+          this.iterated = this.batchSize;
+          this.values?.return();
+        }
+      },
+      pushMany: { value: throwUnsupportedError },
+      push: { value: throwUnsupportedError }
+    });
+  }
+
+  static isCursorResponse(value: unknown): value is CursorResponse {
+    return value instanceof CursorResponse;
+  }
+}

src/cursor/abstract_cursor.ts

@@ -1,6 +1,7 @@
 import { Readable, Transform } from 'stream';
 
 import { type BSONSerializeOptions, type Document, Long, pluckBSONSerializeOptions } from '../bson';
+import { CursorResponse } from '../cmap/wire_protocol/responses';
 import {
   type AnyError,
   MongoAPIError,
@@ -144,7 +145,14 @@ export abstract class AbstractCursor<
   /** @internal */
   [kNamespace]: MongoDBNamespace;
   /** @internal */
-  [kDocuments]: List<TSchema>;
+  [kDocuments]: {
+    length: number;
+    shift(bsonOptions?: any): TSchema | null;
+    unshift(doc: TSchema): void;
+    clear(): void;
+    pushMany(many: Iterable<TSchema>): void;
+    push(item: TSchema): void;
+  };
   /** @internal */
   [kClient]: MongoClient;
   /** @internal */
@@ -286,7 +294,7 @@ export abstract class AbstractCursor<
     const documentsToRead = Math.min(number ?? this[kDocuments].length, this[kDocuments].length);
 
     for (let count = 0; count < documentsToRead; count++) {
-      const document = this[kDocuments].shift();
+      const document = this[kDocuments].shift(this[kOptions]);
       if (document != null) {
         bufferedDocs.push(document);
       }
@@ -633,12 +641,13 @@ export abstract class AbstractCursor<
   protected abstract _initialize(session: ClientSession | undefined): Promise<ExecutionResult>;
 
   /** @internal */
-  async getMore(batchSize: number): Promise<Document | null> {
+  async getMore(batchSize: number, useCursorResponse = false): Promise<Document | null> {
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     const getMoreOperation = new GetMoreOperation(this[kNamespace], this[kId]!, this[kServer]!, {
       ...this[kOptions],
       session: this[kSession],
-      batchSize
+      batchSize,
+      useCursorResponse
     });
 
     return await executeOperation(this[kClient], getMoreOperation);
@@ -656,7 +665,11 @@ export abstract class AbstractCursor<
       const state = await this._initialize(this[kSession]);
       const response = state.response;
       this[kServer] = state.server;
-      if (response.cursor) {
+      if (CursorResponse.isCursorResponse(response)) {
+        this[kId] = response.id;
+        if (response.ns) this[kNamespace] = response.ns;
+        this[kDocuments] = response.documents;
+      } else if (response.cursor) {
         // TODO(NODE-2674): Preserve int64 sent from MongoDB
         this[kId] =
           typeof response.cursor.id === 'number'
@@ -730,7 +743,7 @@ async function next<T>(
     }
 
     if (cursor[kDocuments].length !== 0) {
-      const doc = cursor[kDocuments].shift();
+      const doc = cursor[kDocuments].shift(cursor[kOptions]);
 
       if (doc != null && transform && cursor[kTransform]) {
         try {
@@ -762,8 +775,10 @@ async function next<T>(
 
     try {
       const response = await cursor.getMore(batchSize);
-
-      if (response) {
+      if (CursorResponse.isCursorResponse(response)) {
+        cursor[kId] = response.id;
+        cursor[kDocuments] = response.documents;
+      } else if (response) {
         const cursorId =
           typeof response.cursor.id === 'number'
             ? Long.fromNumber(response.cursor.id)

src/cursor/find_cursor.ts

@@ -111,10 +111,10 @@ export class FindCursor<TSchema = any> extends AbstractCursor<TSchema> {
       }
     }
 
-    const response = await super.getMore(batchSize);
+    const response = await super.getMore(batchSize, true);
     // TODO: wrap this in some logic to prevent it from happening if we don't need this support
     if (response) {
-      this[kNumReturned] = this[kNumReturned] + response.cursor.nextBatch.length;
+      this[kNumReturned] = this[kNumReturned] + response.batchLength;
     }
 
     return response;

src/operations/execute_operation.ts

@@ -1,4 +1,5 @@
 import type { Document } from '../bson';
+import { type CursorResponse } from '../cmap/wire_protocol/responses';
 import {
   isRetryableReadError,
   isRetryableWriteError,
@@ -44,7 +45,7 @@ export interface ExecutionResult {
   /** The session used for this operation, may be implicitly created */
   session?: ClientSession;
   /** The raw server response for the operation */
-  response: Document;
+  response: Document | CursorResponse;
 }
 
 /**

src/operations/find.ts

@@ -1,4 +1,5 @@
 import type { Document } from '../bson';
+import { CursorResponse } from '../cmap/wire_protocol/responses';
 import type { Collection } from '../collection';
 import { MongoInvalidArgumentError } from '../error';
 import { ReadConcern } from '../read_concern';
@@ -111,12 +112,17 @@ export class FindOperation extends CommandOperation<Document> {
       findCommand = decorateWithExplain(findCommand, this.explain);
     }
 
-    return await server.command(this.ns, findCommand, {
-      ...this.options,
-      ...this.bsonOptions,
-      documentsReturnedIn: 'firstBatch',
-      session
-    });
+    return await server.command(
+      this.ns,
+      findCommand,
+      {
+        ...this.options,
+        ...this.bsonOptions,
+        documentsReturnedIn: 'firstBatch',
+        session
+      },
+      this.explain ? undefined : CursorResponse
+    );
   }
 }
 

src/operations/get_more.ts

@@ -1,4 +1,5 @@
 import type { Document, Long } from '../bson';
+import { CursorResponse } from '../cmap/wire_protocol/responses';
 import { MongoRuntimeError } from '../error';
 import type { Server } from '../sdam/server';
 import type { ClientSession } from '../sessions';
@@ -19,6 +20,8 @@ export interface GetMoreOptions extends OperationOptions {
   maxTimeMS?: number;
   /** TODO(NODE-4413): Address bug with maxAwaitTimeMS not being passed in from the cursor correctly */
   maxAwaitTimeMS?: number;
+
+  useCursorResponse: boolean;
 }
 
 /**
@@ -96,7 +99,12 @@ export class GetMoreOperation extends AbstractOperation {
       ...this.options
     };
 
-    return await server.command(this.ns, getMoreCmd, commandOptions);
+    return await server.command(
+      this.ns,
+      getMoreCmd,
+      commandOptions,
+      this.options.useCursorResponse ? CursorResponse : undefined
+    );
   }
 }
 

src/sdam/server.ts

@@ -7,6 +7,7 @@ import {
   type ConnectionPoolOptions
 } from '../cmap/connection_pool';
 import { PoolClearedError } from '../cmap/errors';
+import { type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses';
 import {
   APM_EVENTS,
   CLOSED,
@@ -262,11 +263,25 @@ export class Server extends TypedEventEmitter<ServerEvents> {
     }
   }
 
-  /**
-   * Execute a command
-   * @internal
-   */
-  async command(ns: MongoDBNamespace, cmd: Document, options: CommandOptions): Promise<Document> {
+  public async command<T extends MongoDBResponseConstructor>(
+    ns: MongoDBNamespace,
+    command: Document,
+    options: CommandOptions | undefined,
+    responseType: T | undefined
+  ): Promise<typeof responseType extends undefined ? Document : InstanceType<T>>;
+
+  public async command(
+    ns: MongoDBNamespace,
+    command: Document,
+    options?: CommandOptions
+  ): Promise<Document>;
+
+  public async command(
+    ns: MongoDBNamespace,
+    cmd: Document,
+    options: CommandOptions,
+    responseType?: MongoDBResponseConstructor
+  ): Promise<Document> {
     if (ns.db == null || typeof ns === 'string') {
       throw new MongoInvalidArgumentError('Namespace must not be a string');
     }
@@ -308,7 +323,7 @@ export class Server extends TypedEventEmitter<ServerEvents> {
 
     try {
       try {
-        return await conn.command(ns, cmd, finalOptions);
+        return await conn.command(ns, cmd, finalOptions, responseType);
       } catch (commandError) {
         throw this.decorateCommandError(conn, cmd, finalOptions, commandError);
       }
@@ -319,7 +334,7 @@ export class Server extends TypedEventEmitter<ServerEvents> {
       ) {
         await this.pool.reauthenticate(conn);
         try {
-          return await conn.command(ns, cmd, finalOptions);
+          return await conn.command(ns, cmd, finalOptions, responseType);
         } catch (commandError) {
           throw this.decorateCommandError(conn, cmd, finalOptions, commandError);
         }