From 973278ffe1abaf2321d631227e19a7549374cf3f Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Sun, 26 May 2024 19:22:29 +0200
Subject: [PATCH] events-backend: graceful shutdown of listener

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 .../src/service/EventsPlugin.ts               | 10 ++++++-
 .../src/service/hub/DatabaseEventBusStore.ts  | 29 +++++++++++++++++--
 .../src/service/hub/createEventBusRouter.ts   |  4 +++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/plugins/events-backend/src/service/EventsPlugin.ts b/plugins/events-backend/src/service/EventsPlugin.ts
index e8769b4a09..1369ee98ca 100644
--- a/plugins/events-backend/src/service/EventsPlugin.ts
+++ b/plugins/events-backend/src/service/EventsPlugin.ts
@@ -78,6 +78,7 @@ export const eventsPlugin = createBackendPlugin({
         database: coreServices.database,
         logger: coreServices.logger,
         scheduler: coreServices.scheduler,
+        lifecycle: coreServices.lifecycle,
         httpAuth: coreServices.httpAuth,
         router: coreServices.httpRouter,
       },
@@ -87,6 +88,7 @@ export const eventsPlugin = createBackendPlugin({
         database,
         logger,
         scheduler,
+        lifecycle,
         httpAuth,
         router,
       }) {
@@ -107,7 +109,13 @@ export const eventsPlugin = createBackendPlugin({
         http.bind(eventsRouter);
 
         router.use(
-          await createEventBusRouter({ database, logger, httpAuth, scheduler }),
+          await createEventBusRouter({
+            database,
+            logger,
+            httpAuth,
+            scheduler,
+            lifecycle,
+          }),
         );
 
         router.use(eventsRouter);
diff --git a/plugins/events-backend/src/service/hub/DatabaseEventBusStore.ts b/plugins/events-backend/src/service/hub/DatabaseEventBusStore.ts
index 720b62bb9c..dfcf10ece4 100644
--- a/plugins/events-backend/src/service/hub/DatabaseEventBusStore.ts
+++ b/plugins/events-backend/src/service/hub/DatabaseEventBusStore.ts
@@ -18,6 +18,7 @@ import { EventBusStore } from './types';
 import { Knex } from 'knex';
 import {
   DatabaseService,
+  LifecycleService,
   LoggerService,
   SchedulerService,
   resolvePackagePath,
@@ -86,6 +87,7 @@ class DatabaseEventBusListener {
     reject: (error: Error) => void;
   }>();
 
+  #isShuttingDown = false;
   #connPromise?: Promise<InternalDbConnection>;
   #connTimeout?: NodeJS.Timeout;
   #keepaliveInterval?: NodeJS.Timeout;
@@ -118,6 +120,17 @@ class DatabaseEventBusListener {
     });
   }
 
+  async shutdown() {
+    if (this.#isShuttingDown) {
+      return;
+    }
+    this.#isShuttingDown = true;
+    const conn = await this.#connPromise?.catch(() => undefined);
+    if (conn) {
+      this.#destroyConnection(conn);
+    }
+  }
+
   #handleNotify(topic: string) {
     this.#logger.debug(`Listener received notification for topic '${topic}'`);
     for (const l of this.#listeners) {
@@ -169,6 +182,9 @@ class DatabaseEventBusListener {
   }
 
   async #ensureConnection() {
+    if (this.#isShuttingDown) {
+      throw new Error('Listener is shutting down');
+    }
     if (this.#connPromise) {
       await this.#connPromise;
       return;
@@ -226,6 +242,7 @@ export class DatabaseEventBusStore implements EventBusStore {
     database: DatabaseService;
     logger: LoggerService;
     scheduler: SchedulerService;
+    lifecycle: LifecycleService;
     window?: {
       /** Events within this range will never be deleted */
       minAge?: HumanDuration;
@@ -250,9 +267,12 @@ export class DatabaseEventBusStore implements EventBusStore {
       options.logger.info('DatabaseEventBusStore migrations ran successfully');
     }
 
+    const listener = new DatabaseEventBusListener(db.client, options.logger);
+
     const store = new DatabaseEventBusStore(
       db,
       options.logger,
+      listener,
       options.window?.size ?? WINDOW_SIZE_DEFAULT,
       durationToMilliseconds(options.window?.minAge ?? WINDOW_MIN_AGE_DEFAULT),
       durationToMilliseconds(options.window?.maxAge ?? WINDOW_MAX_AGE_DEFAULT),
@@ -265,29 +285,34 @@ export class DatabaseEventBusStore implements EventBusStore {
       fn: store.#cleanup,
     });
 
+    options.lifecycle.addShutdownHook(async () => {
+      await listener.shutdown();
+    });
+
     return store;
   }
 
   readonly #db: Knex;
   readonly #logger: LoggerService;
+  readonly #listener: DatabaseEventBusListener;
   readonly #windowSize: number;
   readonly #windowMinAge: number;
   readonly #windowMaxAge: number;
-  readonly #listener: DatabaseEventBusListener;
 
   private constructor(
     db: Knex,
     logger: LoggerService,
+    listener: DatabaseEventBusListener,
     windowSize: number,
     windowMinAge: number,
     windowMaxAge: number,
   ) {
     this.#db = db;
     this.#logger = logger;
+    this.#listener = listener;
     this.#windowSize = windowSize;
     this.#windowMinAge = windowMinAge;
     this.#windowMaxAge = windowMaxAge;
-    this.#listener = new DatabaseEventBusListener(db.client, logger);
   }
 
   async publish(options: {
diff --git a/plugins/events-backend/src/service/hub/createEventBusRouter.ts b/plugins/events-backend/src/service/hub/createEventBusRouter.ts
index a5f67538bf..1a702c624b 100644
--- a/plugins/events-backend/src/service/hub/createEventBusRouter.ts
+++ b/plugins/events-backend/src/service/hub/createEventBusRouter.ts
@@ -17,6 +17,7 @@
 import {
   DatabaseService,
   HttpAuthService,
+  LifecycleService,
   LoggerService,
   SchedulerService,
 } from '@backstage/backend-plugin-api';
@@ -121,6 +122,7 @@ export async function createEventBusRouter(options: {
   logger: LoggerService;
   database: DatabaseService;
   scheduler: SchedulerService;
+  lifecycle: LifecycleService;
   httpAuth: HttpAuthService;
   notifyTimeoutMs?: number; // for testing
 }): Promise<Handler> {
@@ -128,6 +130,7 @@ export async function createEventBusRouter(options: {
     database,
     httpAuth,
     scheduler,
+    lifecycle,
     notifyTimeoutMs = DEFAULT_NOTIFY_TIMEOUT_MS,
   } = options;
   const logger = options.logger.child({ type: 'EventBus' });
@@ -141,6 +144,7 @@ export async function createEventBusRouter(options: {
       database,
       logger,
       scheduler,
+      lifecycle,
     });
   } else {
     logger.info('Database is not PostgreSQL, using memory store');