From 1afa4bb8847fb4aa1dac49a826e7ca14936375c8 Mon Sep 17 00:00:00 2001
From: jonirap <joni.rapoport@gmail.com>
Date: Mon, 17 Jul 2023 12:21:31 -0600
Subject: [PATCH] Allow for async init and shutdown functions

---
 index.js       | 60 ++++++++++++++++++++++++++++++++------------------
 lib/types.d.ts |  4 ++--
 lib/utils.js   |  9 ++++++++
 package.json   |  1 +
 tsconfig.json  | 11 +++------
 5 files changed, 53 insertions(+), 32 deletions(-)
 create mode 100644 lib/utils.js

diff --git a/index.js b/index.js
index 1905c71..36f32ba 100644
--- a/index.js
+++ b/index.js
@@ -7,6 +7,7 @@ const eventHandler = require('./lib/event-handler');
 const Context = require('./lib/context');
 const shutdown = require('death')({ uncaughtException: true });
 const fastifyRawBody = require('fastify-raw-body');
+const { isPromise } = require('./lib/utils');
 
 // HTTP framework
 const fastify = require('fastify');
@@ -41,7 +42,10 @@ async function start(func, options) {
     throw new TypeError('Function must export a handle function');
   }
   if (typeof func.init === 'function') {
-    func.init();
+    const initRet = func.init();
+    if (isPromise(initRet)) {
+      await initRet;
+    }
   }
   if (typeof func.shutdown === 'function') {
     options.shutdown = func.shutdown;
@@ -79,10 +83,10 @@ async function __start(func, options) {
   try {
     await server.listen({
       port: config.port,
-      host: '::'
+      host: '::',
     });
     return server.server;
-  } catch(err) {
+  } catch (err) {
     console.error('Error starting server', err);
     process.exit(1);
   }
@@ -100,12 +104,12 @@ function initializeServer(config) {
       level: config.logLevel,
       formatters: {
         bindings: bindings => ({
-            pid: bindings.pid,
-            hostname: bindings.hostname,
-            node_version: process.version
-        })
-      }
-    }
+          pid: bindings.pid,
+          hostname: bindings.hostname,
+          node_version: process.version,
+        }),
+      },
+    },
   });
 
   if (config.includeRaw) {
@@ -118,16 +122,20 @@ function initializeServer(config) {
   }
 
   // Give the Function an opportunity to clean up before the process exits
-  shutdown(_ => {
+  shutdown(async _ => {
     if (typeof config.shutdown === 'function') {
-      config.shutdown();
+      const shutdownRet = config.shutdown();
+      if (isPromise(shutdownRet)) {
+        await shutdownRet;
+      }
     }
     server.close();
     process.exit(0);
   });
 
   // Add a parser for application/x-www-form-urlencoded
-  server.addContentTypeParser('application/x-www-form-urlencoded',
+  server.addContentTypeParser(
+    'application/x-www-form-urlencoded',
     function(_, payload, done) {
       var body = '';
       payload.on('data', data => (body += data));
@@ -140,18 +148,23 @@ function initializeServer(config) {
         }
       });
       payload.on('error', done);
-    });
+    }
+  );
 
   // Add a parser for everything else - parse it as a buffer and
   // let this framework's router handle it
-  server.addContentTypeParser('*', { parseAs: 'buffer' }, function(req, body, done) {
-    try {
-      done(null, body);
-    } catch (err) {
-      err.statusCode = 500;
-      done(err, undefined);
+  server.addContentTypeParser(
+    '*',
+    { parseAs: 'buffer' },
+    function(req, body, done) {
+      try {
+        done(null, body);
+      } catch (err) {
+        err.statusCode = 500;
+        done(err, undefined);
+      }
     }
-  });
+  );
 
   // Initialize the invocation context
   // This is passed as a parameter to the function when it's invoked
@@ -210,11 +223,14 @@ function readFuncYaml(fileOrDirPath) {
     if (!!maybeYaml && maybeYaml.isFile()) {
       try {
         return yaml.load(fs.readFileSync(yamlFile, 'utf8'));
-      } catch(err) {
+      } catch (err) {
         console.warn(err);
       }
     }
   }
 }
 
-module.exports = exports = { start, defaults: { LOG_LEVEL, PORT, INCLUDE_RAW } };
+module.exports = exports = {
+  start,
+  defaults: { LOG_LEVEL, PORT, INCLUDE_RAW },
+};
diff --git a/lib/types.d.ts b/lib/types.d.ts
index ec9a107..3aee80b 100644
--- a/lib/types.d.ts
+++ b/lib/types.d.ts
@@ -10,11 +10,11 @@ import { Http2ServerRequest, Http2ServerResponse } from 'http2';
 export interface Function {
   // The initialization function, called before the server is started
   // This function is optional and should be synchronous.
-  init?: () => any;
+  init?: () => any | Promise<any>;
 
   // The shutdown function, called after the server is stopped
   // This function is optional and should be synchronous.
-  shutdown?: () => any;
+  shutdown?: () => any | Promise<any>;
 
   // The liveness function, called to check if the server is alive
   // This function is optional and should return 200/OK if the server is alive.
diff --git a/lib/utils.js b/lib/utils.js
new file mode 100644
index 0000000..90d386c
--- /dev/null
+++ b/lib/utils.js
@@ -0,0 +1,9 @@
+function isPromise(p) {
+  if (typeof p === 'object' && typeof p.then === 'function') {
+    return true;
+  }
+
+  return false;
+}
+
+module.exports = exports = { isPromise };
diff --git a/package.json b/package.json
index 5d13af6..5a38d27 100644
--- a/package.json
+++ b/package.json
@@ -13,6 +13,7 @@
   "type": "commonjs",
   "scripts": {
     "lint": "eslint --ignore-path .gitignore .",
+    "fix-lint": "eslint --fix --ignore-path .gitignore .",
     "test": "npm run test:source && npm run test:types",
     "test:source": "nyc --reporter=lcovonly tape test/test*.js | colortape",
     "test:types": "tsd",
diff --git a/tsconfig.json b/tsconfig.json
index a245265..f4ef7be 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -1,9 +1,4 @@
 {
-    "include": [
-      "./index.d.ts",
-      "./lib/**/*.ts"
-    ],
-    "exclude": [
-      "node_modules"
-    ]
-  }
+  "include": ["./index.d.ts", "./lib/**/*.ts"],
+  "exclude": ["node_modules"],
+}