refactor: create pre shutdown lifecycle

Signed-off-by: Camila Belo <camilaibs@gmail.com>

packages/backend-app-api/src/wiring/BackendInitializer.ts

@@ -65,18 +65,9 @@ const instanceRegistry = new (class InstanceRegistry {
     if (!this.#registered) {
       this.#registered = true;
 
-      process.addListener('SIGTERM', () => {
-        this.#exitHandler();
-      });
-      process.addListener('SIGINT', () => {
-        this.#exitHandler();
-      });
-      process.addListener('SIGQUIT', () => {
-        process.exit(0);
-      });
-      process.addListener('beforeExit', () => {
-        this.#exitHandler();
-      });
+      process.addListener('SIGTERM', this.#exitHandler);
+      process.addListener('SIGINT', this.#exitHandler);
+      process.addListener('beforeExit', this.#exitHandler);
     }
 
     this.#instances.add(instance);
@@ -88,18 +79,12 @@ const instanceRegistry = new (class InstanceRegistry {
 
   #exitHandler = async () => {
     try {
-      // This signals to the healthcheck service that the process is shutting down
-      process.exitCode = 0;
-
-      const results = await Promise.race([
-        // Give the backend 30 seconds to shut down, then force exit
-        new Promise(resolve => setTimeout(resolve, 30000)),
-        Promise.allSettled(Array.from(this.#instances).map(b => b.stop())),
-      ]);
-
-      const errors = Array.isArray(results)
-        ? results.flatMap(r => (r.status === 'rejected' ? [r.reason] : []))
-        : [];
+      const results = await Promise.allSettled(
+        Array.from(this.#instances).map(b => b.stop()),
+      );
+      const errors = results.flatMap(r =>
+        r.status === 'rejected' ? [r.reason] : [],
+      );
 
       if (errors.length > 0) {
         for (const error of errors) {
@@ -464,6 +449,11 @@ export class BackendInitializer {
       // The startup failed, but we may still want to do cleanup so we continue silently
     }
 
+    const rootLifecycleService = await this.#getRootLifecycleImpl();
+
+    // Root services like the health one need to immediatelly be notified of the shutdown
+    await rootLifecycleService.preShutdown();
+
     // Get all plugins.
     const allPlugins = new Set<string>();
     for (const feature of this.#registrations) {
@@ -483,14 +473,14 @@ export class BackendInitializer {
     );
 
     // Once all plugin shutdown hooks are done, run root shutdown hooks.
-    const lifecycleService = await this.#getRootLifecycleImpl();
-    await lifecycleService.shutdown();
+    await rootLifecycleService.shutdown();
   }
 
   // Bit of a hacky way to grab the lifecycle services, potentially find a nicer way to do this
   async #getRootLifecycleImpl(): Promise<
     RootLifecycleService & {
       startup(): Promise<void>;
+      preShutdown(): Promise<void>;
       shutdown(): Promise<void>;
     }
   > {

packages/backend-app-api/src/wiring/createSpecializedBackend.test.ts

@@ -36,6 +36,7 @@ describe('createSpecializedBackend', () => {
             deps: {},
             factory: async () => ({
               addStartupHook: () => {},
+              addBeforeShutdownHook: () => {},
               addShutdownHook: () => {},
             }),
           }),
@@ -44,6 +45,7 @@ describe('createSpecializedBackend', () => {
             deps: {},
             factory: async () => ({
               addStartupHook: () => {},
+              addBeforeShutdownHook: () => {},
               addShutdownHook: () => {},
             }),
           }),

packages/backend-defaults/src/CreateBackend.test.ts

@@ -47,6 +47,7 @@ describe('createBackend', () => {
         deps: {},
         factory: async () => ({
           addStartupHook: () => {},
+          addBeforeShutdownHook: () => {},
           addShutdownHook: () => {},
         }),
       }),
@@ -57,6 +58,7 @@ describe('createBackend', () => {
         deps: {},
         factory: async () => ({
           addStartupHook: () => {},
+          addBeforeShutdownHook: () => {},
           addShutdownHook: () => {},
         }),
       }),

packages/backend-defaults/src/entrypoints/rootHealth/rootHealthServiceFactory.test.ts

@@ -51,31 +51,13 @@ describe('DefaultRootHealthService', () => {
       });
     });
 
-    it(`should return a 503 response if the server is stopping`, async () => {
-      const service = new DefaultRootHealthService({
-        lifecycle: mockServices.rootLifecycle.mock(),
-      });
-
-      process.exitCode = 1;
-
-      await expect(service.getReadiness()).resolves.toEqual({
-        status: 503,
-        payload: {
-          message: 'Backend has not started yet',
-          status: 'error',
-        },
-      });
-
-      process.exitCode = undefined; // Reset exitCode for other tests
-    });
-
     it(`should return a 500 response if the server has stopped`, async () => {
       let mockServerStartedFn = () => {};
-      let mockServerStoppedFn = () => {};
+      let mockServerBeforeStoppedFn = () => {};
 
       const lifecycle = mockServices.rootLifecycle.mock({
         addStartupHook: jest.fn(fn => (mockServerStartedFn = fn)),
-        addShutdownHook: jest.fn(fn => (mockServerStoppedFn = fn)),
+        addBeforeShutdownHook: jest.fn(fn => (mockServerBeforeStoppedFn = fn)),
       });
 
       const service = new DefaultRootHealthService({
@@ -83,7 +65,7 @@ describe('DefaultRootHealthService', () => {
       });
 
       mockServerStartedFn();
-      mockServerStoppedFn();
+      mockServerBeforeStoppedFn();
       await expect(service.getReadiness()).resolves.toEqual({
         status: 503,
         payload: {

packages/backend-defaults/src/entrypoints/rootHealth/rootHealthServiceFactory.ts

@@ -29,7 +29,7 @@ export class DefaultRootHealthService implements RootHealthService {
     options.lifecycle.addStartupHook(() => {
       this.#isRunning = true;
     });
-    options.lifecycle.addShutdownHook(() => {
+    options.lifecycle.addPreShutdownHook(() => {
       this.#isRunning = false;
     });
   }
@@ -39,8 +39,7 @@ export class DefaultRootHealthService implements RootHealthService {
   }
 
   async getReadiness(): Promise<{ status: number; payload?: any }> {
-    // If the process has been told to exit, or if the backend has not started yet
-    if (process.exitCode !== undefined || !this.#isRunning) {
+    if (!this.#isRunning) {
       return {
         status: 503,
         payload: { message: 'Backend has not started yet', status: 'error' },

packages/backend-defaults/src/entrypoints/rootHttpRouter/rootHttpRouterServiceFactory.ts

@@ -120,7 +120,7 @@ const rootHttpRouterServiceFactoryWithOptions = (
         },
       });
 
-      lifecycle.addShutdownHook(() => server.stop());
+      lifecycle.addPreShutdownHook(() => server.stop());
 
       await server.start();
 

packages/backend-defaults/src/entrypoints/rootLifecycle/rootLifecycleServiceFactory.ts

@@ -65,6 +65,37 @@ export class BackendLifecycleImpl implements RootLifecycleService {
     );
   }
 
+  #hasPreShutdown = false;
+  #preShutdownTasks: Array<{ hook: () => void }> = [];
+
+  addPreShutdownHook(hook: () => void): void {
+    if (this.#hasPreShutdown) {
+      throw new Error('Attempted to add pre shutdown hook after pre shutdown');
+    }
+    this.#preShutdownTasks.push({ hook });
+  }
+
+  async preShutdown(): Promise<void> {
+    if (this.#hasPreShutdown) {
+      return;
+    }
+    this.#hasPreShutdown = true;
+
+    this.logger.debug(
+      `Running ${this.#preShutdownTasks.length} pre shutdown tasks...`,
+    );
+    await Promise.all(
+      this.#preShutdownTasks.map(async ({ hook }) => {
+        try {
+          await hook();
+          this.logger.debug(`Pre shutdown hook succeeded`);
+        } catch (error) {
+          this.logger.error(`Pre shutdown hook failed, ${error}`);
+        }
+      }),
+    );
+  }
+
   #hasShutdown = false;
   #shutdownTasks: Array<{
     hook: LifecycleServiceShutdownHook;

packages/backend-defaults/src/entrypoints/scheduler/lib/PluginTaskSchedulerImpl.test.ts

@@ -55,7 +55,11 @@ describe('PluginTaskManagerImpl', () => {
     const manager = new PluginTaskSchedulerImpl(
       async () => knex,
       mockServices.logger.mock(),
-      { addShutdownHook, addStartupHook: jest.fn() },
+      {
+        addShutdownHook,
+        addBeforeShutdownHook: jest.fn(),
+        addStartupHook: jest.fn(),
+      },
     );
     return { knex, manager };
   }

packages/backend-plugin-api/src/services/definitions/RootLifecycleService.ts

@@ -23,4 +23,6 @@ import { LifecycleService } from './LifecycleService';
  *
  * @public
  */
-export interface RootLifecycleService extends LifecycleService {}
+export interface RootLifecycleService extends LifecycleService {
+  addPreShutdownHook(hook: () => void): void;
+}

packages/backend-test-utils/src/next/services/mockServices.ts

@@ -472,6 +472,7 @@ export namespace mockServices {
     export const factory = () => rootLifecycleServiceFactory;
     export const mock = simpleMock(coreServices.rootLifecycle, () => ({
       addShutdownHook: jest.fn(),
+      addBeforeShutdownHook: jest.fn(),
       addStartupHook: jest.fn(),
     }));
   }