Add retry logic to S3 metadata

apps/backend/src/services/storage.test.ts

@@ -1,5 +1,5 @@
-import { describe, expect, it } from 'vitest';
-import { buildRuleKey, buildVersionedSubmissionKey } from './storage.js';
+import { describe, expect, it, vi } from 'vitest';
+import { buildRuleKey, buildVersionedSubmissionKey, getObjectMetadata } from './storage.js';
 
 describe('storage key builders', () => {
   it('builds versioned submission key with v prefix', () => {
@@ -20,4 +20,67 @@ describe('storage key builders', () => {
     expect(key).toContain('rules/ed-1/');
     expect(key.endsWith('_rule_doc.pdf')).toBe(true);
   });
+
+  it('retries metadata lookup until the object becomes visible', async () => {
+    const fetchMetadata = vi
+      .fn<
+        (
+          bucket: string,
+          key: string,
+        ) => Promise<{ contentLength: number | null; contentType: string | null }>
+      >()
+      .mockRejectedValueOnce(
+        Object.assign(new Error('Not found yet'), {
+          name: 'NotFound',
+          $metadata: { httpStatusCode: 404 },
+        }),
+      )
+      .mockResolvedValueOnce({
+        contentLength: 1024,
+        contentType: null,
+      })
+      .mockResolvedValueOnce({
+        contentLength: 1024,
+        contentType: 'application/pdf',
+      });
+    const sleep = vi.fn(async (_ms: number) => {});
+
+    const metadata = await getObjectMetadata('bucket', 'key', {
+      maxAttempts: 3,
+      retryDelayMs: 10,
+      fetchMetadata,
+      sleep,
+    });
+
+    expect(metadata).toEqual({
+      contentLength: 1024,
+      contentType: 'application/pdf',
+    });
+    expect(fetchMetadata).toHaveBeenCalledTimes(3);
+    expect(sleep).toHaveBeenCalledTimes(2);
+  });
+
+  it('does not retry non-retryable metadata lookup errors', async () => {
+    const fetchMetadata = vi
+      .fn<
+        (
+          bucket: string,
+          key: string,
+        ) => Promise<{ contentLength: number | null; contentType: string | null }>
+      >()
+      .mockRejectedValueOnce(new Error('Access denied'));
+    const sleep = vi.fn(async (_ms: number) => {});
+
+    await expect(
+      getObjectMetadata('bucket', 'key', {
+        maxAttempts: 3,
+        retryDelayMs: 10,
+        fetchMetadata,
+        sleep,
+      }),
+    ).rejects.toThrow('Access denied');
+
+    expect(fetchMetadata).toHaveBeenCalledTimes(1);
+    expect(sleep).not.toHaveBeenCalled();
+  });
 });

apps/backend/src/services/storage.ts

@@ -2,6 +2,7 @@ import { randomUUID } from 'node:crypto';
 import {
   GetObjectCommand,
   HeadObjectCommand,
+  type HeadObjectCommandOutput,
   PutObjectCommand,
   S3Client,
 } from '@aws-sdk/client-s3';
@@ -35,6 +36,18 @@ type PresignUploadByKeyInput = {
   expiresIn?: number;
 };
 
+type ObjectMetadata = {
+  contentLength: number | null;
+  contentType: string | null;
+};
+
+type GetObjectMetadataOptions = {
+  maxAttempts?: number;
+  retryDelayMs?: number;
+  sleep?: (ms: number) => Promise<void>;
+  fetchMetadata?: (bucket: string, key: string) => Promise<ObjectMetadata>;
+};
+
 const encodeName = (name: string): string => {
   return name.replace(/[^a-zA-Z0-9._-]/g, '_');
 };
@@ -129,19 +142,74 @@ export const presignDownload = async (
   return { presignedUrl, expiresIn };
 };
 
-export const getObjectMetadata = async (
-  bucket: string,
-  key: string,
-): Promise<{ contentLength: number | null; contentType: string | null }> => {
-  const response = await s3.send(
+const sleep = async (ms: number): Promise<void> => {
+  await new Promise((resolve) => setTimeout(resolve, ms));
+};
+
+const headObject = async (bucket: string, key: string): Promise<HeadObjectCommandOutput> => {
+  return await s3.send(
     new HeadObjectCommand({
       Bucket: bucket,
       Key: key,
     }),
   );
+};
 
+const toObjectMetadata = (response: HeadObjectCommandOutput): ObjectMetadata => {
   return {
     contentLength: response.ContentLength ?? null,
     contentType: response.ContentType ?? null,
   };
 };
+
+const hasCompleteObjectMetadata = (metadata: ObjectMetadata): boolean => {
+  return metadata.contentLength !== null && metadata.contentType !== null;
+};
+
+const isRetryableMetadataLookupError = (error: unknown): boolean => {
+  if (!(error instanceof Error)) {
+    return false;
+  }
+
+  const metadata = (error as Error & { $metadata?: { httpStatusCode?: number } }).$metadata;
+  const statusCode = metadata?.httpStatusCode;
+  return statusCode === 404 || error.name === 'NotFound' || error.name === 'NoSuchKey';
+};
+
+export const getObjectMetadata = async (
+  bucket: string,
+  key: string,
+  options: GetObjectMetadataOptions = {},
+): Promise<ObjectMetadata> => {
+  const maxAttempts = options.maxAttempts ?? 5;
+  const retryDelayMs = options.retryDelayMs ?? 1000;
+  const wait = options.sleep ?? sleep;
+  const fetchMetadata =
+    options.fetchMetadata ??
+    (async (innerBucket: string, innerKey: string): Promise<ObjectMetadata> => {
+      return toObjectMetadata(await headObject(innerBucket, innerKey));
+    });
+
+  let lastIncompleteMetadata: ObjectMetadata | null = null;
+
+  for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+    try {
+      const metadata = await fetchMetadata(bucket, key);
+      if (hasCompleteObjectMetadata(metadata)) {
+        return metadata;
+      }
+
+      lastIncompleteMetadata = metadata;
+    } catch (error) {
+      if (!isRetryableMetadataLookupError(error) || attempt === maxAttempts) {
+        throw error;
+      }
+    }
+
+    if (attempt < maxAttempts) {
+      await wait(retryDelayMs);
+    }
+  }
+
+  return lastIncompleteMetadata ?? { contentLength: null, contentType: null };
+};