fix(idempotency): bind operations to their own key prefix on a shared store

packages/idempotency/src/IdempotencyHandler.ts

@@ -52,9 +52,9 @@ export class IdempotencyHandler<Func extends AnyFunction> {
    */
   readonly #idempotencyConfig: IdempotencyConfig;
   /**
-   * Custom prefix to be used when generating the idempotency key.
+   * Key prefix resolved by the persistence store for this operation.
    */
-  readonly #keyPrefix: string | undefined;
+  readonly #resolvedKeyPrefix: string;
   /**
    * Persistence layer used to store the idempotency records.
    */
@@ -79,16 +79,16 @@ export class IdempotencyHandler<Func extends AnyFunction> {
     this.#functionToMakeIdempotent = functionToMakeIdempotent;
     this.#functionPayloadToBeHashed = functionPayloadToBeHashed;
     this.#idempotencyConfig = idempotencyConfig;
-    this.#keyPrefix = keyPrefix;
     this.#functionArguments = functionArguments;
     this.#thisArg = thisArg;
 
     this.#persistenceStore = persistenceStore;
 
     this.#persistenceStore.configure({
       config: this.#idempotencyConfig,
-      keyPrefix: this.#keyPrefix,
+      keyPrefix,
     });
+    this.#resolvedKeyPrefix = this.#persistenceStore.idempotencyKeyPrefix;
   }
 
   /**
@@ -336,13 +336,24 @@ export class IdempotencyHandler<Func extends AnyFunction> {
     return false;
   }
 
+  /**
+   * Restore this operation's key prefix on the persistence store.
+   *
+   * Operations sharing a store reconfigure its prefix on construction, so it
+   * is re-applied immediately before each call that hashes a key.
+   */
+  readonly #applyKeyPrefix = (): void => {
+    this.#persistenceStore.idempotencyKeyPrefix = this.#resolvedKeyPrefix;
+  };
+
   /**
    * Delete an in progress record from the idempotency store.
    *
    * This is called when the handler throws an error.
    */
   readonly #deleteInProgressRecord = async (): Promise<void> => {
     try {
+      this.#applyKeyPrefix();
       await this.#persistenceStore.deleteRecord(
         this.#functionPayloadToBeHashed
       );
@@ -376,6 +387,7 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       result: undefined,
     };
     try {
+      this.#applyKeyPrefix();
       await this.#persistenceStore.saveInProgress(
         this.#functionPayloadToBeHashed,
         this.#idempotencyConfig.lambdaContext?.getRemainingTimeInMillis()
@@ -399,6 +411,7 @@ export class IdempotencyHandler<Func extends AnyFunction> {
           // If the error doesn't include the existing record, we need to fetch
           // it from the persistence layer. In doing so, we also call the processExistingRecord
           // method to validate the record and cache it in memory.
+          this.#applyKeyPrefix();
           idempotencyRecord = await this.#persistenceStore.getRecord(
             this.#functionPayloadToBeHashed
           );
@@ -435,6 +448,7 @@ export class IdempotencyHandler<Func extends AnyFunction> {
     result: ReturnType<Func>
   ): Promise<void> => {
     try {
+      this.#applyKeyPrefix();
       await this.#persistenceStore.saveSuccess(
         this.#functionPayloadToBeHashed,
         result

packages/idempotency/src/persistence/BasePersistenceLayer.ts

@@ -66,6 +66,8 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
       this.idempotencyKeyPrefix = keyPrefix.trim();
     } else if (functionName?.trim()) {
       this.idempotencyKeyPrefix = `${this.#keyPrefixBase}.${functionName.trim()}`;
+    } else {
+      this.idempotencyKeyPrefix = this.#keyPrefixBase;
     }
 
     // Prevent reconfiguration

packages/idempotency/tests/unit/makeIdempotent.test.ts

@@ -672,6 +672,72 @@ describe('Function: makeIdempotent', () => {
     expect(saveSuccessSpy).toHaveBeenCalledWith(event, '123456');
   });
 
+  it('completes nested operations under their own key prefix when sharing a persistence store', async () => {
+    // Prepare
+    const persistenceStore = new PersistenceLayerTestClass();
+    const config = new IdempotencyConfig({});
+    config.registerLambdaContext(context);
+    const inner = makeIdempotent(async (_event: unknown) => 'inner', {
+      persistenceStore,
+      config,
+      keyPrefix: 'inner',
+    });
+    const outer = makeIdempotent(
+      async (event: unknown) => `outer:${await inner(event)}`,
+      { persistenceStore, config, keyPrefix: 'outer' }
+    );
+    const event = { id: 'order-1' };
+
+    // Act
+    const result = await outer(event);
+
+    // Assess
+    expect(result).toBe('outer:inner');
+    const putKeys = persistenceStore._putRecord.mock.calls.map(
+      ([record]) => record.idempotencyKey
+    );
+    const updateKeys = persistenceStore._updateRecord.mock.calls.map(
+      ([record]) => record.idempotencyKey
+    );
+    expect(putKeys).toEqual([
+      expect.stringMatching(/^outer#/),
+      expect.stringMatching(/^inner#/),
+    ]);
+    expect(updateKeys).toEqual([putKeys[1], putKeys[0]]);
+  });
+
+  it('completes a nested operation without a key prefix under the default prefix', async () => {
+    // Prepare
+    const persistenceStore = new PersistenceLayerTestClass();
+    const config = new IdempotencyConfig({});
+    config.registerLambdaContext(context);
+    const inner = makeIdempotent(async (_event: unknown) => 'inner', {
+      persistenceStore,
+      config,
+    });
+    const outer = makeIdempotent(
+      async (event: unknown) => `outer:${await inner(event)}`,
+      { persistenceStore, config, keyPrefix: 'outer' }
+    );
+
+    // Act
+    const result = await outer({ id: 'order-1' });
+
+    // Assess
+    expect(result).toBe('outer:inner');
+    const putKeys = persistenceStore._putRecord.mock.calls.map(
+      ([record]) => record.idempotencyKey
+    );
+    const updateKeys = persistenceStore._updateRecord.mock.calls.map(
+      ([record]) => record.idempotencyKey
+    );
+    expect(putKeys).toEqual([
+      expect.stringMatching(/^outer#/),
+      expect.stringMatching(/^my-lambda-function#/),
+    ]);
+    expect(updateKeys).toEqual([putKeys[1], putKeys[0]]);
+  });
+
   it('uses the specified argument as payload when wrapping an arbitrary function', async () => {
     // Prepare
     const config = new IdempotencyConfig({});

packages/idempotency/tests/unit/persistence/BasePersistenceLayer.test.ts

@@ -99,6 +99,19 @@ describe('Class: BasePersistenceLayer', () => {
       );
     });
 
+    it('resets the idempotency key prefix when configured without a prefix or function name', () => {
+      // Prepare
+      const config = new IdempotencyConfig({});
+      const persistenceLayer = new PersistenceLayerTestClass();
+      persistenceLayer.configure({ config, keyPrefix: 'custom' });
+
+      // Act
+      persistenceLayer.configure({ config });
+
+      // Assess
+      expect(persistenceLayer.idempotencyKeyPrefix).toBe('my-lambda-function');
+    });
+
     it('trims the function name before appending as key prefix', () => {
       // Prepare
       const config = new IdempotencyConfig({});
@@ -573,6 +586,26 @@ describe('Class: BasePersistenceLayer', () => {
       );
     });
 
+    it('hashes the idempotency key before yielding to the event loop', async () => {
+      // Prepare
+      const persistenceLayer = new PersistenceLayerTestClass();
+      persistenceLayer.configure({
+        config: new IdempotencyConfig({}),
+        keyPrefix: 'first',
+      });
+      const putRecordSpy = vi.spyOn(persistenceLayer, '_putRecord');
+
+      // Act
+      const pending = persistenceLayer.saveInProgress({ foo: 'bar' }, 2000);
+      persistenceLayer.idempotencyKeyPrefix = 'second';
+      await pending;
+
+      // Assess
+      expect(putRecordSpy).toHaveBeenCalledWith(
+        expect.objectContaining({ idempotencyKey: 'first#mocked-hash' })
+      );
+    });
+
     it('logs a warning when unable to call remainingTimeInMillis() from the context', async () => {
       // Prepare
       const persistenceLayer = new PersistenceLayerTestClass();