Add upload deduping to avoid race. Uploads are idempotent but races could still occur and waste time/bandwidth.

Author: cconstable

contrib/external-storage-s3/src/__tests__/test-driver.ts

@@ -136,6 +136,18 @@ test('identical payloads in the same scope deduplicate to one upload', async (t)
   t.is(client.putCount, 1);
 });
 
+test('concurrent identical payloads in one batch upload once', async (t) => {
+  const client = new FakeS3Client();
+  const driver = new S3StorageDriver({ client, bucket: 'b' });
+
+  const [first, second] = await driver.store(workflowContext, [makePayload('"hello"'), makePayload('"hello"')]);
+  assert(first);
+  assert(second);
+
+  t.is(client.putCount, 1);
+  t.is(first.claimData.key, second.claimData.key);
+});
+
 test('retrieve rejects when stored bytes fail the integrity check', async (t) => {
   const client = new FakeS3Client();
   const driver = new S3StorageDriver({ client, bucket: 'b' });

contrib/external-storage-s3/src/driver.ts

@@ -142,9 +142,10 @@ export class S3StorageDriver implements StorageDriver {
 
   async store(context: StorageDriverStoreContext, payloads: Payload[]): Promise<StorageDriverClaim[]> {
     const contextSegments = buildContextSegments(context.target);
-    return runAllWithAbortOnError(context.abortSignal, (signal) =>
-      payloads.map((payload) => this.storePayload(context, payload, contextSegments, signal))
-    );
+    return runAllWithAbortOnError(context.abortSignal, (signal) => {
+      const uploads = new Map<string, Promise<void>>();
+      return payloads.map((payload) => this.storePayload(context, payload, contextSegments, signal, uploads));
+    });
   }
 
   async retrieve(context: StorageDriverRetrieveContext, claims: StorageDriverClaim[]): Promise<Payload[]> {
@@ -157,7 +158,8 @@ export class S3StorageDriver implements StorageDriver {
     context: StorageDriverStoreContext,
     payload: Payload,
     contextSegments: string,
-    abortSignal: AbortSignal
+    abortSignal: AbortSignal,
+    uploads: Map<string, Promise<void>>
   ): Promise<StorageDriverClaim> {
     const bucket = this.bucket(context, payload);
 
@@ -172,9 +174,13 @@ export class S3StorageDriver implements StorageDriver {
     const key = `v0${contextSegments}/d/sha256/${hashValue}`;
 
     try {
-      if (!(await this.client.objectExists(bucket, key, { abortSignal }))) {
-        await this.client.putObject(bucket, key, payloadBytes, { abortSignal });
+      const dedupeKey = `${bucket} ${key}`;
+      let upload = uploads.get(dedupeKey);
+      if (!upload) {
+        upload = this.uploadIfAbsent(bucket, key, payloadBytes, abortSignal);
+        uploads.set(dedupeKey, upload);
       }
+      await upload;
     } catch (err) {
       throw new Error(
         `S3StorageDriver store failed [bucket=${bucket}, key=${key}${formatClientContext(this.client)}]`,
@@ -185,6 +191,17 @@ export class S3StorageDriver implements StorageDriver {
     return new StorageDriverClaim({ bucket, key, hashAlgorithm: 'sha256', hashValue });
   }
 
+  private async uploadIfAbsent(
+    bucket: string,
+    key: string,
+    payloadBytes: Uint8Array,
+    abortSignal: AbortSignal
+  ): Promise<void> {
+    if (!(await this.client.objectExists(bucket, key, { abortSignal }))) {
+      await this.client.putObject(bucket, key, payloadBytes, { abortSignal });
+    }
+  }
+
   private async retrievePayload(claim: StorageDriverClaim, abortSignal: AbortSignal): Promise<Payload> {
     const { bucket, key, hashAlgorithm, hashValue: expectedHash } = claim.claimData;
     if (!bucket || !key) {