Add app-level startup hooks (#577)

* add startup and teardown hook support

* add initial types

* move hook data handling to wokrer channel

* add app startup hooks to worker init handler

* run new script file on functionEnvironmentReloadRequest

* add logger to base hook

* refactor

* update comment

* create copies of hook data for invocation contexts

* don't copy for app startup hook data + add function app directory

* function environment reload logic

* implement logger interface in hook context

* add host version to startup hooks

* fix tests

* rename app startup function

* remove teardown hooks

* proper message category and functionInvocationId in hook logs

* make sure it compiles

* remove logger

* refactor app startup code

* don't reset the hook data

* simplify copying hook data

* add specialization test for environment reload handler

* add test for loading entry point in specialization scenario

* add test for running app startup hook in non-specialization scenario

* assert called with right context

* add test for specialization scenario

* add test for persisting hook data from startup hooks

* add comment

* allow functionAppDirectory to be undefined in workerInitRequest

* rename to appStart

* one more file rename

* add appHookData and hookData distinction + tests

* fix tests

* remove outdated comment

* non null hostVersion and remove log warning

* fix tests

* old hookData behavior

* address PR comments

Author: hossam-nasr

src/WorkerChannel.ts

@@ -1,7 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License.
 
-import { HookCallback, HookContext } from '@azure/functions-core';
+import { HookCallback, HookContext, HookData } from '@azure/functions-core';
 import { AzureFunctionsRpcMessages as rpc } from '../azure-functions-language-worker-protobuf/src/rpc';
 import { Disposable } from './Disposable';
 import { IFunctionLoader } from './FunctionLoader';
@@ -15,8 +15,21 @@ export class WorkerChannel {
     eventStream: IEventStream;
     functionLoader: IFunctionLoader;
     packageJson: PackageJson;
+    /**
+     * This will only be set after worker init request is received
+     */
+    hostVersion?: string;
+    /**
+     * this hook data will be passed to (and set by) all hooks in all scopes
+     */
+    appHookData: HookData = {};
+    /**
+     * this hook data is limited to the app-level scope and persisted only for app-level hooks
+     */
+    appLevelOnlyHookData: HookData = {};
     #preInvocationHooks: HookCallback[] = [];
     #postInvocationHooks: HookCallback[] = [];
+    #appStartHooks: HookCallback[] = [];
 
     constructor(eventStream: IEventStream, functionLoader: IFunctionLoader) {
         this.eventStream = eventStream;
@@ -80,6 +93,8 @@ export class WorkerChannel {
                 return this.#preInvocationHooks;
             case 'postInvocation':
                 return this.#postInvocationHooks;
+            case 'appStart':
+                return this.#appStartHooks;
             default:
                 throw new RangeError(`Unrecognized hook "${hookName}"`);
         }

src/eventHandlers/FunctionEnvironmentReloadHandler.ts

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 import { AzureFunctionsRpcMessages as rpc } from '../../azure-functions-language-worker-protobuf/src/rpc';
+import { startApp } from '../startApp';
 import { WorkerChannel } from '../WorkerChannel';
 import { EventHandler } from './EventHandler';
 import LogCategory = rpc.RpcLog.RpcLogCategory;
@@ -35,6 +36,7 @@ export class FunctionEnvironmentReloadHandler extends EventHandler<
         });
 
         process.env = Object.assign({}, msg.environmentVariables);
+
         // Change current working directory
         if (msg.functionAppDirectory) {
             channel.log({
@@ -43,7 +45,7 @@ export class FunctionEnvironmentReloadHandler extends EventHandler<
                 logCategory: LogCategory.System,
             });
             process.chdir(msg.functionAppDirectory);
-            await channel.updatePackageJson(msg.functionAppDirectory);
+            await startApp(msg.functionAppDirectory, channel);
         }
 
         return response;

src/eventHandlers/InvocationHandler.ts

@@ -90,16 +90,18 @@ export class InvocationHandler extends EventHandler<'invocationRequest', 'invoca
                 });
             });
 
-            const hookData: HookData = {};
             let userFunction = channel.functionLoader.getFunc(functionId);
+
+            const invocationHookData: HookData = {};
+
             const preInvocContext: PreInvocationContext = {
-                hookData,
+                hookData: invocationHookData,
+                appHookData: channel.appHookData,
                 invocationContext: context,
                 functionCallback: <AzureFunction>userFunction,
                 inputs,
             };
-
-            await channel.executeHooks('preInvocation', preInvocContext, msg.invocationId, msgCategory);
+            await channel.executeHooks('preInvocation', preInvocContext, invocationId, msgCategory);
             inputs = preInvocContext.inputs;
             userFunction = preInvocContext.functionCallback;
 
@@ -117,7 +119,8 @@ export class InvocationHandler extends EventHandler<'invocationRequest', 'invoca
             }
 
             const postInvocContext: PostInvocationContext = {
-                hookData,
+                hookData: invocationHookData,
+                appHookData: channel.appHookData,
                 invocationContext: context,
                 inputs,
                 result: null,

src/eventHandlers/WorkerInitHandler.ts

@@ -4,7 +4,9 @@
 import { access, constants } from 'fs';
 import * as path from 'path';
 import { AzureFunctionsRpcMessages as rpc } from '../../azure-functions-language-worker-protobuf/src/rpc';
+import { startApp } from '../startApp';
 import { isError } from '../utils/ensureErrorType';
+import { nonNullProp } from '../utils/nonNull';
 import { WorkerChannel } from '../WorkerChannel';
 import { EventHandler } from './EventHandler';
 import LogCategory = rpc.RpcLog.RpcLogCategory;
@@ -30,9 +32,11 @@ export class WorkerInitHandler extends EventHandler<'workerInitRequest', 'worker
         });
 
         logColdStartWarning(channel);
-        const functionAppDirectory = msg.functionAppDirectory;
-        if (functionAppDirectory) {
-            await channel.updatePackageJson(functionAppDirectory);
+
+        channel.hostVersion = nonNullProp(msg, 'hostVersion');
+
+        if (msg.functionAppDirectory) {
+            await startApp(msg.functionAppDirectory, channel);
         }
 
         response.capabilities = {

src/startApp.ts

@@ -0,0 +1,62 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License.
+
+import { AppStartContext } from '@azure/functions-core';
+import { pathExists } from 'fs-extra';
+import { AzureFunctionsRpcMessages as rpc } from '../azure-functions-language-worker-protobuf/src/rpc';
+import { loadScriptFile } from './loadScriptFile';
+import { ensureErrorType } from './utils/ensureErrorType';
+import { nonNullProp } from './utils/nonNull';
+import { WorkerChannel } from './WorkerChannel';
+import path = require('path');
+import LogLevel = rpc.RpcLog.Level;
+import LogCategory = rpc.RpcLog.RpcLogCategory;
+
+/**
+ * Starting an app can happen in two places, depending on if the worker was specialized or not
+ * 1. The worker can start in "normal" mode, meaning `workerInitRequest` will reference the user's app
+ * 2. The worker can start in "placeholder" mode, meaning `workerInitRequest` will reference a dummy app to "warm up" the worker and `functionEnvironmentReloadRequest` will be sent with the user's actual app.
+ *    This process is called worker specialization and it helps with cold start times.
+ *    The dummy app should never have actual startup code, so it should be safe to call `startApp` twice in this case
+ *    Worker specialization happens only once, so we don't need to worry about cleaning up resources from previous `functionEnvironmentReloadRequest`s.
+ */
+export async function startApp(functionAppDirectory: string, channel: WorkerChannel): Promise<void> {
+    await channel.updatePackageJson(functionAppDirectory);
+    await loadEntryPointFile(functionAppDirectory, channel);
+    const appStartContext: AppStartContext = {
+        hookData: channel.appLevelOnlyHookData,
+        appHookData: channel.appHookData,
+        functionAppDirectory,
+        hostVersion: nonNullProp(channel, 'hostVersion'),
+    };
+    await channel.executeHooks('appStart', appStartContext);
+}
+
+async function loadEntryPointFile(functionAppDirectory: string, channel: WorkerChannel): Promise<void> {
+    const entryPointFile = channel.packageJson.main;
+    if (entryPointFile) {
+        channel.log({
+            message: `Loading entry point "${entryPointFile}"`,
+            level: LogLevel.Debug,
+            logCategory: LogCategory.System,
+        });
+        try {
+            const entryPointFullPath = path.join(functionAppDirectory, entryPointFile);
+            if (!(await pathExists(entryPointFullPath))) {
+                throw new Error(`file does not exist`);
+            }
+
+            await loadScriptFile(entryPointFullPath, channel.packageJson);
+            channel.log({
+                message: `Loaded entry point "${entryPointFile}"`,
+                level: LogLevel.Debug,
+                logCategory: LogCategory.System,
+            });
+        } catch (err) {
+            const error = ensureErrorType(err);
+            error.isAzureFunctionsInternalException = true;
+            error.message = `Worker was unable to load entry point "${entryPointFile}": ${error.message}`;
+            throw error;
+        }
+    }
+}

test/eventHandlers/FunctionEnvironmentReloadHandler.test.ts

@@ -8,11 +8,12 @@ import { AzureFunctionsRpcMessages as rpc } from '../../azure-functions-language
 import { WorkerChannel } from '../../src/WorkerChannel';
 import { beforeEventHandlerSuite } from './beforeEventHandlerSuite';
 import { TestEventStream } from './TestEventStream';
+import { Msg as WorkerInitMsg } from './WorkerInitHandler.test';
 import path = require('path');
 import LogCategory = rpc.RpcLog.RpcLogCategory;
 import LogLevel = rpc.RpcLog.Level;
 
-namespace Msg {
+export namespace Msg {
     export function reloadEnvVarsLog(numVars: number): rpc.IStreamingMessage {
         return {
             rpcLog: {
@@ -69,11 +70,14 @@ describe('FunctionEnvironmentReloadHandler', () => {
     let stream: TestEventStream;
     let channel: WorkerChannel;
 
-    // Reset `process.env` after this test suite so it doesn't affect other tests
+    // Reset `process.env` and process.cwd() after this test suite so it doesn't affect other tests
     let originalEnv: NodeJS.ProcessEnv;
+    let originalCwd: string;
     before(() => {
+        originalCwd = process.cwd();
         originalEnv = process.env;
         ({ stream, channel } = beforeEventHandlerSuite());
+        channel.hostVersion = '2.7.0';
     });
 
     after(() => {
@@ -82,6 +86,7 @@ describe('FunctionEnvironmentReloadHandler', () => {
 
     afterEach(async () => {
         mock.restore();
+        process.chdir(originalCwd);
         await stream.afterEachEventHandlerTest();
     });
 
@@ -229,6 +234,83 @@ describe('FunctionEnvironmentReloadHandler', () => {
         });
         await stream.assertCalledWith(Msg.reloadEnvVarsLog(0), Msg.changingCwdLog(newDirAbsolute), Msg.reloadSuccess);
         expect(channel.packageJson).to.deep.equal(newPackageJson);
-        process.chdir(cwd);
     });
+
+    it('correctly loads package.json in specialization scenario', async () => {
+        const cwd = process.cwd();
+        const tempDir = 'temp';
+        const appDir = 'app';
+        const packageJson = {
+            type: 'module',
+            hello: 'world',
+        };
+
+        mock({
+            [tempDir]: {},
+            [appDir]: {
+                'package.json': JSON.stringify(packageJson),
+            },
+        });
+
+        stream.addTestMessage(WorkerInitMsg.init(path.join(cwd, tempDir)));
+        await stream.assertCalledWith(
+            WorkerInitMsg.receivedInitLog,
+            WorkerInitMsg.warning(`Worker failed to load package.json: file does not exist`),
+            WorkerInitMsg.response
+        );
+        expect(channel.packageJson).to.be.empty;
+
+        stream.addTestMessage({
+            requestId: 'id',
+            functionEnvironmentReloadRequest: {
+                functionAppDirectory: path.join(cwd, appDir),
+            },
+        });
+        await stream.assertCalledWith(
+            Msg.reloadEnvVarsLog(0),
+            Msg.changingCwdLog(path.join(cwd, appDir)),
+            Msg.reloadSuccess
+        );
+        expect(channel.packageJson).to.deep.equal(packageJson);
+    });
+
+    for (const extension of ['.js', '.mjs', '.cjs']) {
+        it(`Loads entry point (${extension}) in specialization scenario`, async () => {
+            const cwd = process.cwd();
+            const tempDir = 'temp';
+            const fileName = `entryPointFiles/doNothing${extension}`;
+            const expectedPackageJson = {
+                main: fileName,
+            };
+            mock({
+                [tempDir]: {},
+                [__dirname]: {
+                    'package.json': JSON.stringify(expectedPackageJson),
+                    // 'require' and 'mockFs' don't play well together so we need these files in both the mock and real file systems
+                    entryPointFiles: mock.load(path.join(__dirname, 'entryPointFiles')),
+                },
+            });
+
+            stream.addTestMessage(WorkerInitMsg.init(path.join(cwd, tempDir)));
+            await stream.assertCalledWith(
+                WorkerInitMsg.receivedInitLog,
+                WorkerInitMsg.warning('Worker failed to load package.json: file does not exist'),
+                WorkerInitMsg.response
+            );
+
+            stream.addTestMessage({
+                requestId: 'id',
+                functionEnvironmentReloadRequest: {
+                    functionAppDirectory: __dirname,
+                },
+            });
+            await stream.assertCalledWith(
+                Msg.reloadEnvVarsLog(0),
+                Msg.changingCwdLog(__dirname),
+                WorkerInitMsg.loadingEntryPoint(fileName),
+                WorkerInitMsg.loadedEntryPoint(fileName),
+                Msg.reloadSuccess
+            );
+        });
+    }
 });