diff --git a/packages/shell-api/src/decorators.ts b/packages/shell-api/src/decorators.ts index 39ed6115c5..ecabf2ba04 100644 --- a/packages/shell-api/src/decorators.ts +++ b/packages/shell-api/src/decorators.ts @@ -161,15 +161,21 @@ function wrapWithApiChecks any>(fn: T, className: s markImplicitlyAwaited(async function(this: any, ...args: any[]): Promise { const internalState = getShellInternalState(this); checkForDeprecation(internalState, className, fn); - const interrupted = checkInterrupted(internalState); + const interruptFlag = checkInterrupted(internalState); + const interrupt = interruptFlag?.asPromise(); + let result: any; try { result = await Promise.race([ - interrupted ? interrupted.asPromise() : new Promise(() => {}), + interrupt?.promise ?? new Promise(() => {}), fn.call(this, ...args) ]); } catch (e) { throw rephraseMongoError(e); + } finally { + if (interrupt) { + interrupt.destroy(); + } } checkInterrupted(internalState); return result; diff --git a/packages/shell-api/src/interruptor.spec.ts b/packages/shell-api/src/interruptor.spec.ts index 3c3db1edec..042c552070 100644 --- a/packages/shell-api/src/interruptor.spec.ts +++ b/packages/shell-api/src/interruptor.spec.ts @@ -4,9 +4,48 @@ import { EventEmitter } from 'events'; import { StubbedInstance, stubInterface } from 'ts-sinon'; import Database from './database'; import Mongo from './mongo'; +import { InterruptFlag, MongoshInterruptedError } from './interruptor'; import ShellInternalState from './shell-internal-state'; +import { promisify } from 'util'; describe('interruptor', () => { + describe('InterruptFlag', () => { + let interruptFlag: InterruptFlag; + + beforeEach(() => { + interruptFlag = new InterruptFlag(); + }); + + describe('asPromise', () => { + let interruptPromise: { destroy: () => void; promise: Promise }; + + it('rejects the promise on interrupt', async() => { + interruptPromise = interruptFlag.asPromise(); + let interruptError: MongoshInterruptedError | undefined; + interruptPromise.promise.catch(e => { + interruptError = e; + }); + expect(interruptError).to.be.undefined; + interruptFlag.set(); + await promisify(process.nextTick)(); + expect(interruptError).to.be.instanceOf(MongoshInterruptedError); + }); + + it('rejects immediately if the interrupt happened before', async() => { + interruptFlag.set(); + + interruptPromise = interruptFlag.asPromise(); + let interruptError: MongoshInterruptedError | undefined; + interruptPromise.promise.catch(e => { + interruptError = e; + }); + + await promisify(process.nextTick)(); + expect(interruptError).to.be.instanceOf(MongoshInterruptedError); + }); + }); + }); + describe('with Shell API functions', () => { let mongo: Mongo; let serviceProvider: StubbedInstance; diff --git a/packages/shell-api/src/interruptor.ts b/packages/shell-api/src/interruptor.ts index b3f4dff7e9..40162d0789 100644 --- a/packages/shell-api/src/interruptor.ts +++ b/packages/shell-api/src/interruptor.ts @@ -1,6 +1,8 @@ import { MongoshBaseError } from '@mongosh/errors'; +import { EventEmitter } from 'events'; import ShellInternalState from './shell-internal-state'; +const interruptEvent = 'interrupted'; const kUncatchable = Symbol.for('@@mongosh.uncatchable'); export class MongoshInterruptedError extends MongoshBaseError { @@ -13,14 +15,7 @@ export class MongoshInterruptedError extends MongoshBaseError { export class InterruptFlag { private interrupted = false; - private deferred: { - reject: (e: MongoshInterruptedError) => void; - promise: Promise; - }; - - constructor() { - this.deferred = this.defer(); - } + private onInterrupt = new EventEmitter(); public isSet(): boolean { return this.interrupted; @@ -32,31 +27,35 @@ export class InterruptFlag { * instance of `MongoshInterruptedError`. * @returns Promise that is rejected when the interrupt is set */ - public asPromise(): Promise { - return this.deferred.promise; + public asPromise(): { destroy: () => void; promise: Promise } { + if (this.interrupted) { + return { + destroy: () => {}, + promise: Promise.reject(new MongoshInterruptedError()) + }; + } + + let destroy: (() => void) | undefined; + const promise = new Promise((_, reject) => { + destroy = () => { + this.onInterrupt.removeListener(interruptEvent, reject); + reject(null); + }; + this.onInterrupt.once(interruptEvent, reject); + }); + return { + destroy: destroy as unknown as () => void, + promise + }; } public set(): void { this.interrupted = true; - this.deferred.reject(new MongoshInterruptedError()); + this.onInterrupt.emit(interruptEvent, new MongoshInterruptedError()); } public reset(): void { this.interrupted = false; - this.deferred = this.defer(); - } - - private defer(): { reject: (e: MongoshInterruptedError) => void; promise: Promise; } { - const result: any = {}; - result.promise = new Promise((_, reject) => { - result.reject = reject; - }); - result.promise.catch(() => { - // we ignore the error here - all others should be notified - // we just have to ensure there's at least one handler for it - // to prevent an UnhandledPromiseRejection - }); - return result; } }