From 8bfe9f0f61aad548b6d129610b05b56aa9f93f2b Mon Sep 17 00:00:00 2001
From: Eric Peterson <ericpeterson@spotify.com>
Date: Tue, 1 Oct 2024 12:57:43 +0200
Subject: [PATCH] Ensure root lifecycle shutdown hooks happen after plugins'

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
---
 .../src/wiring/BackendInitializer.ts          | 24 ++++++++-
 .../lifecycle/lifecycleServiceFactory.ts      | 51 +++++++++++++++----
 2 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/packages/backend-app-api/src/wiring/BackendInitializer.ts b/packages/backend-app-api/src/wiring/BackendInitializer.ts
index 86771fbb74..e580109987 100644
--- a/packages/backend-app-api/src/wiring/BackendInitializer.ts
+++ b/packages/backend-app-api/src/wiring/BackendInitializer.ts
@@ -407,6 +407,26 @@ export class BackendInitializer {
       // The startup failed, but we may still want to do cleanup so we continue silently
     }
 
+    // Get all plugins.
+    const pluginMap = new Map<string, boolean>();
+    for (const feature of this.#registrations) {
+      for (const r of feature.getRegistrations()) {
+        if (r.type === 'plugin') {
+          pluginMap.set(r.pluginId, true);
+        }
+      }
+    }
+    const allPluginIds = Array.from(pluginMap.keys());
+
+    // Iterate through all plugins and run their shutdown hooks.
+    await Promise.allSettled(
+      allPluginIds.map(async pluginId => {
+        const lifecycleService = await this.#getPluginLifecycleImpl(pluginId);
+        await lifecycleService.shutdown();
+      }),
+    );
+
+    // Once all plugin shutdown hooks are done, run root shutdown hooks.
     const lifecycleService = await this.#getRootLifecycleImpl();
     await lifecycleService.shutdown();
   }
@@ -437,7 +457,9 @@ export class BackendInitializer {
 
   async #getPluginLifecycleImpl(
     pluginId: string,
-  ): Promise<LifecycleService & { startup(): Promise<void> }> {
+  ): Promise<
+    LifecycleService & { startup(): Promise<void>; shutdown(): Promise<void> }
+  > {
     const lifecycleService = await this.#serviceRegistry.get(
       coreServices.lifecycle,
       pluginId,
diff --git a/packages/backend-defaults/src/entrypoints/lifecycle/lifecycleServiceFactory.ts b/packages/backend-defaults/src/entrypoints/lifecycle/lifecycleServiceFactory.ts
index 192886ca61..02195d0cd0 100644
--- a/packages/backend-defaults/src/entrypoints/lifecycle/lifecycleServiceFactory.ts
+++ b/packages/backend-defaults/src/entrypoints/lifecycle/lifecycleServiceFactory.ts
@@ -22,7 +22,6 @@ import {
   LifecycleServiceStartupOptions,
   LoggerService,
   PluginMetadataService,
-  RootLifecycleService,
   coreServices,
   createServiceFactory,
 } from '@backstage/backend-plugin-api';
@@ -31,15 +30,19 @@ import {
 export class BackendPluginLifecycleImpl implements LifecycleService {
   constructor(
     private readonly logger: LoggerService,
-    private readonly rootLifecycle: RootLifecycleService,
     private readonly pluginMetadata: PluginMetadataService,
   ) {}
 
   #hasStarted = false;
+  #hasShutdown = false;
   #startupTasks: Array<{
     hook: LifecycleServiceStartupHook;
     options?: LifecycleServiceStartupOptions;
   }> = [];
+  #shutdownTasks: Array<{
+    hook: LifecycleServiceShutdownHook;
+    options?: LifecycleServiceShutdownOptions;
+  }> = [];
 
   addStartupHook(
     hook: LifecycleServiceStartupHook,
@@ -77,11 +80,42 @@ export class BackendPluginLifecycleImpl implements LifecycleService {
     hook: LifecycleServiceShutdownHook,
     options?: LifecycleServiceShutdownOptions,
   ): void {
+    if (this.#hasShutdown) {
+      throw new Error('Attempted to add shutdown hook after shutdown');
+    }
     const plugin = this.pluginMetadata.getId();
-    this.rootLifecycle.addShutdownHook(hook, {
-      logger: options?.logger?.child({ plugin }) ?? this.logger,
+    const logger = options?.logger?.child({ plugin }) ?? this.logger;
+    this.#shutdownTasks.push({
+      hook,
+      options: {
+        ...options,
+        logger,
+      },
     });
   }
+
+  async shutdown(): Promise<void> {
+    if (this.#hasShutdown) {
+      return;
+    }
+    this.#hasShutdown = true;
+
+    this.logger.debug(
+      `Running ${this.#shutdownTasks.length} plugin shutdown tasks...`,
+    );
+
+    await Promise.all(
+      this.#shutdownTasks.map(async ({ hook, options }) => {
+        const logger = options?.logger ?? this.logger;
+        try {
+          await hook();
+          logger.debug(`Plugin shutdown hook succeeded`);
+        } catch (error) {
+          logger.error(`Plugin shutdown hook failed, ${error}`);
+        }
+      }),
+    );
+  }
 }
 
 /**
@@ -97,14 +131,9 @@ export const lifecycleServiceFactory = createServiceFactory({
   service: coreServices.lifecycle,
   deps: {
     logger: coreServices.logger,
-    rootLifecycle: coreServices.rootLifecycle,
     pluginMetadata: coreServices.pluginMetadata,
   },
-  async factory({ rootLifecycle, logger, pluginMetadata }) {
-    return new BackendPluginLifecycleImpl(
-      logger,
-      rootLifecycle,
-      pluginMetadata,
-    );
+  async factory({ logger, pluginMetadata }) {
+    return new BackendPluginLifecycleImpl(logger, pluginMetadata);
   },
 });