From 48f39f5de9b14f33ac65523534ebeb1406c67880 Mon Sep 17 00:00:00 2001 From: Anna Henningsen Date: Thu, 28 Mar 2024 15:37:23 +0100 Subject: [PATCH 1/2] fix(shell-api): improve AbstractCursor iteration performance MONGOSH-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. --- packages/shell-api/src/abstract-cursor.ts | 23 +++- packages/shell-api/src/cursor.spec.ts | 137 ++++++++++++++++------ 2 files changed, 122 insertions(+), 38 deletions(-) diff --git a/packages/shell-api/src/abstract-cursor.ts b/packages/shell-api/src/abstract-cursor.ts index b2f5e65f7..014d01e82 100644 --- a/packages/shell-api/src/abstract-cursor.ts +++ b/packages/shell-api/src/abstract-cursor.ts @@ -83,6 +83,10 @@ export abstract class AbstractCursor< @returnsPromise async tryNext(): Promise { + return this._tryNext(); + } + + async _tryNext(): Promise { let result = await this._cursor.tryNext(); if (result !== null && this._transform !== null) { result = await this._transform(result); @@ -90,15 +94,27 @@ export abstract class AbstractCursor< 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; } } @@ -114,7 +130,7 @@ export abstract class AbstractCursor< @returnsPromise async itcount(): Promise { let count = 0; - while (await this.tryNext()) { + while (await this._tryNext()) { count++; } return count; @@ -122,6 +138,9 @@ export abstract class AbstractCursor< @returnsPromise async toArray(): Promise { + if (this._canDelegateIterationToUnderlyingCursor()) + return await this._cursor.toArray(); + const result = []; for await (const doc of this) { result.push(doc); diff --git a/packages/shell-api/src/cursor.spec.ts b/packages/shell-api/src/cursor.spec.ts index 9e7363c5c..7cb4ad2de 100644 --- a/packages/shell-api/src/cursor.spec.ts +++ b/packages/shell-api/src/cursor.spec.ts @@ -22,6 +22,14 @@ import { chai.use(sinonChai); const { expect } = chai; +async function allItemsFromAsyncIterable( + iterable: AsyncIterable +): Promise { + const list: T[] = []; + for await (const item of iterable) list.push(item); + return list; +} + describe('Cursor', function () { describe('help', function () { const apiClass = new Cursor( @@ -40,7 +48,7 @@ describe('Cursor', function () { expect(signatures.Cursor.type).to.equal('Cursor'); }); it('map signature', function () { - expect(signatures.Cursor.attributes.map).to.deep.equal({ + expect(signatures.Cursor.attributes?.map).to.deep.equal({ type: 'function', returnsPromise: false, deprecated: false, @@ -107,7 +115,7 @@ describe('Cursor', function () { } as any; describe('#addOption', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -147,7 +155,7 @@ describe('Cursor', function () { describe('#allowPartialResults', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -162,7 +170,7 @@ describe('Cursor', function () { describe('#allowDiskUse', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -187,7 +195,7 @@ describe('Cursor', function () { describe('#batchSize', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -202,22 +210,22 @@ describe('Cursor', function () { describe('#close', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); shellApiCursor = new Cursor(mongo, spCursor); }); - it('closes the cursor', function () { - shellApiCursor.close(); + it('closes the cursor', async function () { + await shellApiCursor.close(); expect(spCursor.close).to.have.been.called; }); }); describe('#collation', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const coll = { locale: 'en' }; beforeEach(function () { @@ -234,7 +242,7 @@ describe('Cursor', function () { describe('#comment', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const cmt = 'hi'; beforeEach(function () { @@ -250,7 +258,7 @@ describe('Cursor', function () { describe('#count', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -271,7 +279,7 @@ describe('Cursor', function () { describe('#hasNext', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -287,7 +295,7 @@ describe('Cursor', function () { describe('#tryNext', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -341,7 +349,7 @@ describe('Cursor', function () { describe('#hint', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const index = 'a_1'; beforeEach(function () { @@ -357,7 +365,7 @@ describe('Cursor', function () { describe('#limit', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = 6; beforeEach(function () { @@ -373,7 +381,7 @@ describe('Cursor', function () { describe('#max', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = { a: 1 }; beforeEach(function () { @@ -389,7 +397,7 @@ describe('Cursor', function () { describe('#maxTimeMS', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = 5000; beforeEach(function () { @@ -405,7 +413,7 @@ describe('Cursor', function () { describe('#maxAwaitTimeMS', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = 5000; beforeEach(function () { @@ -421,7 +429,7 @@ describe('Cursor', function () { describe('#min', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = { a: 1 }; beforeEach(function () { @@ -437,7 +445,7 @@ describe('Cursor', function () { describe('#noCursorTimeout', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -455,7 +463,7 @@ describe('Cursor', function () { describe('#oplogReplay', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -473,7 +481,7 @@ describe('Cursor', function () { describe('#projection', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = { a: 1 }; beforeEach(function () { @@ -489,7 +497,7 @@ describe('Cursor', function () { describe('#readPref', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; let fromOptionsStub; const value = 'primary'; const tagSet = [{ nodeType: 'ANALYTICS' }]; @@ -523,7 +531,7 @@ describe('Cursor', function () { describe('#readConcern', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = 'local'; beforeEach(function () { @@ -541,7 +549,7 @@ describe('Cursor', function () { describe('#returnKey', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = true; beforeEach(function () { @@ -557,7 +565,7 @@ describe('Cursor', function () { describe('#showRecordId', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = true; beforeEach(function () { @@ -573,7 +581,7 @@ describe('Cursor', function () { describe('#objsLeftInBatch', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -589,7 +597,7 @@ describe('Cursor', function () { describe('#skip', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = 6; beforeEach(function () { @@ -605,7 +613,7 @@ describe('Cursor', function () { describe('#sort', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; const value = { a: 1 }; beforeEach(function () { @@ -621,7 +629,7 @@ describe('Cursor', function () { describe('#tailable', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -653,7 +661,7 @@ describe('Cursor', function () { describe('#itcount', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -670,8 +678,8 @@ describe('Cursor', function () { }); describe('#explain', function () { - let nativeCursorStub; - let shellApiCursor; + let nativeCursorStub: StubbedInstance; + let shellApiCursor: Cursor; beforeEach(function () { nativeCursorStub = stubInterface(); @@ -784,7 +792,7 @@ describe('Cursor', function () { describe('#maxScan', function () { let spCursor: StubbedInstance; - let shellApiCursor; + let shellApiCursor: Cursor; beforeEach(function () { spCursor = stubInterface(); @@ -805,8 +813,8 @@ describe('Cursor', function () { }); describe('toShellResult', function () { - let shellApiCursor; - let i; + let shellApiCursor: Cursor; + let i: number; beforeEach(function () { i = 0; @@ -846,5 +854,62 @@ describe('Cursor', function () { expect(result).to.have.nested.property('documents.length', 20); }); }); + + describe('#toArray', function () { + let spCursor: StubbedInstance; + let shellApiCursor: Cursor; + + beforeEach(function () { + spCursor = stubInterface(); + shellApiCursor = new Cursor(mongo, spCursor); + }); + + it('delegates to the underlying cursor if no transform method was specified', async function () { + const docs = [{ a: 1 }, { a: 2 }, { a: 3 }]; + spCursor.toArray.resolves(docs); + const result = await shellApiCursor.toArray(); + expect(result).to.deep.equal(docs); + expect(spCursor.tryNext).to.not.have.been.called; + }); + + it('performs manual iteration if a transform method was specified', async function () { + const docs = [{ a: 1 }, { a: 2 }, { a: 3 }, null]; + spCursor.tryNext.callsFake(() => Promise.resolve(docs.shift())); + shellApiCursor.map(({ a }) => ({ b: a })); + const result = await shellApiCursor.toArray(); + expect(result).to.deep.equal([{ b: 1 }, { b: 2 }, { b: 3 }]); + expect(spCursor.toArray).to.not.have.been.called; + }); + }); + + describe('#async iteration', function () { + let spCursor: StubbedInstance; + let shellApiCursor: Cursor; + + beforeEach(function () { + spCursor = stubInterface(); + shellApiCursor = new Cursor(mongo, spCursor); + }); + + it('delegates to the underlying cursor if no transform method was specified', async function () { + const docs = [{ a: 1 }, { a: 2 }, { a: 3 }]; + // eslint-disable-next-line @typescript-eslint/require-await + spCursor[Symbol.asyncIterator].callsFake(async function* () { + yield* docs; + }); + const result = await allItemsFromAsyncIterable(shellApiCursor); + expect(result).to.deep.equal(docs); + expect(spCursor.tryNext).to.not.have.been.called; + }); + + it('performs manual iteration if a transform method was specified', async function () { + const docs = [{ a: 1 }, { a: 2 }, { a: 3 }, null]; + spCursor.tryNext.callsFake(() => Promise.resolve(docs.shift())); + shellApiCursor.map(({ a }) => ({ b: a })); + const result = await allItemsFromAsyncIterable(shellApiCursor); + expect(result).to.deep.equal([{ b: 1 }, { b: 2 }, { b: 3 }]); + expect(spCursor[Symbol.asyncIterator]).to.not.have.been.called; + }); + }); }); }); From 7284d065b3a4121a5879f260f9d00fe2ab3072f8 Mon Sep 17 00:00:00 2001 From: Anna Henningsen Date: Thu, 28 Mar 2024 16:52:53 +0100 Subject: [PATCH 2/2] fixup: tests --- packages/shell-api/src/abstract-cursor.ts | 8 +++++++- packages/shell-api/src/collection.spec.ts | 3 +-- packages/shell-api/src/database.spec.ts | 3 +-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/packages/shell-api/src/abstract-cursor.ts b/packages/shell-api/src/abstract-cursor.ts index 014d01e82..b3af0ec8a 100644 --- a/packages/shell-api/src/abstract-cursor.ts +++ b/packages/shell-api/src/abstract-cursor.ts @@ -138,8 +138,14 @@ export abstract class AbstractCursor< @returnsPromise async toArray(): Promise { - if (this._canDelegateIterationToUnderlyingCursor()) + // 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) { diff --git a/packages/shell-api/src/collection.spec.ts b/packages/shell-api/src/collection.spec.ts index 85ec1dfcf..0f90746d6 100644 --- a/packages/shell-api/src/collection.spec.ts +++ b/packages/shell-api/src/collection.spec.ts @@ -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([ diff --git a/packages/shell-api/src/database.spec.ts b/packages/shell-api/src/database.spec.ts index 193730ec3..72061775d 100644 --- a/packages/shell-api/src/database.spec.ts +++ b/packages/shell-api/src/database.spec.ts @@ -415,8 +415,7 @@ describe('Database', 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.aggregateDb.returns(serviceProviderCursor); const cursor = await database.aggregate([{ $piplelineStage: {} }]);