Drain jobs cleanly on shutdown

ts/util/StartupQueue.ts

@@ -1,20 +1,26 @@
 // Copyright 2021 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
+import PQueue from 'p-queue';
 import * as Errors from '../types/errors';
 import * as log from '../logging/log';
 
 type EntryType = Readonly<{
   value: number;
-  callback(): void;
+  callback(): Promise<void>;
 }>;
 
 let startupProcessingQueue: StartupQueue | undefined;
 
 export class StartupQueue {
   private readonly map = new Map<string, EntryType>();
+  private readonly running: PQueue = new PQueue({
+    // mostly io-bound work that is not very parallelizable
+    // small number should be sufficient
+    concurrency: 5,
+  });
 
-  public add(id: string, value: number, f: () => void): void {
+  public add(id: string, value: number, f: () => Promise<void>): void {
     const existing = this.map.get(id);
     if (existing && existing.value >= value) {
       return;
@@ -30,26 +36,36 @@ export class StartupQueue {
     this.map.clear();
 
     for (const { callback } of values) {
-      try {
-        callback();
-      } catch (error) {
-        log.error(
-          'StartupQueue: Failed to process item due to error',
-          Errors.toLogFormat(error)
-        );
-      }
+      void this.running.add(async () => {
+        try {
+          return callback();
+        } catch (error) {
+          log.error(
+            'StartupQueue: Failed to process item due to error',
+            Errors.toLogFormat(error)
+          );
+          throw error;
+        }
+      });
     }
   }
 
+  private shutdown(): Promise<void> {
+    log.info(
+      `StartupQueue: Waiting for ${this.running.pending} tasks to drain`
+    );
+    return this.running.onIdle();
+  }
+
   static initialize(): void {
     startupProcessingQueue = new StartupQueue();
   }
 
-  static isReady(): boolean {
+  static isAvailable(): boolean {
     return Boolean(startupProcessingQueue);
   }
 
-  static add(id: string, value: number, f: () => void): void {
+  static add(id: string, value: number, f: () => Promise<void>): void {
     startupProcessingQueue?.add(id, value, f);
   }
 
@@ -57,4 +73,8 @@ export class StartupQueue {
     startupProcessingQueue?.flush();
     startupProcessingQueue = undefined;
   }
+
+  static async shutdown(): Promise<void> {
+    await startupProcessingQueue?.shutdown();
+  }
 }

ts/util/sleeper.ts

@@ -0,0 +1,98 @@
+// Copyright 2023 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import * as log from '../logging/log';
+import * as Errors from '../types/errors';
+
+/**
+ * Provides a way to delay tasks
+ * but also a way to force sleeping tasks to immediately resolve/reject on shutdown
+ */
+export class Sleeper {
+  private shuttingDown = false;
+  private shutdownCallbacks: Set<() => void> = new Set();
+
+  /**
+   * delay by ms, careful when using on a loop if resolving on shutdown (default)
+   */
+  sleep(
+    ms: number,
+    reason: string,
+    options?: { resolveOnShutdown?: boolean }
+  ): Promise<void> {
+    log.info(`Sleeper: sleeping for ${ms}ms. Reason: ${reason}`);
+    const resolveOnShutdown = options?.resolveOnShutdown ?? true;
+
+    return new Promise((resolve, reject) => {
+      let timeout: NodeJS.Timeout | undefined;
+
+      const shutdownCallback = () => {
+        if (timeout) {
+          clearTimeout(timeout);
+        }
+        log.info(
+          `Sleeper: resolving sleep task on shutdown. Original reason: ${reason}`
+        );
+        if (resolveOnShutdown) {
+          setTimeout(resolve, 0);
+        } else {
+          setTimeout(() => {
+            reject(
+              new Error(
+                `Sleeper: rejecting sleep task during shutdown. Original reason: ${reason}`
+              )
+            );
+          }, 0);
+        }
+      };
+
+      if (this.shuttingDown) {
+        log.info(
+          `Sleeper: sleep called when shutdown is in progress, scheduling immediate ${
+            resolveOnShutdown ? 'resolution' : 'rejection'
+          }. Original reason: ${reason}`
+        );
+        shutdownCallback();
+        return;
+      }
+
+      timeout = setTimeout(() => {
+        resolve();
+        this.removeShutdownCallback(shutdownCallback);
+      }, ms);
+
+      this.addShutdownCallback(shutdownCallback);
+    });
+  }
+
+  private addShutdownCallback(callback: () => void) {
+    this.shutdownCallbacks.add(callback);
+  }
+
+  private removeShutdownCallback(callback: () => void) {
+    this.shutdownCallbacks.delete(callback);
+  }
+
+  shutdown(): void {
+    if (this.shuttingDown) {
+      return;
+    }
+    log.info(
+      `Sleeper: shutting down, settling ${this.shutdownCallbacks.size} in-progress sleep calls`
+    );
+    this.shuttingDown = true;
+    this.shutdownCallbacks.forEach(cb => {
+      try {
+        cb();
+      } catch (error) {
+        log.error(
+          'Sleeper: Error executing shutdown callback',
+          Errors.toLogFormat(error)
+        );
+      }
+    });
+    log.info('Sleeper: sleep tasks settled');
+  }
+}
+
+export const sleeper = new Sleeper();