fix(NODE-5502): recursive calls to next cause memory leak

Author: nbbeeken

src/cursor/abstract_cursor.ts

@@ -129,6 +129,13 @@ export type AbstractCursorEvents = {
   [AbstractCursor.CLOSE](): void;
 };
 
+export interface AbstractCursor {
+  /** @internal */
+  _getMoreAsync: (batchSize: number) => Promise<Document>;
+  /** @internal */
+  kInit: () => Promise<void>;
+}
+
 /** @public */
 export abstract class AbstractCursor<
   TSchema = any,
@@ -220,6 +227,11 @@ export abstract class AbstractCursor<
     return this[kId] ?? undefined;
   }
 
+  /** @internal */
+  get isDead() {
+    return this[kId]?.isZero() ?? true;
+  }
+
   /** @internal */
   get client(): MongoClient {
     return this[kClient];
@@ -632,25 +644,17 @@ export abstract class AbstractCursor<
    * operation.  We cannot refactor to use the abstract _initialize method without
    * a significant refactor.
    */
-  [kInit](callback: Callback<TSchema | null>): void {
+  [kInit](callback: Callback<void>): void {
     this._initialize(this[kSession], (error, state) => {
       if (state) {
         const response = state.response;
         this[kServer] = state.server;
 
         if (response.cursor) {
-          // TODO(NODE-2674): Preserve int64 sent from MongoDB
-          this[kId] =
-            typeof response.cursor.id === 'number'
-              ? Long.fromNumber(response.cursor.id)
-              : typeof response.cursor.id === 'bigint'
-              ? Long.fromBigInt(response.cursor.id)
-              : response.cursor.id;
-
+          this[kId] = getCursorId(response);
           if (response.cursor.ns) {
             this[kNamespace] = ns(response.cursor.ns);
           }
-
           this[kDocuments].pushMany(response.cursor.firstBatch);
         }
 
@@ -671,7 +675,7 @@ export abstract class AbstractCursor<
         return cleanupCursor(this, { error }, () => callback(error, undefined));
       }
 
-      if (cursorIsDead(this)) {
+      if (this.isDead) {
         return cleanupCursor(this, undefined, () => callback());
       }
 
@@ -680,6 +684,22 @@ export abstract class AbstractCursor<
   }
 }
 
+// @ts-expect-error: Callback mutual exclusion type issue
+AbstractCursor.prototype._getMoreAsync = promisify(function (
+  this: AbstractCursor,
+  batchSize: number,
+  callback: Callback<Document>
+) {
+  return this._getMore(batchSize, callback);
+});
+
+AbstractCursor.prototype.kInit = promisify(function (
+  this: AbstractCursor,
+  callback: Callback<void>
+) {
+  return this[kInit](callback);
+});
+
 /**
  * @param cursor - the cursor on which to call `next`
  * @param blocking - a boolean indicating whether or not the cursor should `block` until data
@@ -701,96 +721,58 @@ async function next<T>(
     transform: boolean;
   }
 ): Promise<T | null> {
-  const cursorId = cursor[kId];
   if (cursor.closed) {
     return null;
   }
 
-  if (cursor[kDocuments].length !== 0) {
-    const doc = cursor[kDocuments].shift();
-
-    if (doc != null && transform && cursor[kTransform]) {
-      try {
-        return cursor[kTransform](doc);
-      } catch (error) {
-        await cleanupCursorAsync(cursor, { error, needsToEmitClosed: true }).catch(() => {
-          // `cleanupCursorAsync` should never throw, but if it does we want to throw the original
-          // error instead.
-        });
-        throw error;
-      }
-    }
-
-    return doc;
+  if (cursor[kId] == null) {
+    await cursor.kInit();
   }
 
-  if (cursorId == null) {
-    // All cursors must operate within a session, one must be made implicitly if not explicitly provided
-    const init = promisify(cb => cursor[kInit](cb));
-    await init();
-    return next(cursor, { blocking, transform });
-  }
+  do {
+    if (cursor[kDocuments].length !== 0) {
+      const doc = cursor[kDocuments].shift();
 
-  if (cursorIsDead(cursor)) {
-    // if the cursor is dead, we clean it up
-    // cleanupCursorAsync should never throw, but if it does it indicates a bug in the driver
-    // and we should surface the error
-    await cleanupCursorAsync(cursor, {});
-    return null;
-  }
+      if (doc != null && transform && cursor[kTransform]) {
+        try {
+          return cursor[kTransform](doc);
+        } catch (error) {
+          // `cleanupCursorAsync` should never throw
+          // but if it does we want to throw the original error instead.
+          await cleanupCursorAsync(cursor, { error, needsToEmitClosed: true }).catch(() => null);
+          throw error;
+        }
+      }
 
-  // otherwise need to call getMore
-  const batchSize = cursor[kOptions].batchSize || 1000;
-  const getMore = promisify((batchSize: number, cb: Callback<Document | undefined>) =>
-    cursor._getMore(batchSize, cb)
-  );
-
-  let response: Document | undefined;
-  try {
-    response = await getMore(batchSize);
-  } catch (error) {
-    if (error) {
-      await cleanupCursorAsync(cursor, { error }).catch(() => {
-        // `cleanupCursorAsync` should never throw, but if it does we want to throw the original
-        // error instead.
-      });
-      throw error;
+      return doc;
     }
-  }
-
-  if (response) {
-    const cursorId =
-      typeof response.cursor.id === 'number'
-        ? Long.fromNumber(response.cursor.id)
-        : typeof response.cursor.id === 'bigint'
-        ? Long.fromBigInt(response.cursor.id)
-        : response.cursor.id;
-
-    cursor[kDocuments].pushMany(response.cursor.nextBatch);
-    cursor[kId] = cursorId;
-  }
-
-  if (cursorIsDead(cursor)) {
-    // If we successfully received a response from a cursor BUT the cursor indicates that it is exhausted,
-    // we intentionally clean up the cursor to release its session back into the pool before the cursor
-    // is iterated.  This prevents a cursor that is exhausted on the server from holding
-    // onto a session indefinitely until the AbstractCursor is iterated.
-    //
-    // cleanupCursorAsync should never throw, but if it does it indicates a bug in the driver
-    // and we should surface the error
-    await cleanupCursorAsync(cursor, {});
-  }
 
-  if (cursor[kDocuments].length === 0 && blocking === false) {
-    return null;
-  }
+    if (cursor.isDead) {
+      // if the cursor is dead, we clean it up
+      await cleanupCursorAsync(cursor, {});
+      return null;
+    }
 
-  return next(cursor, { blocking, transform });
-}
+    try {
+      const { batchSize = 1000 } = cursor[kOptions];
+      const response = await cursor._getMoreAsync(batchSize);
+      cursor[kId] = getCursorId(response);
+      cursor[kDocuments].pushMany(response.cursor.nextBatch);
+      if (cursor[kDocuments].length === 0 && blocking === false) {
+        return null;
+      }
+      if (cursor.isDead) {
+        // Received more documents to return but the cursor is done server side.
+        // So release the session as early as possible.
+        await cleanupCursorAsync(cursor, {});
+      }
+    } catch (error) {
+      await cleanupCursorAsync(cursor, { error }).catch(() => null);
+      throw error;
+    }
+  } while (!cursor.isDead || cursor[kDocuments].length !== 0);
 
-function cursorIsDead(cursor: AbstractCursor): boolean {
-  const cursorId = cursor[kId];
-  return !!cursorId && cursorId.isZero();
+  return null;
 }
 
 const cleanupCursorAsync = promisify(cleanupCursor);
@@ -955,3 +937,12 @@ class ReadableCursorStream extends Readable {
     );
   }
 }
+
+function getCursorId(response: { cursor: { id: number | bigint | Long } } | Document) {
+  // TODO(NODE-2674): Preserve int64 sent from MongoDB
+  return typeof response.cursor.id === 'number'
+    ? Long.fromNumber(response.cursor.id)
+    : typeof response.cursor.id === 'bigint'
+    ? Long.fromBigInt(response.cursor.id)
+    : response.cursor.id;
+}