Ensure messages are sent in order, even with errors

ts/jobs/JobQueue.ts

@@ -1,6 +1,7 @@
 // Copyright 2021 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
+import PQueue from 'p-queue';
 import { v4 as uuid } from 'uuid';
 import { noop } from 'lodash';
 
@@ -60,6 +61,8 @@ export abstract class JobQueue<T> {
     }
   >();
 
+  private readonly defaultInMemoryQueue = new PQueue();
+
   private started = false;
 
   constructor(options: Readonly<JobQueueOptions>) {
@@ -172,6 +175,10 @@ export abstract class JobQueue<T> {
     return new Job(id, timestamp, this.queueType, data, completion);
   }
 
+  protected getInMemoryQueue(_parsedJob: ParsedJob<T>): PQueue {
+    return this.defaultInMemoryQueue;
+  }
+
   private async enqueueStoredJob(storedJob: Readonly<StoredJob>) {
     assert(
       storedJob.queueType === this.queueType,
@@ -205,14 +212,17 @@ export abstract class JobQueue<T> {
       data: parsedData,
     };
 
+    const queue: PQueue = this.getInMemoryQueue(parsedJob);
+
     const logger = new JobLogger(parsedJob, this.logger);
 
-    let result:
+    const result:
       | undefined
       | { success: true }
-      | { success: false; err: unknown };
+      | { success: false; err: unknown } = await queue.add(async () => {
-
       for (let attempt = 1; attempt <= this.maxAttempts; attempt += 1) {
+        const isFinalAttempt = attempt === this.maxAttempts;
+
         logger.attempt = attempt;
 
         log.info(
@@ -223,20 +233,25 @@ export abstract class JobQueue<T> {
           //   than once at a time. Ideally, the job will succeed on the first attempt.
           // eslint-disable-next-line no-await-in-loop
           await this.run(parsedJob, { attempt, log: logger });
-        result = { success: true };
           log.info(
             `${this.logPrefix} job ${storedJob.id} succeeded on attempt ${attempt}`
           );
-        break;
+          return { success: true };
         } catch (err: unknown) {
-        result = { success: false, err };
           log.error(
             `${this.logPrefix} job ${
               storedJob.id
             } failed on attempt ${attempt}. ${Errors.toLogFormat(err)}`
           );
+          if (isFinalAttempt) {
+            return { success: false, err };
           }
         }
+      }
+
+      // This should never happen. See the assertion below.
+      return undefined;
+    });
 
     await this.store.delete(storedJob.id);
 

ts/jobs/normalMessageSendJobQueue.ts

@@ -95,25 +95,25 @@ export class NormalMessageSendJobQueue extends JobQueue<NormalMessageSendJobData
     return { messageId, conversationId };
   }
 
-  private getQueue(queueKey: string): PQueue {
+  protected getInMemoryQueue({
-    const existingQueue = this.queues.get(queueKey);
+    data,
+  }: Readonly<{ data: NormalMessageSendJobData }>): PQueue {
+    const { conversationId } = data;
+
+    const existingQueue = this.queues.get(conversationId);
     if (existingQueue) {
       return existingQueue;
     }
 
     const newQueue = new PQueue({ concurrency: 1 });
     newQueue.once('idle', () => {
-      this.queues.delete(queueKey);
+      this.queues.delete(conversationId);
     });
 
-    this.queues.set(queueKey, newQueue);
+    this.queues.set(conversationId, newQueue);
     return newQueue;
   }
 
-  private enqueue(queueKey: string, fn: () => Promise<void>): Promise<void> {
-    return this.getQueue(queueKey).add(fn);
-  }
-
   protected async run(
     {
       data,
@@ -121,9 +121,8 @@ export class NormalMessageSendJobQueue extends JobQueue<NormalMessageSendJobData
     }: Readonly<{ data: NormalMessageSendJobData; timestamp: number }>,
     { attempt, log }: Readonly<{ attempt: number; log: LoggerType }>
   ): Promise<void> {
-    const { messageId, conversationId } = data;
+    const { messageId } = data;
 
-    await this.enqueue(conversationId, async () => {
     const timeRemaining = timestamp + MAX_RETRY_TIME - Date.now();
     const isFinalAttempt = attempt >= MAX_ATTEMPTS;
 
@@ -174,9 +173,7 @@ export class NormalMessageSendJobQueue extends JobQueue<NormalMessageSendJobData
         };
 
     if (!shouldContinue) {
-        log.info(
+      log.info(`message ${messageId} ran out of time. Giving up on sending it`);
-          `message ${messageId} ran out of time. Giving up on sending it`
-        );
       await markMessageFailed(message, messageSendErrors);
       return;
     }
@@ -362,7 +359,6 @@ export class NormalMessageSendJobQueue extends JobQueue<NormalMessageSendJobData
 
       throw err;
     }
-    });
   }
 }
 

ts/test-node/jobs/JobQueue_test.ts

@@ -8,6 +8,7 @@ import EventEmitter, { once } from 'events';
 import { z } from 'zod';
 import { noop, groupBy } from 'lodash';
 import { v4 as uuid } from 'uuid';
+import PQueue from 'p-queue';
 import { JobError } from '../../jobs/JobError';
 import { TestJobQueueStore } from './TestJobQueueStore';
 import { missingCaseError } from '../../util/missingCaseError';
@@ -67,6 +68,98 @@ describe('JobQueue', () => {
       assert.isEmpty(store.storedJobs);
     });
 
+    it('by default, kicks off multiple jobs in parallel', async () => {
+      let activeJobCount = 0;
+      const eventBus = new EventEmitter();
+      const updateActiveJobCount = (incrementBy: number): void => {
+        activeJobCount += incrementBy;
+        eventBus.emit('updated');
+      };
+
+      class Queue extends JobQueue<number> {
+        parseData(data: unknown): number {
+          return z.number().parse(data);
+        }
+
+        async run(): Promise<void> {
+          try {
+            updateActiveJobCount(1);
+            await new Promise<void>(resolve => {
+              eventBus.on('updated', () => {
+                if (activeJobCount === 4) {
+                  eventBus.emit('got to 4');
+                  resolve();
+                }
+              });
+            });
+          } finally {
+            updateActiveJobCount(-1);
+          }
+        }
+      }
+
+      const store = new TestJobQueueStore();
+
+      const queue = new Queue({
+        store,
+        queueType: 'test queue',
+        maxAttempts: 100,
+      });
+      queue.streamJobs();
+
+      queue.add(1);
+      queue.add(2);
+      queue.add(3);
+      queue.add(4);
+
+      await once(eventBus, 'got to 4');
+    });
+
+    it('can override the in-memory queue', async () => {
+      let jobsAdded = 0;
+      const testQueue = new PQueue();
+      testQueue.on('add', () => {
+        jobsAdded += 1;
+      });
+
+      class Queue extends JobQueue<number> {
+        parseData(data: unknown): number {
+          return z.number().parse(data);
+        }
+
+        protected getInMemoryQueue(parsedJob: ParsedJob<number>): PQueue {
+          assert(
+            new Set([1, 2, 3, 4]).has(parsedJob.data),
+            'Bad data passed to `getInMemoryQueue`'
+          );
+          return testQueue;
+        }
+
+        run(): Promise<void> {
+          return Promise.resolve();
+        }
+      }
+
+      const store = new TestJobQueueStore();
+
+      const queue = new Queue({
+        store,
+        queueType: 'test queue',
+        maxAttempts: 100,
+      });
+      queue.streamJobs();
+
+      const jobs = await Promise.all([
+        queue.add(1),
+        queue.add(2),
+        queue.add(3),
+        queue.add(4),
+      ]);
+      await Promise.all(jobs.map(job => job.completion));
+
+      assert.strictEqual(jobsAdded, 4);
+    });
+
     it('writes jobs to the database correctly', async () => {
       const store = new TestJobQueueStore();