chore: Improve graceful server shutdown (#4625)

* chore: Improve graceful server shutdown

* Replace node timers with custom promise timeout

server/utils/RateLimiter.ts

@@ -0,0 +1,67 @@
+import {
+  IRateLimiterStoreOptions,
+  RateLimiterRedis,
+} from "rate-limiter-flexible";
+import env from "@server/env";
+import Redis from "@server/redis";
+
+export default class RateLimiter {
+  constructor() {
+    throw Error(`Cannot instantiate class!`);
+  }
+
+  static readonly RATE_LIMITER_REDIS_KEY_PREFIX = "rl";
+
+  static readonly rateLimiterMap = new Map<string, RateLimiterRedis>();
+
+  static readonly defaultRateLimiter = new RateLimiterRedis({
+    storeClient: Redis.defaultClient,
+    points: env.RATE_LIMITER_REQUESTS,
+    duration: env.RATE_LIMITER_DURATION_WINDOW,
+    keyPrefix: this.RATE_LIMITER_REDIS_KEY_PREFIX,
+  });
+
+  static getRateLimiter(path: string): RateLimiterRedis {
+    return this.rateLimiterMap.get(path) || this.defaultRateLimiter;
+  }
+
+  static setRateLimiter(path: string, config: IRateLimiterStoreOptions): void {
+    const rateLimiter = new RateLimiterRedis(config);
+    this.rateLimiterMap.set(path, rateLimiter);
+  }
+
+  static hasRateLimiter(path: string): boolean {
+    return this.rateLimiterMap.has(path);
+  }
+}
+
+/**
+ * Re-useable configuration for rate limiter middleware.
+ */
+export const RateLimiterStrategy = {
+  /** Allows five requests per minute, per IP address */
+  FivePerMinute: {
+    duration: 60,
+    requests: 5,
+  },
+  /** Allows ten requests per minute, per IP address */
+  TenPerMinute: {
+    duration: 60,
+    requests: 10,
+  },
+  /** Allows one thousand requests per hour, per IP address */
+  OneThousandPerHour: {
+    duration: 3600,
+    requests: 1000,
+  },
+  /** Allows ten requests per hour, per IP address */
+  TenPerHour: {
+    duration: 3600,
+    requests: 10,
+  },
+  /** Allows five requests per hour, per IP address */
+  FivePerHour: {
+    duration: 3600,
+    requests: 5,
+  },
+};

server/utils/ShutdownHelper.ts

@@ -0,0 +1,96 @@
+import { groupBy } from "lodash";
+import Logger from "@server/logging/Logger";
+import { timeout } from "./timers";
+
+export enum ShutdownOrder {
+  first = 0,
+  normal = 1,
+  last = 2,
+}
+
+type Handler = {
+  name: string;
+  order: ShutdownOrder;
+  callback: () => Promise<unknown>;
+};
+
+export default class ShutdownHelper {
+  /**
+   * The amount of time to wait for connections to close before forcefully
+   * closing them. This allows for regular HTTP requests to complete but
+   * prevents long running requests from blocking shutdown.
+   */
+  public static readonly connectionGraceTimeout = 5 * 1000;
+
+  /**
+   * The maximum amount of time to wait for ongoing work to finish before
+   * force quitting the process. In the event of a force quit, the process
+   * will exit with a non-zero exit code.
+   */
+  public static readonly forceQuitTimeout = 60 * 1000;
+
+  /** Whether the server is currently shutting down */
+  private static isShuttingDown = false;
+
+  /** List of shutdown handlers to execute */
+  private static handlers: Handler[] = [];
+
+  /**
+   * Add a shutdown handler to be executed when the process is exiting
+   *
+   * @param name The name of the handler
+   * @param callback The callback to execute
+   */
+  public static add(
+    name: string,
+    order: ShutdownOrder,
+    callback: () => Promise<unknown>
+  ) {
+    this.handlers.push({ name, order, callback });
+  }
+
+  /**
+   * Exit the process after all shutdown handlers have completed
+   */
+  public static async execute() {
+    if (this.isShuttingDown) {
+      return;
+    }
+    this.isShuttingDown = true;
+
+    // Start the shutdown timer
+    void timeout(this.forceQuitTimeout).then(() => {
+      Logger.info("lifecycle", "Force quitting");
+      process.exit(1);
+    });
+
+    // Group handlers by order
+    const shutdownGroups = groupBy(this.handlers, "order");
+    const orderedKeys = Object.keys(shutdownGroups).sort();
+
+    // Execute handlers in order
+    for (const key of orderedKeys) {
+      Logger.debug("lifecycle", `Running shutdown group ${key}`);
+      const handlers = shutdownGroups[key];
+
+      await Promise.allSettled(
+        handlers.map(async (handler) => {
+          Logger.debug("lifecycle", `Running shutdown handler ${handler.name}`);
+
+          await handler.callback().catch((error) => {
+            Logger.error(
+              `Error inside shutdown handler ${handler.name}`,
+              error,
+              {
+                name: handler.name,
+              }
+            );
+          });
+        })
+      );
+    }
+
+    Logger.info("lifecycle", "Gracefully quitting");
+    process.exit(0);
+  }
+}

server/utils/queue.ts

@@ -3,6 +3,7 @@ import { snakeCase } from "lodash";
 import env from "@server/env";
 import Metrics from "@server/logging/Metrics";
 import Redis from "../redis";
+import ShutdownHelper, { ShutdownOrder } from "./ShutdownHelper";
 
 export function createQueue(
   name: string,
@@ -57,5 +58,9 @@ export function createQueue(
     }, 5 * 1000);
   }
 
+  ShutdownHelper.add(name, ShutdownOrder.normal, async () => {
+    await queue.close();
+  });
+
   return queue;
 }

server/utils/timers.ts

@@ -0,0 +1,8 @@
+/**
+ * Returns a promise that resolves after a specified number of milliseconds.
+ *
+ * @param [delay=1] The number of milliseconds to wait before fulfilling the promise.
+ */
+export function timeout(ms = 1) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}