Skip to content

Commit

Permalink
fix(shell-api): improve AbstractCursor iteration performance MONGOSH-…
Browse files Browse the repository at this point in the history
…1688

- Delegate cursor iteration to the driver where possible
- Do not call the wrapped `.tryNext()` method from inside other
  shell API methods, and instead only call an internal unwrapped
  version of it (similar to `db.runCommand()` vs `db._runCommand()`

This results in a 60% improvement of runtime in local testing on the
`db_cursor_iteration_plainvm` benchmark.
  • Loading branch information
addaleax committed Mar 28, 2024
1 parent 276de8a commit 48f39f5
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 38 deletions.
23 changes: 21 additions & 2 deletions packages/shell-api/src/abstract-cursor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -83,22 +83,38 @@ export abstract class AbstractCursor<

@returnsPromise
async tryNext(): Promise<Document | null> {
return this._tryNext();
}

async _tryNext(): Promise<Document | null> {
let result = await this._cursor.tryNext();
if (result !== null && this._transform !== null) {
result = await this._transform(result);
}
return result;
}

_canDelegateIterationToUnderlyingCursor(): boolean {
return this._transform === null;
}

get [Symbol.for('@@mongosh.syntheticAsyncIterable')]() {
return true;
}

async *[Symbol.asyncIterator]() {
if (
this._cursor[Symbol.asyncIterator] &&
this._canDelegateIterationToUnderlyingCursor()
) {
yield* this._cursor;
return;
}

let doc;
// !== null should suffice, but some stubs in our tests return 'undefined'
// eslint-disable-next-line eqeqeq
while ((doc = await this.tryNext()) != null) {
while ((doc = await this._tryNext()) != null) {
yield doc;
}
}
Expand All @@ -114,14 +130,17 @@ export abstract class AbstractCursor<
@returnsPromise
async itcount(): Promise<number> {
let count = 0;
while (await this.tryNext()) {
while (await this._tryNext()) {
count++;
}
return count;
}

@returnsPromise
async toArray(): Promise<Document[]> {
if (this._canDelegateIterationToUnderlyingCursor())
return await this._cursor.toArray();

const result = [];
for await (const doc of this) {
result.push(doc);
Expand Down
Loading

0 comments on commit 48f39f5

Please sign in to comment.