From ee335ba00f3a44e0cc49da774c28c82125f420ae Mon Sep 17 00:00:00 2001
From: Haili Zhang <haili.zhang@outlook.com>
Date: Sat, 30 Jul 2022 16:33:17 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20enable=20graceful=20shutdow?=
 =?UTF-8?q?n?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Haili Zhang <haili.zhang@outlook.com>
---
 package-lock.json             |  1 +
 package.json                  |  1 +
 src/main.ts                   | 42 ++++++++++++++++++++++++++++++-----
 test/conformance/package.json |  5 +++--
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 12f97f9b..f3d638ce 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -15,6 +15,7 @@
         "debug": "^4.3.4",
         "express": "^4.18.1",
         "express-interceptor": "^1.2.0",
+        "http-terminator": "^3.2.0",
         "lodash": "^4.17.21",
         "minimist": "^1.2.6",
         "on-finished": "^2.4.1",
diff --git a/package.json b/package.json
index e03bf3e1..fb193a45 100644
--- a/package.json
+++ b/package.json
@@ -19,6 +19,7 @@
     "debug": "^4.3.4",
     "express": "^4.18.1",
     "express-interceptor": "^1.2.0",
+    "http-terminator": "^3.2.0",
     "lodash": "^4.17.21",
     "minimist": "^1.2.6",
     "on-finished": "^2.4.1",
diff --git a/src/main.ts b/src/main.ts
index 1622a9c5..bee7a1c1 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -16,17 +16,22 @@
 
 // Functions framework entry point that configures and starts Node.js server
 // that runs user's code on HTTP request.
-import {getUserFunction} from './loader';
-import {ErrorHandler} from './invoker';
-import {getServer} from './server';
-import {parseOptions, helpText, OptionsError} from './options';
-import {OpenFunction} from './functions';
+import * as process from 'process';
+
+import {createHttpTerminator} from 'http-terminator';
+
 import getAysncServer from './openfunction/async_server';
 import {
   OpenFunctionContext,
   ContextUtils,
 } from './openfunction/function_context';
 
+import {getUserFunction} from './loader';
+import {ErrorHandler} from './invoker';
+import {getServer} from './server';
+import {parseOptions, helpText, OptionsError} from './options';
+import {OpenFunction} from './functions';
+
 /**
  * Main entrypoint for the functions framework that loads the user's function
  * and starts the HTTP server.
@@ -51,6 +56,8 @@ export const main = async () => {
     }
     const {userFunction, signatureType} = loadedFunction;
 
+    // Try to determine the server runtime
+    // Considering the async runtime in the first place
     if (ContextUtils.IsAsyncRuntime(options.context as OpenFunctionContext)) {
       options.context!.port = options.port;
 
@@ -59,7 +66,12 @@ export const main = async () => {
         options.context!
       );
       await server.start();
-    } else {
+
+      // DaprServer uses httpTerminator in server.stop()
+      handleShutdown(async () => await server.stop());
+    }
+    // Then taking sync runtime as the fallback
+    else {
       const server = getServer(userFunction!, signatureType, options.context);
       const errorHandler = new ErrorHandler(server);
       server
@@ -73,6 +85,12 @@ export const main = async () => {
           }
         })
         .setTimeout(0); // Disable automatic timeout on incoming connections.
+
+      // Create and use httpTerminator for Express
+      const terminator = createHttpTerminator({
+        server,
+      });
+      handleShutdown(async () => await terminator.terminate());
     }
   } catch (e) {
     if (e instanceof OptionsError) {
@@ -86,3 +104,15 @@ export const main = async () => {
 
 // Call the main method to load the user code and start the http server.
 main();
+
+function handleShutdown(handler: () => Promise<void>): void {
+  if (!handler) return;
+
+  const shutdown = async (code: string) => {
+    console.log(`🛑 Terminating OpenFunction server on code ${code}...`);
+    await handler();
+  };
+
+  process.on('SIGTERM', shutdown);
+  process.on('SIGINT', shutdown);
+}
diff --git a/test/conformance/package.json b/test/conformance/package.json
index f4eb9778..3145dd38 100644
--- a/test/conformance/package.json
+++ b/test/conformance/package.json
@@ -8,11 +8,12 @@
   },
   "scripts": {
     "start": "functions-framework --target=writeHttp",
-    "knative:async": "concurrently npm:knative:async:run:* npm:knative:async:test",
+    "compile": "cd .. && npm run compile",
+    "knative:async": "npm run compile && concurrently npm:knative:async:run:* npm:knative:async:test",
     "knative:async:run:func": "cross-env DEBUG=test:*,common:*,ofn:* env-cmd -e knative functions-framework --signature-type=openfunction --target=tryKnativeAsync",
     "knative:async:run:dapr": "dapr run -H 3500 -d ../data/components/http --log-level warn",
     "knative:async:test": "wait-on tcp:3500 tcp:8080 && curl -s -d '{\"data\": \"hello\"}' -H 'Content-Type: application/json' localhost:8080",
-    "async": "concurrently npm:async:run:*",
+    "async": "npm run compile && concurrently npm:async:run:*",
     "async:run:func": "cross-env DEBUG=test:*,common:*,ofn:* env-cmd -e async functions-framework --target=tryAsync",
     "async:run:dapr": "dapr run -H 3500 -p 8080 -d ../data/components/cron --log-level info"
   }