CCM-13555: Improve enqueue performance (#336)

Author: francisco-videira-nhs

lambdas/api-handler/src/services/__tests__/letter-operations.test.ts

@@ -192,26 +192,19 @@ describe("getLetterDataUrl function", () => {
   });
 });
 
+function makeUpdateLetterCommand(n: number): UpdateLetterCommand {
+  return {
+    id: `letter${n}`,
+    status: "PENDING",
+    supplierId: `testSupplier}`,
+  };
+}
+
 describe("enqueueLetterUpdateRequests function", () => {
   beforeEach(() => {
     jest.clearAllMocks();
   });
 
-  const updateLetterCommands: UpdateLetterCommand[] = [
-    {
-      id: "id1",
-      status: "REJECTED",
-      supplierId: "s1",
-      reasonCode: "123",
-      reasonText: "Reason text",
-    },
-    {
-      id: "id2",
-      status: "ACCEPTED",
-      supplierId: "s1",
-    },
-  ];
-
   it("should update the letter status successfully", async () => {
     const sqsClient = { send: jest.fn() } as unknown as SQSClient;
     const logger = { error: jest.fn() } as unknown as pino.Logger;
@@ -220,6 +213,13 @@ describe("enqueueLetterUpdateRequests function", () => {
     };
     const deps: Deps = { sqsClient, logger, env } as Deps;
 
+    const updateLetterCommands = Array.from({ length: 25 }, (_, i) =>
+      makeUpdateLetterCommand(i),
+    );
+
+    const sqsClientSendMock = sqsClient.send as jest.Mock;
+    sqsClientSendMock.mockResolvedValue({ Failed: [] });
+
     const result = await enqueueLetterUpdateRequests(
       updateLetterCommands,
       "correlationId1",
@@ -228,63 +228,55 @@ describe("enqueueLetterUpdateRequests function", () => {
 
     expect(result).toBeUndefined();
 
-    expect(deps.sqsClient.send).toHaveBeenNthCalledWith(
-      1,
-      expect.objectContaining({
-        input: {
-          QueueUrl: deps.env.QUEUE_URL,
-          MessageAttributes: {
-            CorrelationId: {
-              DataType: "String",
-              StringValue: "correlationId1",
-            },
-          },
-          MessageBody: JSON.stringify({
-            id: updateLetterCommands[0].id,
-            status: updateLetterCommands[0].status,
-            supplierId: updateLetterCommands[0].supplierId,
-            reasonCode: updateLetterCommands[0].reasonCode,
-            reasonText: updateLetterCommands[0].reasonText,
-          }),
-        },
-      }),
-    );
+    // processes 10 at a time (25 -> 10+10+5)
+    expect(sqsClientSendMock).toHaveBeenCalledTimes(3);
 
-    expect(deps.sqsClient.send).toHaveBeenNthCalledWith(
-      2,
-      expect.objectContaining({
-        input: {
-          QueueUrl: deps.env.QUEUE_URL,
-          MessageAttributes: {
-            CorrelationId: {
-              DataType: "String",
-              StringValue: "correlationId1",
-            },
-          },
-          MessageBody: JSON.stringify({
-            id: updateLetterCommands[1].id,
-            status: updateLetterCommands[1].status,
-            supplierId: updateLetterCommands[1].supplierId,
-          }),
-        },
-      }),
-    );
+    const firstCallArg = sqsClientSendMock.mock.calls[0][0];
+    const firstInput = firstCallArg.input;
+
+    expect(firstInput.QueueUrl).toBe(deps.env.QUEUE_URL);
+    expect(Array.isArray(firstInput.Entries)).toBe(true);
+    expect(firstInput.Entries.length).toBe(10);
+
+    expect(firstInput.Entries[0].Id).toBe("0-0");
+    expect(firstInput.Entries[9].Id).toBe("0-9");
+
+    expect(
+      firstInput.Entries[0].MessageAttributes.CorrelationId.StringValue,
+    ).toBe("correlationId1");
+
+    const parsed = JSON.parse(firstInput.Entries[0].MessageBody);
+    expect(parsed.id).toBe("letter0");
+
+    // check last batch had 5 entries
+    const thirdCallArg = sqsClientSendMock.mock.calls[2][0];
+    const thirdInput = thirdCallArg.input;
+    expect(thirdInput.Entries.length).toBe(5);
+    // ids in third batch should start "2-0"
+    expect(thirdInput.Entries[0].Id).toBe("2-0");
   });
 
-  it("should log error if enqueueing fails", async () => {
-    const mockError = new Error("error");
+  it("should log error when SendMessageBatch returns Failed entries", async () => {
     const sqsClient = {
       send: jest
         .fn()
-        .mockRejectedValueOnce(mockError)
-        .mockResolvedValueOnce({ MessageId: "m1" }),
+        .mockResolvedValueOnce({ Failed: [] }) // first batch succeeds
+        .mockResolvedValueOnce({
+          Failed: [
+            { Id: "1-1", SenderFault: false, Code: "Err", Message: "failed" },
+          ],
+        }),
     } as unknown as SQSClient;
     const logger = { error: jest.fn() } as unknown as pino.Logger;
     const env = {
       QUEUE_URL: "sqsUrl",
     };
     const deps: Deps = { sqsClient, logger, env } as Deps;
 
+    const updateLetterCommands = Array.from({ length: 12 }, (_, i) =>
+      makeUpdateLetterCommand(i),
+    );
+
     const result = await enqueueLetterUpdateRequests(
       updateLetterCommands,
       "correlationId1",
@@ -293,36 +285,44 @@ describe("enqueueLetterUpdateRequests function", () => {
 
     expect(result).toBeUndefined();
 
-    expect(deps.sqsClient.send).toHaveBeenNthCalledWith(
-      2,
-      expect.objectContaining({
-        input: {
-          QueueUrl: deps.env.QUEUE_URL,
-          MessageAttributes: {
-            CorrelationId: {
-              DataType: "String",
-              StringValue: "correlationId1",
-            },
-          },
-          MessageBody: JSON.stringify({
-            id: updateLetterCommands[1].id,
-            status: updateLetterCommands[1].status,
-            supplierId: updateLetterCommands[1].supplierId,
-          }),
-        },
-      }),
-    );
+    // 12 = 10 + 2
+    expect(deps.sqsClient.send).toHaveBeenCalledTimes(2);
 
     expect(deps.logger.error).toHaveBeenCalledTimes(1);
-    expect(deps.logger.error).toHaveBeenCalledWith(
-      {
-        err: mockError,
-        correlationId: "correlationId1",
-        letterId: updateLetterCommands[0].id,
-        letterStatus: updateLetterCommands[0].status,
-        supplierId: updateLetterCommands[0].supplierId,
-      },
-      "Error enqueuing letter status update",
+    const errorArgs = (deps.logger.error as jest.Mock).mock.calls[0][0];
+    expect(errorArgs.failed).toBeDefined();
+    expect(Array.isArray(errorArgs.failed)).toBe(true);
+    expect(errorArgs.failed[0].Id).toBe("1-1");
+  });
+
+  it("should log error if enqueueing fails", async () => {
+    const sqsClient = {
+      send: jest
+        .fn()
+        .mockResolvedValueOnce({ Failed: [] }) // batch 0
+        .mockImplementationOnce(() => {
+          throw new Error("some failure");
+        }) // batch 1
+        .mockResolvedValueOnce({ Failed: [] }), // batch 2
+    } as unknown as SQSClient;
+    const logger = { error: jest.fn() } as unknown as pino.Logger;
+    const env = {
+      QUEUE_URL: "sqsUrl",
+    };
+    const deps: Deps = { sqsClient, logger, env } as Deps;
+
+    const lettersToUpdate = Array.from({ length: 21 }, (_, i) =>
+      makeUpdateLetterCommand(i),
     );
+
+    await enqueueLetterUpdateRequests(lettersToUpdate, "correlationId1", deps);
+
+    // all 3 attempted
+    expect(deps.sqsClient.send).toHaveBeenCalledTimes(3);
+
+    expect(deps.logger.error).toHaveBeenCalledTimes(1);
+    const logged = (deps.logger.error as jest.Mock).mock.calls[0][0];
+    expect(logged.correlationId).toBe("correlationId1");
+    expect(logged.err).toBeInstanceOf(Error);
   });
 });

lambdas/api-handler/src/services/letter-operations.ts

@@ -1,7 +1,7 @@
 import { LetterBase, LetterRepository } from "@internal/datastore";
 import { getSignedUrl } from "@aws-sdk/s3-request-presigner";
 import { GetObjectCommand, S3Client } from "@aws-sdk/client-s3";
-import { SendMessageCommand } from "@aws-sdk/client-sqs";
+import { SendMessageBatchCommand } from "@aws-sdk/client-sqs";
 import NotFoundError from "../errors/not-found-error";
 import { UpdateLetterCommand } from "../contracts/letters";
 import { ApiErrorDetail } from "../contracts/errors";
@@ -81,36 +81,64 @@ export const getLetterDataUrl = async (
   }
 };
 
+function chunk(
+  arr: UpdateLetterCommand[],
+  size: number,
+): UpdateLetterCommand[][] {
+  const chunks: UpdateLetterCommand[][] = [];
+  for (let i = 0; i < arr.length; i += size)
+    chunks.push(arr.slice(i, i + size));
+  return chunks;
+}
+
 export async function enqueueLetterUpdateRequests(
   updateLetterCommands: UpdateLetterCommand[],
   correlationId: string,
   deps: Deps,
 ) {
-  const tasks = updateLetterCommands.map(
-    async (request: UpdateLetterCommand) => {
-      try {
-        const command = new SendMessageCommand({
-          QueueUrl: deps.env.QUEUE_URL,
+  const BATCH_SIZE = 10; // SQS SendMessageBatch max
+  const CONCURRENCY = 5; // number of parallel batch API calls
+
+  const batches = chunk(updateLetterCommands, BATCH_SIZE);
+
+  // send batches in groups with limited concurrency
+  // BATCH_SIZE * CONCURRENCY is the number of total updates / db calls in-flight
+  for (let i = 0; i < batches.length; i += CONCURRENCY) {
+    const window = batches.slice(i, i + CONCURRENCY);
+
+    await Promise.all(
+      window.map(async (batch, batchIdx) => {
+        const entries = batch.map((request, idx) => ({
+          Id: `${i + batchIdx}-${idx}`, // unique per batch entry
+          MessageBody: JSON.stringify(request),
           MessageAttributes: {
             CorrelationId: { DataType: "String", StringValue: correlationId },
           },
-          MessageBody: JSON.stringify(request),
+        }));
+
+        const cmd = new SendMessageBatchCommand({
+          QueueUrl: deps.env.QUEUE_URL,
+          Entries: entries,
         });
-        await deps.sqsClient.send(command);
-      } catch (error) {
-        deps.logger.error(
-          {
-            err: error,
-            correlationId,
-            letterId: request.id,
-            letterStatus: request.status,
-            supplierId: request.supplierId,
-          },
-          "Error enqueuing letter status update",
-        );
-      }
-    },
-  );
 
-  await Promise.all(tasks);
+        try {
+          const result = await deps.sqsClient.send(cmd);
+          if (result.Failed && result.Failed.length > 0) {
+            deps.logger.error(
+              { failed: result.Failed },
+              "Some batch entries failed",
+            );
+          }
+        } catch (error) {
+          deps.logger.error(
+            {
+              err: error,
+              correlationId,
+            },
+            "Error enqueuing letter status updates",
+          );
+        }
+      }),
+    );
+  }
 }