diff --git a/lib/common/promise.ts b/lib/common/promise.ts index ce765a260..7558349a1 100644 --- a/lib/common/promise.ts +++ b/lib/common/promise.ts @@ -207,6 +207,9 @@ Zone.__load_patch('ZoneAwarePromise', (global: any, Zone: ZoneType, api: _ZonePr api.scheduleMicroTask(); // to make sure that it is running } } + if ((promise as any)['__zone_symbol__isAsync'] === true && (promise as any)['__zone_symbol__outsideAsync'] !== true) { + Zone.setAsyncFrame(); + } } } // Resolving an already resolved promise is a noop. @@ -450,9 +453,16 @@ Zone.__load_patch('ZoneAwarePromise', (global: any, Zone: ZoneType, api: _ZonePr proto[symbolThen] = originalThen; Ctor.prototype.then = function(onResolve: any, onReject: any) { + let isNativePromise = false; + if (this && !(this instanceof ZoneAwarePromise)) { + isNativePromise = true; + } const wrapped = new ZoneAwarePromise((resolve, reject) => { originalThen.call(this, resolve, reject); }); + if (isNativePromise) { + (wrapped as any)['__zone_symbol__isAsync'] = true; + } return wrapped.then(onResolve, onReject); }; (Ctor as any)[symbolThenPatched] = true; diff --git a/lib/zone.ts b/lib/zone.ts index 4de4e077d..656e2a7c7 100644 --- a/lib/zone.ts +++ b/lib/zone.ts @@ -307,6 +307,9 @@ interface ZoneType { /** @internal */ __symbol__(name: string): string; + + /** @internal */ + setAsyncFrame(): void; } /** @internal */ @@ -661,6 +664,11 @@ const Zone: ZoneType = (function(global: any) { } } + const detectAsyncFunction = async function() {}; + const AsyncFunction = (detectAsyncFunction.constructor as any).name === 'AsyncFunction' ? + detectAsyncFunction.constructor : + null; + class Zone implements AmbientZone { static __symbol__: (name: string) => string = __symbol__; @@ -691,6 +699,10 @@ const Zone: ZoneType = (function(global: any) { return _currentTask; } + static setAsyncFrame() { + _currentZoneFrame = _asyncZoneFrame!; + } + static __load_patch(name: string, fn: _PatchFn): void { if (patches.hasOwnProperty(name)) { throw Error('Already loaded patch: ' + name); @@ -761,9 +773,27 @@ const Zone: ZoneType = (function(global: any) { callback: (...args: any[]) => T, applyThis?: any, applyArgs?: any[], source?: string): T { _currentZoneFrame = {parent: _currentZoneFrame, zone: this}; try { + if (callback && callback.constructor === AsyncFunction) { + const r = this._zoneDelegate.invoke(this, callback, applyThis, applyArgs, source); + if (r && typeof r.then === 'function') { + r['__zone_symbol__outsideAsync'] = true; + _asyncZoneFrame = _currentZoneFrame; + return r.then((result: any) => { + _currentZoneFrame = _asyncZoneFrame!.parent!; + _isAsyncSet = true; + _asyncZoneFrame = null; + return result; + }); + } + return r; + } return this._zoneDelegate.invoke(this, callback, applyThis, applyArgs, source); } finally { - _currentZoneFrame = _currentZoneFrame.parent!; + if (!_isAsyncSet) { + _currentZoneFrame = _currentZoneFrame.parent!; + } else { + _isAsyncSet = false; + } } } @@ -1353,6 +1383,8 @@ const Zone: ZoneType = (function(global: any) { }, }; let _currentZoneFrame: _ZoneFrame = {parent: null, zone: new Zone(null, null)}; + let _asyncZoneFrame: _ZoneFrame|null = null; + let _isAsyncSet = false; let _currentTask: Task|null = null; let _numberOfNestedTaskFrames = 0;