From d5bc3585aaab44b4cd819bb31d9bb0a479dd4220 Mon Sep 17 00:00:00 2001
From: Christopher Kolstad <chriswk@getunleash.ai>
Date: Fri, 15 Mar 2024 14:05:35 +0100
Subject: [PATCH] feat: Added shutdown hook option (#6585)

An entrypoint for enterprise to register a hook which will be called
before the database and scheduler services are torn down. That way
enterprise can also perform graceful shutdown.
---
 src/lib/server-impl.ts  | 9 ++++++++-
 src/lib/types/option.ts | 3 ++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/lib/server-impl.ts b/src/lib/server-impl.ts
index 4456bc537a..52f593e40f 100644
--- a/src/lib/server-impl.ts
+++ b/src/lib/server-impl.ts
@@ -12,13 +12,13 @@ import { createDb } from './db/db-pool';
 import sessionDb from './middleware/session-db';
 // Types
 import {
+    CustomAuthHandler,
     IAuthType,
     IUnleash,
     IUnleashConfig,
     IUnleashOptions,
     IUnleashServices,
     RoleName,
-    CustomAuthHandler,
     SYSTEM_USER,
 } from './types';
 
@@ -59,6 +59,13 @@ async function createApp(
             const stopServer = promisify(server.stop);
             await stopServer();
         }
+        if (typeof config.shutdownHook === 'function') {
+            try {
+                await config.shutdownHook();
+            } catch (e) {
+                logger.error('Failure when executing shutdown hook', e);
+            }
+        }
         services.schedulerService.stop();
         services.addonService.destroy();
         await db.destroy();
diff --git a/src/lib/types/option.ts b/src/lib/types/option.ts
index c3bf8e856d..4930d2a8ca 100644
--- a/src/lib/types/option.ts
+++ b/src/lib/types/option.ts
@@ -2,7 +2,7 @@ import { Express } from 'express';
 import EventEmitter from 'events';
 import { LogLevel, LogProvider } from '../logger';
 import { ILegacyApiTokenCreate } from './models/api-token';
-import { IFlagResolver, IExperimentalOptions, IFlags } from './experimental';
+import { IExperimentalOptions, IFlagResolver, IFlags } from './experimental';
 import SMTPTransport from 'nodemailer/lib/smtp-transport';
 import { IUnleashServices } from './services';
 import { ResourceLimitsSchema } from '../openapi/spec/resource-limits-schema';
@@ -231,6 +231,7 @@ export interface IUnleashConfig {
     enableOAS: boolean;
     preHook?: Function;
     preRouterHook?: Function;
+    shutdownHook?: Function;
     enterpriseVersion?: string;
     eventBus: EventEmitter;
     environmentEnableOverrides?: string[];