feat(node): patch outgoing http requests to capture the zone (#430)

* feat(node): patch outgoing http requests to capture the zone

* feat(node): patch all EventEmitters

Author: alxhub

lib/common/utils.ts

@@ -108,10 +108,11 @@ export function patchOnProperties(obj: any, properties: string[]) {
 };
 
 const EVENT_TASKS = zoneSymbol('eventTasks');
+
+// For EventTarget
 const ADD_EVENT_LISTENER = 'addEventListener';
 const REMOVE_EVENT_LISTENER = 'removeEventListener';
-const SYMBOL_ADD_EVENT_LISTENER = zoneSymbol(ADD_EVENT_LISTENER);
-const SYMBOL_REMOVE_EVENT_LISTENER = zoneSymbol(REMOVE_EVENT_LISTENER);
+
 
 interface ListenerTaskMeta extends TaskData {
   useCapturing: boolean;
@@ -151,83 +152,113 @@ function attachRegisteredEvent(target: any, eventTask: Task): void {
   eventTasks.push(eventTask);
 }
 
-function scheduleEventListener(eventTask: Task): any {
-  const meta = <ListenerTaskMeta>eventTask.data;
-  attachRegisteredEvent(meta.target, eventTask);
-  return meta.target[SYMBOL_ADD_EVENT_LISTENER](meta.eventName, eventTask.invoke,
-      meta.useCapturing);
-}
+export function makeZoneAwareAddListener(addFnName: string, removeFnName: string, useCapturingParam: boolean = true, allowDuplicates: boolean = false) {
+  const addFnSymbol = zoneSymbol(addFnName);
+  const removeFnSymbol = zoneSymbol(removeFnName);
+  const defaultUseCapturing = useCapturingParam ? false : undefined;
 
-function cancelEventListener(eventTask: Task): void {
-  const meta = <ListenerTaskMeta>eventTask.data;
-  findExistingRegisteredTask(meta.target, eventTask.invoke, meta.eventName,
-      meta.useCapturing, true);
-  meta.target[SYMBOL_REMOVE_EVENT_LISTENER](meta.eventName, eventTask.invoke,
-      meta.useCapturing);
-}
-
-function zoneAwareAddEventListener(self: any, args: any[]) {
-  const eventName: string = args[0];
-  const handler: EventListenerOrEventListenerObject = args[1];
-  const useCapturing: boolean = args[2] || false;
-  // - Inside a Web Worker, `this` is undefined, the context is `global`
-  // - When `addEventListener` is called on the global context in strict mode, `this` is undefined
-  // see https://github.com/angular/zone.js/issues/190
-  const target = self || _global;
-  let delegate: EventListener = null;
-  if (typeof handler == 'function') {
-    delegate = <EventListener>handler;
-  } else if (handler && (<EventListenerObject>handler).handleEvent) {
-    delegate = (event) => (<EventListenerObject>handler).handleEvent(event);
-  }
-  var validZoneHandler = false;
-  try {
-    // In cross site contexts (such as WebDriver frameworks like Selenium),
-    // accessing the handler object here will cause an exception to be thrown which
-    // will fail tests prematurely.
-    validZoneHandler = handler && handler.toString() === "[object FunctionWrapper]";
-  } catch(e) {
-    // Returning nothing here is fine, because objects in a cross-site context are unusable
-    return;
+  function scheduleEventListener(eventTask: Task): any {
+    const meta = <ListenerTaskMeta>eventTask.data;
+    attachRegisteredEvent(meta.target, eventTask);
+    return meta.target[addFnSymbol](meta.eventName, eventTask.invoke,
+        meta.useCapturing);
   }
-  // Ignore special listeners of IE11 & Edge dev tools, see https://github.com/angular/zone.js/issues/150
-  if (!delegate || validZoneHandler) {
-    return target[SYMBOL_ADD_EVENT_LISTENER](eventName, handler, useCapturing);
-  }
-  const eventTask: Task
-      = findExistingRegisteredTask(target, handler, eventName, useCapturing, false);
-  if (eventTask) {
-    // we already registered, so this will have noop.
-    return target[SYMBOL_ADD_EVENT_LISTENER](eventName, eventTask.invoke, useCapturing);
+
+  function cancelEventListener(eventTask: Task): void {
+    const meta = <ListenerTaskMeta>eventTask.data;
+    findExistingRegisteredTask(meta.target, eventTask.invoke, meta.eventName,
+        meta.useCapturing, true);
+    meta.target[removeFnSymbol](meta.eventName, eventTask.invoke,
+        meta.useCapturing);
   }
-  const zone: Zone = Zone.current;
-  const source = target.constructor['name'] + '.addEventListener:' + eventName;
-  const data: ListenerTaskMeta = {
-    target: target,
-    eventName: eventName,
-    name: eventName,
-    useCapturing: useCapturing,
-    handler: handler
+
+  return function zoneAwareAddListener(self: any, args: any[]) {
+    const eventName: string = args[0];
+    const handler: EventListenerOrEventListenerObject = args[1];
+    const useCapturing: boolean = args[2] || defaultUseCapturing;
+    // - Inside a Web Worker, `this` is undefined, the context is `global`
+    // - When `addEventListener` is called on the global context in strict mode, `this` is undefined
+    // see https://github.com/angular/zone.js/issues/190
+    const target = self || _global;
+    let delegate: EventListener = null;
+    if (typeof handler == 'function') {
+      delegate = <EventListener>handler;
+    } else if (handler && (<EventListenerObject>handler).handleEvent) {
+      delegate = (event) => (<EventListenerObject>handler).handleEvent(event);
+    }
+    var validZoneHandler = false;
+    try {
+      // In cross site contexts (such as WebDriver frameworks like Selenium),
+      // accessing the handler object here will cause an exception to be thrown which
+      // will fail tests prematurely.
+      validZoneHandler = handler && handler.toString() === "[object FunctionWrapper]";
+    } catch(e) {
+      // Returning nothing here is fine, because objects in a cross-site context are unusable
+      return;
+    }
+    // Ignore special listeners of IE11 & Edge dev tools, see https://github.com/angular/zone.js/issues/150
+    if (!delegate || validZoneHandler) {
+      return target[addFnSymbol](eventName, handler, useCapturing);
+    }
+
+    if (!allowDuplicates) {
+      const eventTask: Task
+          = findExistingRegisteredTask(target, handler, eventName, useCapturing, false);
+      if (eventTask) {
+        // we already registered, so this will have noop.
+        return target[addFnSymbol](eventName, eventTask.invoke, useCapturing);
+      }
+    }
+  
+    const zone: Zone = Zone.current;
+    const source = target.constructor['name'] + '.' + addFnName + ':' + eventName;
+    const data: ListenerTaskMeta = {
+      target: target,
+      eventName: eventName,
+      name: eventName,
+      useCapturing: useCapturing,
+      handler: handler
+    };
+    zone.scheduleEventTask(source, delegate, data, scheduleEventListener, cancelEventListener);
   };
-  zone.scheduleEventTask(source, delegate, data, scheduleEventListener, cancelEventListener);
 }
 
-function zoneAwareRemoveEventListener(self: any, args: any[]) {
-  const eventName: string = args[0];
-  const handler: EventListenerOrEventListenerObject = args[1];
-  const useCapturing: boolean = args[2] || false;
-  // - Inside a Web Worker, `this` is undefined, the context is `global`
-  // - When `addEventListener` is called on the global context in strict mode, `this` is undefined
-  // see https://github.com/angular/zone.js/issues/190
-  const target = self || _global;
-  const eventTask = findExistingRegisteredTask(target, handler, eventName, useCapturing, true);
-  if (eventTask) {
-    eventTask.zone.cancelTask(eventTask);
-  } else {
-    target[SYMBOL_REMOVE_EVENT_LISTENER](eventName, handler, useCapturing);
+export function makeZoneAwareRemoveListener(fnName: string, useCapturingParam: boolean = true) {
+  const symbol = zoneSymbol(fnName);
+  const defaultUseCapturing = useCapturingParam ? false : undefined;
+
+  return function zoneAwareRemoveListener(self: any, args: any[]) {
+    const eventName: string = args[0];
+    const handler: EventListenerOrEventListenerObject = args[1];
+    const useCapturing: boolean = args[2] || defaultUseCapturing;
+    // - Inside a Web Worker, `this` is undefined, the context is `global`
+    // - When `addEventListener` is called on the global context in strict mode, `this` is undefined
+    // see https://github.com/angular/zone.js/issues/190
+    const target = self || _global;
+    const eventTask = findExistingRegisteredTask(target, handler, eventName, useCapturing, true);
+    if (eventTask) {
+      eventTask.zone.cancelTask(eventTask);
+    } else {
+      target[symbol](eventName, handler, useCapturing);
+    }
+  };
+}
+
+export function makeZoneAwareListeners(fnName: string) {
+  const symbol = zoneSymbol(fnName);
+
+  return function zoneAwareEventListeners(self: any, args: any[]) {
+    const eventName: string = args[0];
+    const target = self || _global;
+    return target[EVENT_TASKS]
+      .filter(task => task.data.eventName === eventName)
+      .map(task => task.data.handler);
   }
 }
 
+const zoneAwareAddEventListener = makeZoneAwareAddListener(ADD_EVENT_LISTENER, REMOVE_EVENT_LISTENER);
+const zoneAwareRemoveEventListener = makeZoneAwareRemoveListener(REMOVE_EVENT_LISTENER);
+
 export function patchEventTargetMethods(obj: any): boolean {
   if (obj && obj.addEventListener) {
     patchMethod(obj, ADD_EVENT_LISTENER, () => zoneAwareAddEventListener);
@@ -236,7 +267,8 @@ export function patchEventTargetMethods(obj: any): boolean {
   } else {
     return false;
   }
-};
+}
+
 
 const originalInstanceKey = zoneSymbol('originalInstance');
 

lib/node/events.ts

@@ -0,0 +1,38 @@
+import {makeZoneAwareAddListener, makeZoneAwareListeners, makeZoneAwareRemoveListener, patchMethod} from '../common/utils';
+
+
+// For EventEmitter
+const EE_ADD_LISTENER = 'addListener';
+const EE_PREPEND_LISTENER = 'prependListener';
+const EE_REMOVE_LISTENER = 'removeListener';
+const EE_LISTENERS = 'listeners';
+const EE_ON = 'on';
+
+
+const zoneAwareAddListener = makeZoneAwareAddListener(EE_ADD_LISTENER, EE_REMOVE_LISTENER, false, true);
+const zoneAwarePrependListener = makeZoneAwareAddListener(EE_PREPEND_LISTENER, EE_REMOVE_LISTENER, false, true);
+const zoneAwareRemoveListener = makeZoneAwareRemoveListener(EE_REMOVE_LISTENER, false);
+const zoneAwareListeners = makeZoneAwareListeners(EE_LISTENERS);
+
+export function patchEventEmitterMethods(obj: any): boolean {
+  if (obj && obj.addListener) {
+    patchMethod(obj, EE_ADD_LISTENER, () => zoneAwareAddListener);
+    patchMethod(obj, EE_PREPEND_LISTENER, () => zoneAwarePrependListener);
+    patchMethod(obj, EE_REMOVE_LISTENER, () => zoneAwareRemoveListener);
+    patchMethod(obj, EE_LISTENERS, () => zoneAwareListeners);
+    obj[EE_ON] = obj[EE_ADD_LISTENER];
+    return true;
+  } else {
+    return false;
+  }
+}
+
+// EventEmitter
+let events;
+try {
+  events = require('events');
+} catch (err) {}
+
+if (events && events.EventEmitter) {
+  patchEventEmitterMethods(events.EventEmitter.prototype);
+}

lib/node/node.ts

@@ -1,6 +1,8 @@
 import '../zone';
 import {patchTimer} from '../common/timers';
 
+import './events';
+
 const set = 'set';
 const clear = 'clear';
 const _global = typeof window === 'object' && window || typeof self === 'object' && self || global;
@@ -52,3 +54,21 @@ if (crypto) {
     }
   }.bind(crypto);
 }
+
+// HTTP Client
+let httpClient;
+try {
+  httpClient = require('_http_client');
+} catch (err) {}
+
+if (httpClient && httpClient.ClientRequest) {
+  let ClientRequest = httpClient.ClientRequest.bind(httpClient);
+  httpClient.ClientRequest = function(options: any, callback?: Function) {
+    if (!callback) {
+      return new ClientRequest(options);
+    } else {
+      let zone = Zone.current;
+      return new ClientRequest(options, zone.wrap(callback, 'http.ClientRequest'));
+    }
+  }
+}

test/node/events.spec.ts

@@ -0,0 +1,53 @@
+import {EventEmitter} from 'events';
+
+describe('nodejs EventEmitter', () => {
+  let zone, zoneA, zoneB, emitter, expectZoneACount;
+  beforeEach(() => {
+    zone = Zone.current;
+    zoneA = zone.fork({ name: 'A' });
+    zoneB = zone.fork({ name: 'B' });
+
+    emitter = new EventEmitter();
+    expectZoneACount = 0;
+  });
+
+  function expectZoneA(value) {
+    expectZoneACount++;
+    expect(Zone.current).toBe(zoneA);
+    expect(value).toBe('test value');
+  }
+
+  function shouldNotRun() {
+    fail('this listener should not run');
+  }
+
+  it('should register listeners in the current zone', () => {
+    zoneA.run(() => {
+      emitter.on('test', expectZoneA);
+      emitter.addListener('test', expectZoneA);
+    });
+    zoneB.run(() => emitter.emit('test', 'test value'));
+    expect(expectZoneACount).toBe(2);
+  });
+  it('should remove listeners properly', () => {
+    zoneA.run(() => {
+      emitter.on('test', shouldNotRun);
+      emitter.on('test2', shouldNotRun);
+      emitter.removeListener('test', shouldNotRun);
+    });
+    zoneB.run(() => {
+      emitter.removeListener('test2', shouldNotRun);
+      emitter.emit('test', 'test value');
+      emitter.emit('test2', 'test value');
+    });
+  });
+  it('should return all listeners for an event', () => {
+    zoneA.run(() => {
+      emitter.on('test', expectZoneA);
+    });
+    zoneB.run(() => {
+      emitter.on('test', shouldNotRun);
+    });
+    expect(emitter.listeners('test')).toEqual([expectZoneA, shouldNotRun]);
+  });
+});

test/node_entry_point.ts

@@ -17,4 +17,4 @@ import './test-env-setup';
 
 // List all tests here:
 import './common_tests';
-
+import './node_tests';

test/node_tests.ts

@@ -0,0 +1 @@
+import './node/events.spec';