Merge pull request #192 from bcgov/AI-1240

Add UploadNormalizationLimiter and overlap writes

Author: kmandryk

apps/backend-services/src/document/document.module.ts

@@ -1,14 +1,20 @@
 import { Module } from "@nestjs/common";
 import { BlobStorageModule } from "../blob-storage/blob-storage.module";
 import { TemporalModule } from "../temporal/temporal.module";
+import { UploadNormalizationLimiter } from "../upload/upload-normalization-limiter";
 import { DocumentController } from "./document.controller";
 import { DocumentService } from "./document.service";
 import { DocumentDbService } from "./document-db.service";
 import { PdfNormalizationService } from "./pdf-normalization.service";
 
 @Module({
   imports: [BlobStorageModule, TemporalModule],
-  providers: [DocumentDbService, DocumentService, PdfNormalizationService],
+  providers: [
+    DocumentDbService,
+    DocumentService,
+    PdfNormalizationService,
+    UploadNormalizationLimiter,
+  ],
   controllers: [DocumentController],
   exports: [DocumentService, PdfNormalizationService],
 })

apps/backend-services/src/document/document.service.spec.ts

@@ -6,6 +6,7 @@ import {
   BLOB_STORAGE,
   BlobStorageInterface,
 } from "../blob-storage/blob-storage.interface";
+import { UploadNormalizationLimiter } from "../upload/upload-normalization-limiter";
 import { DocumentService } from "./document.service";
 import { DocumentDbService } from "./document-db.service";
 import { PdfNormalizationService } from "./pdf-normalization.service";
@@ -15,7 +16,13 @@ describe("DocumentService", () => {
   let documentDbService: DocumentDbService;
   let blobStorage: BlobStorageInterface;
   let pdfNormalization: jest.Mocked<
-    Pick<PdfNormalizationService, "validateForUpload" | "normalizeToPdf">
+    Pick<
+      PdfNormalizationService,
+      "validateForUpload" | "normalizeToPdf" | "generateThumbnailWebp"
+    >
+  >;
+  let uploadNormalizationLimiter: jest.Mocked<
+    Pick<UploadNormalizationLimiter, "run">
   >;
 
   beforeEach(async () => {
@@ -24,7 +31,11 @@ describe("DocumentService", () => {
       normalizeToPdf: jest
         .fn()
         .mockImplementation((buf: Buffer) => Promise.resolve(Buffer.from(buf))),
+      generateThumbnailWebp: jest.fn().mockResolvedValue(Buffer.from("webp")),
     };
+    uploadNormalizationLimiter = {
+      run: jest.fn((task: () => Promise<unknown>) => task()),
+    } as jest.Mocked<Pick<UploadNormalizationLimiter, "run">>;
     documentDbService = {
       createDocument: jest.fn(),
       findDocument: jest.fn(),
@@ -35,7 +46,7 @@ describe("DocumentService", () => {
       upsertOcrResult: jest.fn(),
     } as any;
     blobStorage = {
-      write: jest.fn(),
+      write: jest.fn().mockResolvedValue(undefined),
       read: jest.fn(),
       exists: jest.fn(),
       delete: jest.fn(),
@@ -49,6 +60,10 @@ describe("DocumentService", () => {
         { provide: DocumentDbService, useValue: documentDbService },
         { provide: BLOB_STORAGE, useValue: blobStorage },
         { provide: PdfNormalizationService, useValue: pdfNormalization },
+        {
+          provide: UploadNormalizationLimiter,
+          useValue: uploadNormalizationLimiter,
+        },
       ],
     }).compile();
     service = module.get<DocumentService>(DocumentService);
@@ -103,6 +118,121 @@ describe("DocumentService", () => {
         expect.stringMatching(/^group-1\/ocr\/.+\/normalized\.pdf$/),
         expect.any(Buffer),
       );
+      expect(uploadNormalizationLimiter.run).toHaveBeenCalledTimes(1);
+      expect(pdfNormalization.generateThumbnailWebp).toHaveBeenCalledWith(
+        expect.any(Buffer),
+        "pdf",
+      );
+    });
+
+    it("writes the normalized blob only after the original write completes", async () => {
+      const pdfBytes = Buffer.from("%PDF-1.4\n%\xe2\xe3\xcf\xd3\n");
+      const base64 = pdfBytes.toString("base64");
+      const mockDoc = {
+        id: "1",
+        title: "Test",
+        original_filename: "file.pdf",
+        file_path: "documents/1/original.pdf",
+        normalized_file_path: "documents/1/normalized.pdf",
+        file_type: "pdf",
+        file_size: pdfBytes.length,
+        metadata: {},
+        source: "api",
+        status: DocumentStatus.ongoing_ocr,
+        created_at: new Date(),
+        updated_at: new Date(),
+        model_id: "test-model-id",
+        group_id: "group-1",
+      };
+      const writeOrder: string[] = [];
+      (documentDbService.createDocument as jest.Mock).mockResolvedValue(
+        mockDoc,
+      );
+      (blobStorage.write as jest.Mock).mockImplementation((key: string) => {
+        if (key.endsWith("/original.pdf")) {
+          writeOrder.push("original");
+          return new Promise<void>((resolve) => {
+            setTimeout(resolve, 10);
+          });
+        }
+        if (key.endsWith("/normalized.pdf")) {
+          writeOrder.push("normalized");
+        }
+        return Promise.resolve();
+      });
+
+      await service.uploadDocument(
+        "Test",
+        base64,
+        "pdf",
+        "file.pdf",
+        "test-model-id",
+        "group-1",
+        {},
+      );
+
+      expect(writeOrder).toEqual(["original", "normalized"]);
+    });
+
+    it("starts PDF normalization before the original blob write finishes", async () => {
+      const pdfBytes = Buffer.from("%PDF-1.4\n%\xe2\xe3\xcf\xd3\n");
+      const base64 = pdfBytes.toString("base64");
+      const mockDoc = {
+        id: "1",
+        title: "Test",
+        original_filename: "file.pdf",
+        file_path: "documents/1/original.pdf",
+        normalized_file_path: "documents/1/normalized.pdf",
+        file_type: "pdf",
+        file_size: pdfBytes.length,
+        metadata: {},
+        source: "api",
+        status: DocumentStatus.ongoing_ocr,
+        created_at: new Date(),
+        updated_at: new Date(),
+        model_id: "test-model-id",
+        group_id: "group-1",
+      };
+      let resolveOriginalWrite: (() => void) | undefined;
+      let normalizationStarted = false;
+      (documentDbService.createDocument as jest.Mock).mockResolvedValue(
+        mockDoc,
+      );
+      (blobStorage.write as jest.Mock).mockImplementation((key: string) => {
+        if (key.endsWith("/original.pdf")) {
+          return new Promise<void>((resolve) => {
+            resolveOriginalWrite = resolve;
+          });
+        }
+        return Promise.resolve();
+      });
+      pdfNormalization.normalizeToPdf.mockImplementation(
+        async (buf: Buffer) => {
+          normalizationStarted = true;
+          return Buffer.from(buf);
+        },
+      );
+
+      const uploadPromise = service.uploadDocument(
+        "Test",
+        base64,
+        "pdf",
+        "file.pdf",
+        "test-model-id",
+        "group-1",
+        {},
+      );
+
+      await Promise.resolve();
+
+      expect(normalizationStarted).toBe(true);
+      expect(documentDbService.createDocument).not.toHaveBeenCalled();
+
+      resolveOriginalWrite?.();
+      const result = await uploadPromise;
+
+      expect(result.kind).toBe("success");
+      expect(documentDbService.createDocument).toHaveBeenCalled();
     });
 
     it("should throw on invalid base64", async () => {

apps/backend-services/src/document/document.service.ts

@@ -17,6 +17,7 @@ import {
   BlobStorageInterface,
 } from "../blob-storage/blob-storage.interface";
 import { AppLoggerService } from "../logging/app-logger.service";
+import { UploadNormalizationLimiter } from "../upload/upload-normalization-limiter";
 import { DocumentDbService } from "./document-db.service";
 import type { DocumentData } from "./document-db.types";
 import { extensionForOriginalBlob } from "./original-blob-key.util";
@@ -55,6 +56,7 @@ export class DocumentService {
     @Inject(BLOB_STORAGE)
     private readonly blobStorage: BlobStorageInterface,
     private readonly pdfNormalization: PdfNormalizationService,
+    private readonly uploadNormalizationLimiter: UploadNormalizationLimiter,
     private readonly logger: AppLoggerService,
   ) {}
 
@@ -136,21 +138,30 @@ export class DocumentService {
         `original.${extension}`,
       );
 
-      await this.blobStorage.write(blobKey, fileBuffer);
-      this.logger.debug(`File saved to blob storage: ${blobKey}`);
-
       const normalizedKey = buildBlobFilePath(
         groupId,
         OperationCategory.OCR,
         [documentId],
         "normalized.pdf",
       );
+      const originalWrite = this.blobStorage.write(blobKey, fileBuffer);
+      // Normalization can run long enough for an early write failure to look
+      // unhandled; the original promise is still awaited below.
+      void originalWrite.catch(() => undefined);
+      this.logger.debug(`Original file write started: ${blobKey}`);
+
       try {
-        const pdfBuffer = await this.pdfNormalization.normalizeToPdf(
-          fileBuffer,
-          fileType,
+        const pdfBuffer = await this.uploadNormalizationLimiter.run(() =>
+          this.pdfNormalization.normalizeToPdf(fileBuffer, fileType),
         );
+        await originalWrite;
+        // Drop the decoded upload buffer before the normalized blob write so
+        // original bytes and pdf-lib workspace are not retained together.
+        fileBuffer = Buffer.alloc(0);
         await this.blobStorage.write(normalizedKey, pdfBuffer);
+        this.logger.debug(
+          `Files saved to blob storage: ${blobKey}, ${normalizedKey}`,
+        );
 
         const thumbnailKey = buildBlobFilePath(
           groupId,
@@ -160,10 +171,7 @@ export class DocumentService {
         );
         try {
           const thumbnailBuffer =
-            await this.pdfNormalization.generateThumbnailWebp(
-              fileBuffer,
-              fileType,
-            );
+            await this.pdfNormalization.generateThumbnailWebp(pdfBuffer, "pdf");
           await this.blobStorage.write(thumbnailKey, thumbnailBuffer);
           this.logger.debug(`Thumbnail saved: ${thumbnailKey}`);
         } catch (thumbErr) {
@@ -172,6 +180,12 @@ export class DocumentService {
           );
         }
       } catch (e) {
+        // Ensure the overlapping original write finished before we record failure.
+        // We intentionally keep the original blob (no rollback): the API returns
+        // conversion_failed, status is conversion_failed, normalized_file_path is
+        // null, OCR is not started, but GET .../download still serves the upload.
+        await originalWrite;
+
         if (e instanceof BadRequestException) {
           throw e;
         }

apps/backend-services/src/document/pdf-normalization.service.spec.ts

@@ -84,19 +84,6 @@ describe("PdfNormalizationService", () => {
       ).rejects.toThrow(BadRequestException);
     });
 
-    it("rejects PDF when header is valid but body is truncated", async () => {
-      const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {});
-      try {
-        const full = await minimalValidPdfBuffer();
-        const truncated = full.subarray(0, Math.floor(full.length / 2));
-        await expect(
-          service.validateForUpload(truncated, "pdf"),
-        ).rejects.toThrow(BadRequestException);
-      } finally {
-        warnSpy.mockRestore();
-      }
-    });
-
     it("accepts a minimal valid PDF", async () => {
       const buf = await minimalValidPdfBuffer();
       await expect(
@@ -125,6 +112,14 @@ describe("PdfNormalizationService", () => {
   });
 
   describe("normalizeToPdf", () => {
+    it("rejects PDF when header is valid but body is truncated", async () => {
+      const full = await minimalValidPdfBuffer();
+      const truncated = full.subarray(0, Math.floor(full.length / 2));
+      await expect(service.normalizeToPdf(truncated, "pdf")).rejects.toThrow(
+        BadRequestException,
+      );
+    });
+
     it("returns a copy of the buffer for pdf", async () => {
       const buf = await minimalValidPdfBuffer();
       const out = await service.normalizeToPdf(buf, "pdf");

apps/backend-services/src/document/pdf-normalization.service.ts

@@ -31,6 +31,11 @@ export class PdfNormalizationService {
 
   /**
    * Validates upload bytes before persisting the original blob.
+   *
+   * PDF/scan inputs are checked with a cheap magic-byte probe only; full
+   * pdf-lib parsing happens once in {@link normalizeToPdf} to avoid holding
+   * two document workspaces in memory at the same time.
+   *
    * @throws BadRequestException for corrupt or unsupported inputs.
    */
   async validateForUpload(fileBuffer: Buffer, fileType: string): Promise<void> {
@@ -44,13 +49,6 @@ export class PdfNormalizationService {
           "The file is not a valid PDF or is corrupted.",
         );
       }
-      try {
-        await PDFDocument.load(fileBuffer, { ignoreEncryption: true });
-      } catch {
-        throw new BadRequestException(
-          "The file is not a valid PDF or is corrupted.",
-        );
-      }
       return;
     }
     if (ft === "image") {
@@ -236,7 +234,14 @@ export class PdfNormalizationService {
    * Keep both sites in sync when changing the math.
    */
   private async normalizePdfPageRotations(fileBuffer: Buffer): Promise<Buffer> {
-    const srcDoc = await PDFDocument.load(fileBuffer);
+    let srcDoc: PDFDocument;
+    try {
+      srcDoc = await PDFDocument.load(fileBuffer, { ignoreEncryption: true });
+    } catch {
+      throw new BadRequestException(
+        "The file is not a valid PDF or is corrupted.",
+      );
+    }
     const pageCount = srcDoc.getPageCount();
 
     const hasRotation = srcDoc