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 (#1920)

- 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 authored Apr 2, 2024
1 parent 159a2a9 commit cdce3b8
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 42 deletions.
29 changes: 27 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,23 @@ 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[]> {
// toArray is always defined for driver cursors, but not necessarily
// in tests
if (
typeof this._cursor.toArray === 'function' &&
this._canDelegateIterationToUnderlyingCursor()
) {
return await this._cursor.toArray();
}

const result = [];
for await (const doc of this) {
result.push(doc);
Expand Down
3 changes: 1 addition & 2 deletions packages/shell-api/src/collection.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -209,8 +209,7 @@ describe('Collection', function () {

it('returns an AggregationCursor that wraps the service provider one', async function () {
const toArrayResult = [{ foo: 'bar' }];
serviceProviderCursor.tryNext.onFirstCall().resolves({ foo: 'bar' });
serviceProviderCursor.tryNext.onSecondCall().resolves(null);
serviceProviderCursor.toArray.resolves(toArrayResult);
serviceProvider.aggregate.returns(serviceProviderCursor);

const cursor = await collection.aggregate([
Expand Down
Loading

0 comments on commit cdce3b8

Please sign in to comment.