Skip to content

Commit cdce3b8

Browse files
authored
fix(shell-api): improve AbstractCursor iteration performance MONGOSH-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.
1 parent 159a2a9 commit cdce3b8

File tree

4 files changed

+130
-42
lines changed

4 files changed

+130
-42
lines changed

packages/shell-api/src/abstract-cursor.ts

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,22 +83,38 @@ export abstract class AbstractCursor<
8383

8484
@returnsPromise
8585
async tryNext(): Promise<Document | null> {
86+
return this._tryNext();
87+
}
88+
89+
async _tryNext(): Promise<Document | null> {
8690
let result = await this._cursor.tryNext();
8791
if (result !== null && this._transform !== null) {
8892
result = await this._transform(result);
8993
}
9094
return result;
9195
}
9296

97+
_canDelegateIterationToUnderlyingCursor(): boolean {
98+
return this._transform === null;
99+
}
100+
93101
get [Symbol.for('@@mongosh.syntheticAsyncIterable')]() {
94102
return true;
95103
}
96104

97105
async *[Symbol.asyncIterator]() {
106+
if (
107+
this._cursor[Symbol.asyncIterator] &&
108+
this._canDelegateIterationToUnderlyingCursor()
109+
) {
110+
yield* this._cursor;
111+
return;
112+
}
113+
98114
let doc;
99115
// !== null should suffice, but some stubs in our tests return 'undefined'
100116
// eslint-disable-next-line eqeqeq
101-
while ((doc = await this.tryNext()) != null) {
117+
while ((doc = await this._tryNext()) != null) {
102118
yield doc;
103119
}
104120
}
@@ -114,14 +130,23 @@ export abstract class AbstractCursor<
114130
@returnsPromise
115131
async itcount(): Promise<number> {
116132
let count = 0;
117-
while (await this.tryNext()) {
133+
while (await this._tryNext()) {
118134
count++;
119135
}
120136
return count;
121137
}
122138

123139
@returnsPromise
124140
async toArray(): Promise<Document[]> {
141+
// toArray is always defined for driver cursors, but not necessarily
142+
// in tests
143+
if (
144+
typeof this._cursor.toArray === 'function' &&
145+
this._canDelegateIterationToUnderlyingCursor()
146+
) {
147+
return await this._cursor.toArray();
148+
}
149+
125150
const result = [];
126151
for await (const doc of this) {
127152
result.push(doc);

packages/shell-api/src/collection.spec.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -209,8 +209,7 @@ describe('Collection', function () {
209209

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

216215
const cursor = await collection.aggregate([

0 commit comments

Comments
 (0)