diff --git a/apps/backend-services/src/document/document.module.ts b/apps/backend-services/src/document/document.module.ts index 0fc8c9449..c9f41b801 100644 --- a/apps/backend-services/src/document/document.module.ts +++ b/apps/backend-services/src/document/document.module.ts @@ -1,6 +1,7 @@ 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"; @@ -8,7 +9,12 @@ import { PdfNormalizationService } from "./pdf-normalization.service"; @Module({ imports: [BlobStorageModule, TemporalModule], - providers: [DocumentDbService, DocumentService, PdfNormalizationService], + providers: [ + DocumentDbService, + DocumentService, + PdfNormalizationService, + UploadNormalizationLimiter, + ], controllers: [DocumentController], exports: [DocumentService, PdfNormalizationService], }) diff --git a/apps/backend-services/src/document/document.service.spec.ts b/apps/backend-services/src/document/document.service.spec.ts index 449f8b629..c94b1e94e 100644 --- a/apps/backend-services/src/document/document.service.spec.ts +++ b/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 + Pick< + PdfNormalizationService, + "validateForUpload" | "normalizeToPdf" | "generateThumbnailWebp" + > + >; + let uploadNormalizationLimiter: jest.Mocked< + Pick >; 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) => task()), + } as jest.Mocked>; 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); @@ -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((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((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 () => { diff --git a/apps/backend-services/src/document/document.service.ts b/apps/backend-services/src/document/document.service.ts index 02cf1ec61..19687a03e 100644 --- a/apps/backend-services/src/document/document.service.ts +++ b/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; } diff --git a/apps/backend-services/src/document/pdf-normalization.service.spec.ts b/apps/backend-services/src/document/pdf-normalization.service.spec.ts index 02dcecf9c..6fc5a3e7b 100644 --- a/apps/backend-services/src/document/pdf-normalization.service.spec.ts +++ b/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"); diff --git a/apps/backend-services/src/document/pdf-normalization.service.ts b/apps/backend-services/src/document/pdf-normalization.service.ts index 027f26f12..db4e4e3f4 100644 --- a/apps/backend-services/src/document/pdf-normalization.service.ts +++ b/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 { @@ -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 { - 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 diff --git a/apps/backend-services/src/upload/upload-normalization-limiter.spec.ts b/apps/backend-services/src/upload/upload-normalization-limiter.spec.ts new file mode 100644 index 000000000..cf118392b --- /dev/null +++ b/apps/backend-services/src/upload/upload-normalization-limiter.spec.ts @@ -0,0 +1,55 @@ +jest.mock("node:os", () => ({ + availableParallelism: () => 2, + cpus: () => [{}, {}, {}, {}], +})); + +import { + getUploadNormalizationConcurrency, + UploadNormalizationLimiter, +} from "./upload-normalization-limiter"; + +describe("getUploadNormalizationConcurrency", () => { + it("returns Math.max(2, availableParallelism)", () => { + expect(getUploadNormalizationConcurrency()).toBe(2); + }); +}); + +describe("UploadNormalizationLimiter", () => { + it("executes the wrapped task and returns its result", async () => { + const limiter = new UploadNormalizationLimiter(); + await expect(limiter.run(async () => "ok")).resolves.toBe("ok"); + }); + + it("queues tasks when the concurrency limit is reached", async () => { + const limiter = new UploadNormalizationLimiter(); + let active = 0; + let maxActive = 0; + const unblock: Array<() => void> = []; + + const task = () => + limiter.run(async () => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => { + unblock.push(resolve); + }); + active -= 1; + }); + + const pending = [task(), task(), task()]; + await new Promise((resolve) => setImmediate(resolve)); + expect(maxActive).toBe(2); + expect(unblock).toHaveLength(2); + + unblock.shift()?.(); + await new Promise((resolve) => setImmediate(resolve)); + expect(unblock).toHaveLength(2); + + unblock.shift()?.(); + unblock.shift()?.(); + unblock.shift()?.(); + + await Promise.all(pending); + expect(maxActive).toBe(2); + }); +}); diff --git a/apps/backend-services/src/upload/upload-normalization-limiter.ts b/apps/backend-services/src/upload/upload-normalization-limiter.ts new file mode 100644 index 000000000..0c5921f1c --- /dev/null +++ b/apps/backend-services/src/upload/upload-normalization-limiter.ts @@ -0,0 +1,85 @@ +import { availableParallelism, cpus } from "node:os"; +import { Injectable } from "@nestjs/common"; + +/** + * In-process counting semaphore for async work. + * + * Tasks beyond the configured limit wait in FIFO order until a slot frees. + * Used by {@link UploadNormalizationLimiter} to cap concurrent PDF/image + * normalization without a cross-request queue or external coordinator. + */ +class Semaphore { + private active = 0; + private readonly waiting: Array<() => void> = []; + + /** @param limit Maximum number of tasks that may run at once (must be ≥ 1). */ + constructor(private readonly limit: number) {} + + /** + * Acquire a slot, run `task`, then release the slot (even when `task` throws). + */ + async run(task: () => Promise): Promise { + await this.acquire(); + try { + return await task(); + } finally { + this.release(); + } + } + + /** Reserve a slot immediately or enqueue until one is available. */ + private acquire(): Promise { + if (this.active < this.limit) { + this.active += 1; + return Promise.resolve(); + } + + return new Promise((resolve) => { + this.waiting.push(() => { + this.active += 1; + resolve(); + }); + }); + } + + /** Free a slot and wake the next queued waiter, if any. */ + private release(): void { + this.active -= 1; + const next = this.waiting.shift(); + if (next) { + next(); + } + } +} + +/** + * Concurrency cap for upload normalization on this process. + * + * Uses `availableParallelism()` when present (Node 19+), otherwise `cpus().length`, + * floored at 2 so small containers still allow limited overlap. + */ +export function getUploadNormalizationConcurrency(): number { + const detectedParallelism = + typeof availableParallelism === "function" + ? availableParallelism() + : cpus().length; + return Math.max(2, detectedParallelism); +} + +/** + * Nest injectable that bounds concurrent `normalizeToPdf` work per backend pod. + * + * JSON/base64 uploads still decode fully before normalization; this limiter only + * prevents many large pdf-lib workspaces from running at once on one process. + */ +@Injectable() +export class UploadNormalizationLimiter { + private readonly semaphore = new Semaphore( + getUploadNormalizationConcurrency(), + ); + + /** Run `task` when a normalization slot is available. */ + run(task: () => Promise): Promise { + return this.semaphore.run(task); + } +} diff --git a/docs-md/DOCUMENT_IMAGE_NORMALIZATION.md b/docs-md/DOCUMENT_IMAGE_NORMALIZATION.md index 86b8e11c3..9c6588adc 100644 --- a/docs-md/DOCUMENT_IMAGE_NORMALIZATION.md +++ b/docs-md/DOCUMENT_IMAGE_NORMALIZATION.md @@ -20,11 +20,11 @@ Database: `file_path` → original blob; `normalized_file_path` → normalized P | `GET .../labeling/projects/:id/documents/:docId/view` | Normalized PDF for labeling. | | `GET .../labeling/projects/:id/documents/:docId/download` | Original labeling upload. | -Upload may return **422** with `code: conversion_failed` when the original was stored but PDF normalization failed; status `conversion_failed` is set and OCR is not started. +Upload may return **422** with `code: conversion_failed` when the original was stored but PDF normalization failed; status `conversion_failed` is set and OCR is not started. The original blob is **not** deleted on normalization failure: `file_path` remains valid so clients can still download the upload; only `normalized_file_path` and downstream OCR/view are absent until a future retry path exists. Labeling project upload (`POST .../labeling/projects/:id/upload`) requires `group_id` in the body to **match** the project’s group; the caller must also be allowed to access that group (same as other labeling routes). -**Invalid or unsupported files** are rejected with **400** (corrupt image, bad PDF signature, etc.) before a successful store, where validation applies. +**Invalid or unsupported files** are rejected with **400** when validation fails before storage (bad PDF signature, corrupt image, etc.). PDFs with a valid `%PDF` header but an unreadable body fail during normalization (400) after the original blob write has started. ## Client @@ -34,7 +34,20 @@ Labeling project upload (`POST .../labeling/projects/:id/upload`) requires `grou ## Dependencies -Backend normalization uses **sharp** (images, including multi-page TIFF) and **pdf-lib** (PDF assembly). PDF uploads must pass a **magic-byte check and a full `pdf-lib` parse** (with `ignoreEncryption: true`) before storage; images are validated with **sharp** metadata. +Backend normalization uses **sharp** (images, including multi-page TIFF) and **pdf-lib** (PDF assembly). PDF uploads must pass a **magic-byte check** (`%PDF`) before the original blob is written; full **`pdf-lib` parsing** happens once during normalization (with `ignoreEncryption: true`). Images are validated with **sharp** metadata before storage. + +## Upload memory and concurrency (`POST /api/upload`) + +The upload API is JSON/base64, so Nest must parse the full request body and the service decodes it to a `Buffer` before normalization. Within that constraint, the upload path minimizes peak memory: + +1. **Cheap pre-write validation** — PDF/scan: `%PDF` header only; corrupt bodies are rejected when `normalizeToPdf` loads the document (400). Images: sharp metadata probe. +2. **Overlapped I/O and CPU** — original blob write starts while `normalizeToPdf` runs under a concurrency cap. +3. **Single pdf-lib load** — validation no longer loads the PDF separately from normalization. +4. **Early buffer release** — after the original write completes, the decoded upload buffer is dropped before the normalized blob write. +5. **Thumbnail from normalized PDF** — `generateThumbnailWebp` uses the normalized PDF bytes, not the original upload buffer. +6. **Normalization limiter** — [`UploadNormalizationLimiter`](../apps/backend-services/src/upload/upload-normalization-limiter.ts) bounds concurrent `normalizeToPdf` calls to `Math.max(2, availableParallelism())`. + +True request-stream normalization would require changing `/api/upload` to multipart or a raw-body contract; see [LOAD_TESTING.md](./LOAD_TESTING.md). ## Two-stage orientation correction diff --git a/docs-md/LOAD_TESTING.md b/docs-md/LOAD_TESTING.md index 30dee829e..3be6969f4 100644 --- a/docs-md/LOAD_TESTING.md +++ b/docs-md/LOAD_TESTING.md @@ -216,7 +216,9 @@ Payload configuration: No proprietary fixtures are checked in or required. The upload path generates a syntactically valid synthetic PDF with generic text and padding; the blob path generates deterministic binary text. If operators provide fixture paths, those files must be generated or openly licensed, and any committed fixture must include clear license documentation. -`POST /api/upload` sends JSON, so raw PDF bytes expand by roughly 4/3 as base64 before Nest applies `BODY_LIMIT`. k6 estimates the JSON body size and aborts setup if the selected tier would exceed the configured limit. The backend then validates the PDF, stores the original blob, normalizes to PDF, creates the document row, and queues OCR/workflow processing. Keep backend-services on `DOCUMENT_INTELLIGENCE_MODE=mock` and Temporal worker on `MOCK_AZURE_OCR=true` for these runs so live Azure Document Intelligence latency or quota errors do not mask platform body-limit, normalization, storage, or queue behavior. +`POST /api/upload` sends JSON, so raw PDF bytes expand by roughly 4/3 as base64 before Nest applies `BODY_LIMIT`. k6 estimates the JSON body size and aborts setup if the selected tier would exceed the configured limit. The backend validates uploads with a cheap header/metadata probe, starts the original blob write, normalizes to PDF under an in-process cap of `Math.max(2, available CPU parallelism)`, awaits the original write, drops the decoded upload buffer, writes the normalized PDF, creates the document row, and queues OCR/workflow processing. Thumbnails are generated from the normalized PDF so the original bytes are not retained for the full request. Keep backend-services on `DOCUMENT_INTELLIGENCE_MODE=mock` and Temporal worker on `MOCK_AZURE_OCR=true` for these runs so live Azure Document Intelligence latency or quota errors do not mask platform body-limit, normalization, storage, or queue behavior. + +The current upload API contract is JSON/base64, not multipart. That means the decoded request bytes must still exist at least long enough for validation, original storage, and `pdf-lib` normalization; true request-stream normalization would require changing `/api/upload` to accept multipart or adding a replacement upload contract. ## Review / HITL API scenario