[worker v2] feature: add before/after InvocationRequest hooks (#342)

* 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

* revert changes

* Re-added before / after invocation tests refactored for this branch

Co-authored-by: Mark Wolff <[email protected]>

Author: mhoeger

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);
+    }
+  }
 }