feature: add before/after InvocationRequest hooks (#339)

* feature: add before/after InvocationRequest hooks

* rm return workerChannel

* refactor to register cb pattern

* remove double callback on promise resolve

Author: Mark Wolff

src/WorkerChannel.ts

@@ -7,9 +7,13 @@ import { toTypedData } from './converters';
 import { augmentTriggerMetadata } from './augmenters';
 import { systemError } from './utils/Logger';
 import { InternalException } from './utils/InternalException';
+import { Context } from './public/Interfaces';
 import LogCategory = rpc.RpcLog.RpcLogCategory;
 import LogLevel = rpc.RpcLog.Level;
 
+type InvocationRequestBefore = (context: Context, userFn: Function) => Function;
+type InvocationRequestAfter = (context: Context) => void;
+
 /**
  * The worker channel should have a way to handle all incoming gRPC messages.
  * This includes all incoming StreamingMessage types (exclude *Response types and RpcLog type)
@@ -25,6 +29,8 @@ interface IWorkerChannel {
   invocationRequest(requestId: string, msg: rpc.InvocationRequest): void;
   invocationCancel(requestId: string, msg: rpc.InvocationCancel): void;
   functionEnvironmentReloadRequest(requestId: string, msg: rpc.IFunctionEnvironmentReloadRequest): void;
+  registerBeforeInvocationRequest(beforeCb: InvocationRequestBefore): void;
+  registerAfterInvocationRequest(afterCb: InvocationRequestAfter): void;
 }
 
 /**
@@ -34,11 +40,15 @@ export class WorkerChannel implements IWorkerChannel {
   private _eventStream: IEventStream;
   private _functionLoader: IFunctionLoader;
   private _workerId: string;
+  private _invocationRequestBefore: InvocationRequestBefore[];
+  private _invocationRequestAfter: InvocationRequestAfter[];
 
   constructor(workerId: string, eventStream: IEventStream, functionLoader: IFunctionLoader) {
     this._workerId = workerId;
     this._eventStream = eventStream;
     this._functionLoader = functionLoader;
+    this._invocationRequestBefore = [];
+    this._invocationRequestAfter = [];
 
     // call the method with the matching 'event' name on this class, passing the requestId and event message
     eventStream.on('data', (msg) => {
@@ -82,6 +92,21 @@ export class WorkerChannel implements IWorkerChannel {
     });
   }
 
+  /**
+   * Register a patching function to be run before User Function is executed.
+   * Hook should return a patched version of User Function.
+   */
+  public registerBeforeInvocationRequest(beforeCb: InvocationRequestBefore): void {
+    this._invocationRequestBefore.push(beforeCb);
+  }
+
+  /**
+   * Register a function to be run after User Function resolves.
+   */
+  public registerAfterInvocationRequest(afterCb: InvocationRequestAfter): void {
+    this._invocationRequestAfter.push(afterCb);
+  }
+
   /**
    * Host sends capabilities/init data to worker and requests the worker to initialize itself 
    * @param requestId gRPC message request id
@@ -160,7 +185,7 @@ export class WorkerChannel implements IWorkerChannel {
         invocationId: msg.invocationId,
         result: this.getStatus(err)
       }
-
+      
       try {
         if (result) {
           if (result.return) {
@@ -183,19 +208,27 @@ export class WorkerChannel implements IWorkerChannel {
         requestId: requestId,
         invocationResponse: response
       });
+      
+      this.runInvocationRequestAfter(context);
     }
 
     let { context, inputs } = CreateContextAndInputs(info, msg, logCallback, resultCallback);
     let userFunction = this._functionLoader.getFunc(<string>msg.functionId);
 
+    userFunction = this.runInvocationRequestBefore(context, userFunction);
+    
     // catch user errors from the same async context in the event loop and correlate with invocation
     // throws from asynchronous work (setTimeout, etc) are caught by 'unhandledException' and cannot be correlated with invocation
     try {
-      let result = userFunction(context, ...inputs);
+        let result = userFunction(context, ...inputs);
 
-      if (result && isFunction(result.then)) {
-        result.then(result => (<any>context.done)(null, result, true))
-          .catch(err => (<any>context.done)(err, null, true));
+        if (result && isFunction(result.then)) {
+        result.then(result => {
+          (<any>context.done)(null, result, true)
+        })
+          .catch(err => {
+            (<any>context.done)(err, null, true)
+          });
       }
     } catch (err) {
       resultCallback(err);
@@ -208,7 +241,7 @@ export class WorkerChannel implements IWorkerChannel {
   public startStream(requestId: string, msg: rpc.StartStream): void {
     // Not yet implemented
   }
-
+  
   /**
    * Message is empty by design - Will add more fields in future if needed
    */ 
@@ -304,4 +337,18 @@ export class WorkerChannel implements IWorkerChannel {
 
     return status;
   }
+
+  private runInvocationRequestBefore(context: Context, userFunction: Function): Function {
+    let wrappedFunction = userFunction;
+    for (let before of this._invocationRequestBefore) {
+      wrappedFunction = before(context, wrappedFunction);
+    }
+    return wrappedFunction;
+  }
+
+  private runInvocationRequestAfter(context: Context) {
+    for (let after of this._invocationRequestAfter) {
+      after(context);
+    }
+  }
 }

test/WorkerChannelTests.ts

@@ -6,12 +6,73 @@ import * as sinon from 'sinon';
 import { AzureFunctionsRpcMessages as rpc } from '../azure-functions-language-worker-protobuf/src/rpc';
 import 'mocha';
 import { load } from 'grpc';
+import { worker } from 'cluster';
 
 describe('WorkerChannel', () => {
   var channel: WorkerChannel;
   var stream: TestEventStream;
   var loader: sinon.SinonStubbedInstance<FunctionLoader>;
   var functions;
+  
+  const runInvokedFunction = () => {
+    const triggerDataMock: { [k: string]: rpc.ITypedData } = {
+      "Headers": {
+          json: JSON.stringify({Connection: 'Keep-Alive'})
+      },
+      "Sys": {
+          json: JSON.stringify({MethodName: 'test-js', UtcNow: '2018', RandGuid: '3212'})
+        }
+    };
+
+    const inputDataValue = {
+      name: "req",
+      data: {
+          data: "http",
+          http: 
+          {
+              body:
+              {
+                  data: "string",
+                  body: "blahh"
+              },
+              rawBody:
+              {
+                  data: "string",
+                  body: "blahh"
+              }
+          }
+      } 
+    };
+
+    const actualInvocationRequest: rpc.IInvocationRequest = <rpc.IInvocationRequest> {
+      functionId: 'id',
+      invocationId: '1',
+      inputData: [inputDataValue],
+      triggerMetadata: triggerDataMock,
+    };
+
+    stream.addTestMessage({
+      invocationRequest: actualInvocationRequest
+    });
+    
+    return [inputDataValue, actualInvocationRequest];
+  }
+  
+  const assertInvokedFunction = (inputDataValue, actualInvocationRequest) => {
+    sinon.assert.calledWithMatch(stream.written, <rpc.IStreamingMessage> {
+      invocationResponse: {
+        invocationId: '1',
+        result:  {
+          status: rpc.StatusResult.Status.Success
+        },
+        outputData: []
+      }
+    });
+
+    // triggerMedata will be augmented with inpuDataValue since "RpcHttpTriggerMetadataRemoved" capability is set to true and therefore not populated by the host.
+    expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.$request)).to.equal(JSON.stringify(inputDataValue.data));
+    expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.req)).to.equal(JSON.stringify(inputDataValue.data));
+  }
 
   beforeEach(() => {
     stream = new TestEventStream();
@@ -226,67 +287,131 @@ describe('WorkerChannel', () => {
       workerStatusResponse: {}
     });
   });
+  
+  describe('#invocationRequestBefore, #invocationRequestAfter', () => {
+    afterEach(() => {
+      channel['_invocationRequestAfter'] = [];
+      channel['_invocationRequestBefore'] = [];
+    });
 
-  it ('invokes function', () => {
-    loader.getFunc.returns((context) => context.done());
-    loader.getInfo.returns({
-      name: 'test',
-      outputBindings: {}
-    })
+    it("should apply hook before user function is executed", () => {
+      channel.registerBeforeInvocationRequest((context, userFunction) => {
+        context['magic_flag'] = 'magic value';
+        return userFunction.bind({ __wrapped: true });
+      });
+      
+      channel.registerBeforeInvocationRequest((context, userFunction) => {
+        context["secondary_flag"] = 'magic value';
+        return userFunction;
+      });
+
+      loader.getFunc.returns(function (this: any, context) {
+        expect(context['magic_flag']).to.equal('magic value');
+        expect(context['secondary_flag']).to.equal('magic value');
+        expect(this.__wrapped).to.equal(true);
+        expect(channel['_invocationRequestBefore'].length).to.equal(2);
+        expect(channel['_invocationRequestAfter'].length).to.equal(0);
+        context.done();
+      });
+      loader.getInfo.returns({
+        name: 'test',
+        outputBindings: {}
+      });
+
+      const [inputDataValue, actualInvocationRequest] = runInvokedFunction();
+      assertInvokedFunction(inputDataValue, actualInvocationRequest);
+    });
 
-    var triggerDataMock: { [k: string]: rpc.ITypedData } = {
-      "Headers": {
-          json: JSON.stringify({Connection: 'Keep-Alive'})
-      },
-      "Sys": {
-          json: JSON.stringify({MethodName: 'test-js', UtcNow: '2018', RandGuid: '3212'})
-        }
-    };
+    it('should apply hook after user function is executed (callback)', (done) => {
+      let finished = false;
+      let count = 0;
+      channel.registerAfterInvocationRequest((context) => {
+        expect(finished).to.equal(true);
+        count += 1;
+      });
 
-    var inputDataValue = {
-      name: "req",
-      data: {
-          data: "http",
-          http: 
-          {
-              body:
-              {
-                  data: "string",
-                  body: "blahh"
-              },
-              rawBody:
-              {
-                  data: "string",
-                  body: "blahh"
-              }
-          }
-      } 
-    };
+      loader.getFunc.returns(function (this: any, context) {
+        finished = true;
+        expect(channel['_invocationRequestBefore'].length).to.equal(0);
+        expect(channel['_invocationRequestAfter'].length).to.equal(1);
+        expect(count).to.equal(0);
+        context.done();
+        expect(count).to.equal(1);
+        done();
+      });
+      loader.getInfo.returns({
+        name: 'test',
+        outputBindings: {}
+      });
 
-    var actualInvocationRequest: rpc.IInvocationRequest = <rpc.IInvocationRequest> {
-      functionId: 'id',
-      invocationId: '1',
-      inputData: [inputDataValue],
-      triggerMetadata: triggerDataMock,
-    };
+      const [inputDataValue, actualInvocationRequest] = runInvokedFunction();
+      assertInvokedFunction(inputDataValue, actualInvocationRequest);
+    });
+    
+    it('should apply hook after user function resolves (promise)', (done) => {
+      let finished = false;
+      let count = 0;
+      let inputDataValue, actualInvocationRequest;
+      channel.registerAfterInvocationRequest((context) => {
+        expect(finished).to.equal(true);
+        count += 1;
+        expect(count).to.equal(1);
+        assertInvokedFunction(inputDataValue, actualInvocationRequest);
+        done();
+      });
 
-    stream.addTestMessage({
-      invocationRequest: actualInvocationRequest
+      loader.getFunc.returns(() => new Promise((resolve) => {
+        finished = true;
+        expect(channel['_invocationRequestBefore'].length).to.equal(0);
+        expect(channel['_invocationRequestAfter'].length).to.equal(1);
+        expect(count).to.equal(0);
+        resolve();
+      }));
+      loader.getInfo.returns({
+        name: 'test',
+        outputBindings: {}
+      });
+
+      [inputDataValue, actualInvocationRequest] = runInvokedFunction();
     });
+    
+    
+    it('should apply hook after user function rejects (promise)', (done) => {
+      let finished = false;
+      let count = 0;
+      channel.registerAfterInvocationRequest((context) => {
+        expect(finished).to.equal(true);
+        count += 1;
+        expect(count).to.equal(1);
+        assertInvokedFunction(inputDataValue, actualInvocationRequest);
+        done();
+      });
 
-    sinon.assert.calledWithMatch(stream.written, <rpc.IStreamingMessage> {
-      invocationResponse: {
-        invocationId: '1',
-        result:  {
-          status: rpc.StatusResult.Status.Success
-        },
-        outputData: []
-      }
+      loader.getFunc.returns((context) => new Promise((_, reject) => {
+        finished = true;
+        expect(channel['_invocationRequestBefore'].length).to.equal(0);
+        expect(channel['_invocationRequestAfter'].length).to.equal(1);
+        expect(count).to.equal(0);
+        reject();
+      }));
+      loader.getInfo.returns({
+        name: 'test',
+        outputBindings: {}
+      });
+
+      const [inputDataValue, actualInvocationRequest] = runInvokedFunction();
     });
+  });
 
-    // triggerMedata will be augmented with inpuDataValue since "RpcHttpTriggerMetadataRemoved" capability is set to true and therefore not populated by the host.
-    expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.$request)).to.equal(JSON.stringify(inputDataValue.data));
-    expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.req)).to.equal(JSON.stringify(inputDataValue.data));
+  it ('invokes function', () => {
+    loader.getFunc.returns((context) => context.done());
+    loader.getInfo.returns({
+      name: 'test',
+      outputBindings: {}
+    })
+
+    const [inputDataValue, actualInvocationRequest] = runInvokedFunction();
+    assertInvokedFunction(inputDataValue, actualInvocationRequest);
   });
 
   it ('throws for malformed messages', () => {