feature: add before/after InvocationRequest hooks (Azure#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, systemWarn } 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;
 }
 
 /**
@@ -35,13 +41,17 @@ export class WorkerChannel implements IWorkerChannel {
   private _functionLoader: IFunctionLoader;
   private _workerId: string;
   private _v1WorkerBehavior: boolean;
+  private _invocationRequestBefore: InvocationRequestBefore[];
+  private _invocationRequestAfter: InvocationRequestAfter[];
 
   constructor(workerId: string, eventStream: IEventStream, functionLoader: IFunctionLoader) {
     this._workerId = workerId;
     this._eventStream = eventStream;
     this._functionLoader = functionLoader;
     // default value
     this._v1WorkerBehavior = false;
+    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) => {
@@ -85,6 +95,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
@@ -252,19 +277,27 @@ export class WorkerChannel implements IWorkerChannel {
         requestId: requestId,
         invocationResponse: response
       });
+      
+      this.runInvocationRequestAfter(context);
     }
 
     let { context, inputs } = CreateContextAndInputs(info, msg, logCallback, resultCallback, this._v1WorkerBehavior);
     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);
@@ -277,7 +310,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
    */ 
@@ -369,4 +402,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,6 +6,7 @@ import * as sinon from 'sinon';
 import { AzureFunctionsRpcMessages as rpc } from '../azure-functions-language-worker-protobuf/src/rpc';
 import 'mocha';
 import { load } from 'grpc';
+<<<<<<< HEAD
 import { FunctionInfo } from '../src/FunctionInfo';
 
 describe('WorkerChannel', () => {
@@ -85,6 +86,75 @@ describe('WorkerChannel', () => {
       test: orchestrationTriggerBinding
     }
   };
+=======
+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));
+  }
+>>>>>>> 0ff6fb4... feature: add before/after InvocationRequest hooks (#339)
 
   beforeEach(() => {
     stream = new TestEventStream();
@@ -378,7 +448,14 @@ describe('WorkerChannel', () => {
     expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.$request)).to.equal(JSON.stringify(httpInputData.data));
     expect(JSON.stringify(actualInvocationRequest.triggerMetadata!.req)).to.equal(JSON.stringify(httpInputData.data));
   });
+  
+  describe('#invocationRequestBefore, #invocationRequestAfter', () => {
+    afterEach(() => {
+      channel['_invocationRequestAfter'] = [];
+      channel['_invocationRequestBefore'] = [];
+    });
 
+<<<<<<< HEAD
   it ('invokes function', () => {
     loader.getFunc.returns((context) => context.done());
     loader.getInfo.returns(new FunctionInfo(orchestratorBinding));
@@ -529,11 +606,102 @@ describe('WorkerChannel', () => {
       inputData: [httpInputData],
       triggerMetadata: getTriggerDataMock(),
     };
+=======
+    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;
+      });
 
-    stream.addTestMessage({
-      invocationRequest: actualInvocationRequest
+      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);
     });
+    
+    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;
+      });
+
+      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: {}
+      });
 
+      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();
+      });
+>>>>>>> 0ff6fb4... feature: add before/after InvocationRequest hooks (#339)
+
+      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();
+      });
+
+<<<<<<< HEAD
     sinon.assert.calledWithMatch(stream.written, <rpc.IStreamingMessage> {
       invocationResponse: {
         invocationId: '1',
@@ -682,6 +850,33 @@ describe('WorkerChannel', () => {
         }
       }
     });
+=======
+      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();
+    });
+  });
+
+  it ('invokes function', () => {
+    loader.getFunc.returns((context) => context.done());
+    loader.getInfo.returns({
+      name: 'test',
+      outputBindings: {}
+    })
+
+    const [inputDataValue, actualInvocationRequest] = runInvokedFunction();
+    assertInvokedFunction(inputDataValue, actualInvocationRequest);
+>>>>>>> 0ff6fb4... feature: add before/after InvocationRequest hooks (#339)
   });
 
   it ('throws for malformed messages', () => {