From d9e0a1e0429a1bba4e49461f939d3f7639d0c6c2 Mon Sep 17 00:00:00 2001
From: Guillaume Quittet <guillaume.quittet@gmail.com>
Date: Fri, 7 Jun 2024 17:29:01 +0200
Subject: [PATCH] feat: make close function as public

The goal is to allow end-user to close the server and trigger
the close promises callbacks.

Issue: #25
---
 src/core/improvedServer.ts      | 41 ++++++++++++++++++++++++---------
 src/core/index.ts               |  5 ++--
 src/index.ts                    |  1 +
 src/interface/core.ts           |  2 +-
 src/interface/gracefulServer.ts |  1 +
 src/util/init.ts                |  2 +-
 6 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/core/improvedServer.ts b/src/core/improvedServer.ts
index 0a4f8c2..5e66956 100644
--- a/src/core/improvedServer.ts
+++ b/src/core/improvedServer.ts
@@ -1,16 +1,16 @@
 import type { Server } from "#interface/server";
 import type { IStatus } from "#interface/status";
 import type http from "node:http";
+import { exit } from "node:process";
 import config from "#config/index";
 import SocketsPool from "#core/socketsPool";
+import State from "#core/state";
 import onRequest from "#util/onRequest";
+import sleep from "#util/sleep";
 
 const { livenessEndpoint, readinessEndpoint } = config;
 
-const improvedServer = <TServer extends Server>(
-  server: TServer,
-  serverStatus: IStatus,
-): TServer & { stop: () => Promise<void> } => {
+const improvedServer = <TServer extends Server>(server: TServer, serverStatus: IStatus) => {
   const { healthCheck, kubernetes } = config;
   const socketsPool = SocketsPool();
   const secureSocketsPool = SocketsPool();
@@ -65,13 +65,30 @@ const improvedServer = <TServer extends Server>(
     );
   }
 
-  const stop = async (): Promise<void> => {
+  const stop = async (
+    args: { value: number; body?: Error; type?: string } = { value: 0 },
+  ): Promise<void> => {
     if (!server.listening || stopping) {
       return;
     }
 
     stopping = true;
 
+    const { timeout, closePromises } = config;
+
+    let error: Error | undefined;
+    if (args.body && args.body.message) {
+      error = args.body;
+    } else if (args.type) {
+      error = new Error(args.type);
+    }
+
+    serverStatus.set(State.SHUTTING_DOWN, error);
+
+    await sleep(timeout);
+
+    await Promise.all(closePromises.map((closePromise) => closePromise()));
+
     server.removeAllListeners("request");
     server.on("request", (_: http.IncomingMessage, res: http.ServerResponse) => {
       if (!res.headersSent) res.setHeader("connection", "close");
@@ -79,14 +96,16 @@ const improvedServer = <TServer extends Server>(
 
     await Promise.all([socketsPool.closeAll(), secureSocketsPool.closeAll()]);
 
-    return new Promise((resolve, reject) => {
-      server.close((error) => {
-        if (error) {
-          reject(error);
-        } else {
-          resolve();
+    await new Promise((resolve, reject) => {
+      server.close((err) => {
+        if (err) {
+          return reject(err);
         }
+        return resolve(error);
       });
+
+      serverStatus.set(State.SHUTDOWN, error);
+      exit(args.value);
     });
   };
 
diff --git a/src/core/index.ts b/src/core/index.ts
index fc0fc33..ad89453 100644
--- a/src/core/index.ts
+++ b/src/core/index.ts
@@ -4,7 +4,6 @@ import { EventEmitter } from "events";
 import ImprovedServer from "#core/improvedServer";
 import Status from "#core/status";
 import init from "#util/init";
-import shutdown from "#util/shutdown";
 
 const core = (server: Server): ICore => {
   const _emitter = new EventEmitter();
@@ -15,8 +14,8 @@ const core = (server: Server): ICore => {
     init: function () {
       return init(this);
     },
-    shutdown: function (type: string, value: number, error?: Error) {
-      return shutdown(_server, this)(type, value, error);
+    stop: function (args?: { type?: string; value: number; body?: Error }) {
+      return _server.stop(args);
     },
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     on: (name: string, callback: (...args: any[]) => void) => _emitter.on(name, callback),
diff --git a/src/index.ts b/src/index.ts
index b62d753..4a1b549 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -12,6 +12,7 @@ const buildGracefulServer = (server: Server, options?: IGracefulServerOptions):
     isReady: () => gracefulServer.status.isReady(),
     setReady: () => gracefulServer.status.setReady(),
     on: gracefulServer.on,
+    stop: async () => await gracefulServer.stop(),
   };
 };
 
diff --git a/src/interface/core.ts b/src/interface/core.ts
index cb8efdf..61bc0ff 100644
--- a/src/interface/core.ts
+++ b/src/interface/core.ts
@@ -4,7 +4,7 @@ import type { EventEmitter } from "events";
 export type ICore = {
   status: IStatus;
   init: () => ICore;
-  shutdown: (type: string, value: number, error?: Error) => Promise<void>;
+  stop: (args?: { type?: string; value: number; body?: Error }) => Promise<void>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   on: (name: string, callback: (...args: any[]) => void) => EventEmitter;
 };
diff --git a/src/interface/gracefulServer.ts b/src/interface/gracefulServer.ts
index db48304..0ec58c5 100644
--- a/src/interface/gracefulServer.ts
+++ b/src/interface/gracefulServer.ts
@@ -5,4 +5,5 @@ export type IGracefulServer = {
   setReady: () => void;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   on: (name: string, callback: (...args: any[]) => void) => EventEmitter;
+  stop: () => Promise<void>;
 };
diff --git a/src/util/init.ts b/src/util/init.ts
index 23733d1..152354d 100644
--- a/src/util/init.ts
+++ b/src/util/init.ts
@@ -6,7 +6,7 @@ import signals from "#core/signals";
 const init = (parent: ICore) => {
   for (const signal of signals) {
     (process as NodeJS.EventEmitter).on(signal.type, async (body) => {
-      await parent.shutdown(signal.type, signal.code, body);
+      await parent.stop({ type: signal.type, value: signal.code, body });
     });
   }
   return parent;