Add pre and post invocation hooks

.eslintrc.json

@@ -43,6 +43,7 @@
         "@typescript-eslint/restrict-template-expressions": "off",
         "@typescript-eslint/unbound-method": "off",
         "no-empty": "off",
+        "prefer-const": ["error", { "destructuring": "all" }],
         "prefer-rest-params": "off",
         "prefer-spread": "off"
     },

src/Context.ts

@@ -22,16 +22,12 @@ import {
 import { FunctionInfo } from './FunctionInfo';
 import { Request } from './http/Request';
 import { Response } from './http/Response';
+import EventEmitter = require('events');
 import LogLevel = rpc.RpcLog.Level;
-import LogCategory = rpc.RpcLog.RpcLogCategory;
 
-export function CreateContextAndInputs(
-    info: FunctionInfo,
-    request: rpc.IInvocationRequest,
-    logCallback: LogCallback,
-    callback: ResultCallback
-) {
-    const context = new InvocationContext(info, request, logCallback, callback);
+export function CreateContextAndInputs(info: FunctionInfo, request: rpc.IInvocationRequest, logCallback: LogCallback) {
+    const doneEmitter = new EventEmitter();
+    const context = new InvocationContext(info, request, logCallback, doneEmitter);
 
     const bindings: ContextBindings = {};
     const inputs: any[] = [];
@@ -76,6 +72,7 @@ export function CreateContextAndInputs(
     return {
         context: <Context>context,
         inputs: inputs,
+        doneEmitter,
     };
 }
 
@@ -95,7 +92,7 @@ class InvocationContext implements Context {
         info: FunctionInfo,
         request: rpc.IInvocationRequest,
         logCallback: LogCallback,
-        callback: ResultCallback
+        doneEmitter: EventEmitter
     ) {
         this.invocationId = <string>request.invocationId;
         this.traceContext = fromRpcTraceContext(request.traceContext);
@@ -107,89 +104,32 @@ class InvocationContext implements Context {
         };
         this.executionContext = executionContext;
         this.bindings = {};
-        let _done = false;
-        let _promise = false;
 
         // Log message that is tied to function invocation
-        this.log = Object.assign(
-            (...args: any[]) => logWithAsyncCheck(_done, logCallback, LogLevel.Information, executionContext, ...args),
-            {
-                error: (...args: any[]) =>
-                    logWithAsyncCheck(_done, logCallback, LogLevel.Error, executionContext, ...args),
-                warn: (...args: any[]) =>
-                    logWithAsyncCheck(_done, logCallback, LogLevel.Warning, executionContext, ...args),
-                info: (...args: any[]) =>
-                    logWithAsyncCheck(_done, logCallback, LogLevel.Information, executionContext, ...args),
-                verbose: (...args: any[]) =>
-                    logWithAsyncCheck(_done, logCallback, LogLevel.Trace, executionContext, ...args),
-            }
-        );
+        this.log = Object.assign((...args: any[]) => logCallback(LogLevel.Information, ...args), {
+            error: (...args: any[]) => logCallback(LogLevel.Error, ...args),
+            warn: (...args: any[]) => logCallback(LogLevel.Warning, ...args),
+            info: (...args: any[]) => logCallback(LogLevel.Information, ...args),
+            verbose: (...args: any[]) => logCallback(LogLevel.Trace, ...args),
+        });
 
         this.bindingData = getNormalizedBindingData(request);
         this.bindingDefinitions = getBindingDefinitions(info);
 
-        // isPromise is a hidden parameter that we set to true in the event of a returned promise
-        this.done = (err?: any, result?: any, isPromise?: boolean) => {
-            _promise = isPromise === true;
-            if (_done) {
-                if (_promise) {
-                    logCallback(
-                        LogLevel.Error,
-                        LogCategory.User,
-                        "Error: Choose either to return a promise or call 'done'.  Do not use both in your script."
-                    );
-                } else {
-                    logCallback(
-                        LogLevel.Error,
-                        LogCategory.User,
-                        "Error: 'done' has already been called. Please check your script for extraneous calls to 'done'."
-                    );
-                }
-                return;
-            }
-            _done = true;
-
-            // Allow HTTP response from context.res if HTTP response is not defined from the context.bindings object
-            if (info.httpOutputName && this.res && this.bindings[info.httpOutputName] === undefined) {
-                this.bindings[info.httpOutputName] = this.res;
-            }
-
-            callback(err, {
-                return: result,
-                bindings: this.bindings,
-            });
+        this.done = (err?: unknown, result?: any) => {
+            doneEmitter.emit('done', err, result);
         };
     }
 }
 
-// Emit warning if trying to log after function execution is done.
-function logWithAsyncCheck(
-    done: boolean,
-    log: LogCallback,
-    level: LogLevel,
-    executionContext: ExecutionContext,
-    ...args: any[]
-) {
-    if (done) {
-        let badAsyncMsg =
-            "Warning: Unexpected call to 'log' on the context object after function execution has completed. Please check for asynchronous calls that are not awaited or calls to 'done' made before function execution completes. ";
-        badAsyncMsg += `Function name: ${executionContext.functionName}. Invocation Id: ${executionContext.invocationId}. `;
-        badAsyncMsg += `Learn more: https://go.microsoft.com/fwlink/?linkid=2097909 `;
-        log(LogLevel.Warning, LogCategory.System, badAsyncMsg);
-    }
-    return log(level, LogCategory.User, ...args);
-}
-
 export interface InvocationResult {
     return: any;
     bindings: ContextBindings;
 }
 
-export type DoneCallback = (err?: Error | string, result?: any) => void;
-
-export type LogCallback = (level: LogLevel, category: rpc.RpcLog.RpcLogCategory, ...args: any[]) => void;
+export type DoneCallback = (err?: unknown, result?: any) => void;
 
-export type ResultCallback = (err?: any, result?: InvocationResult) => void;
+export type LogCallback = (level: LogLevel, ...args: any[]) => void;
 
 export interface Dict<T> {
     [key: string]: T;

src/Disposable.ts

@@ -0,0 +1,35 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License.
+
+/**
+ * Based off of VS Code
+ * https://github.com/microsoft/vscode/blob/a64e8e5673a44e5b9c2d493666bde684bd5a135c/src/vs/workbench/api/common/extHostTypes.ts#L32
+ */
+export class Disposable {
+    static from(...inDisposables: { dispose(): any }[]): Disposable {
+        let disposables: ReadonlyArray<{ dispose(): any }> | undefined = inDisposables;
+        return new Disposable(function () {
+            if (disposables) {
+                for (const disposable of disposables) {
+                    if (disposable && typeof disposable.dispose === 'function') {
+                        disposable.dispose();
+                    }
+                }
+                disposables = undefined;
+            }
+        });
+    }
+
+    #callOnDispose?: () => any;
+
+    constructor(callOnDispose: () => any) {
+        this.#callOnDispose = callOnDispose;
+    }
+
+    dispose(): any {
+        if (this.#callOnDispose instanceof Function) {
+            this.#callOnDispose();
+            this.#callOnDispose = undefined;
+        }
+    }
+}

src/Worker.ts

@@ -5,6 +5,7 @@ import * as parseArgs from 'minimist';
 import { FunctionLoader } from './FunctionLoader';
 import { CreateGrpcEventStream } from './GrpcClient';
 import { setupEventStream } from './setupEventStream';
+import { setupWorkerModule } from './setupWorkerModule';
 import { ensureErrorType } from './utils/ensureErrorType';
 import { InternalException } from './utils/InternalException';
 import { systemError, systemLog } from './utils/Logger';
@@ -42,6 +43,7 @@ export function startNodeWorker(args) {
 
     const channel = new WorkerChannel(eventStream, new FunctionLoader());
     setupEventStream(workerId, channel);
+    setupWorkerModule(channel);
 
     eventStream.write({
         requestId: requestId,

src/WorkerChannel.ts

@@ -1,25 +1,21 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License.
 
-import { Context } from '@azure/functions';
+import { HookCallback, HookContext } from '@azure/functions-worker';
 import { AzureFunctionsRpcMessages as rpc } from '../azure-functions-language-worker-protobuf/src/rpc';
+import { Disposable } from './Disposable';
 import { IFunctionLoader } from './FunctionLoader';
 import { IEventStream } from './GrpcClient';
 
-type InvocationRequestBefore = (context: Context, userFn: Function) => Function;
-type InvocationRequestAfter = (context: Context) => void;
-
 export class WorkerChannel {
     public eventStream: IEventStream;
     public functionLoader: IFunctionLoader;
-    private _invocationRequestBefore: InvocationRequestBefore[];
-    private _invocationRequestAfter: InvocationRequestAfter[];
+    private _preInvocationHooks: HookCallback[] = [];
+    private _postInvocationHooks: HookCallback[] = [];
 
     constructor(eventStream: IEventStream, functionLoader: IFunctionLoader) {
         this.eventStream = eventStream;
         this.functionLoader = functionLoader;
-        this._invocationRequestBefore = [];
-        this._invocationRequestAfter = [];
     }
 
     /**
@@ -33,32 +29,32 @@ export class WorkerChannel {
         });
     }
 
-    /**
-     * 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);
+    public registerHook(hookName: string, callback: HookCallback): Disposable {
+        const hooks = this.getHooks(hookName);
+        hooks.push(callback);
+        return new Disposable(() => {
+            const index = hooks.indexOf(callback);
+            if (index > -1) {
+                hooks.splice(index, 1);
+            }
+        });
     }
 
-    public runInvocationRequestBefore(context: Context, userFunction: Function): Function {
-        let wrappedFunction = userFunction;
-        for (const before of this._invocationRequestBefore) {
-            wrappedFunction = before(context, wrappedFunction);
+    public async executeHooks(hookName: string, context: HookContext): Promise<void> {
+        const callbacks = this.getHooks(hookName);
+        for (const callback of callbacks) {
+            await callback(context);
         }
-        return wrappedFunction;
     }
 
-    public runInvocationRequestAfter(context: Context) {
-        for (const after of this._invocationRequestAfter) {
-            after(context);
+    private getHooks(hookName: string): HookCallback[] {
+        switch (hookName) {
+            case 'preInvocation':
+                return this._preInvocationHooks;
+            case 'postInvocation':
+                return this._postInvocationHooks;
+            default:
+                throw new RangeError(`Unrecognized hook "${hookName}"`);
         }
     }
 }