fix(memoize): correctly throw error on memoize when the original method throws

Author: tabkram

src/execution/memoize.ts

@@ -59,15 +59,23 @@ export function executeMemoize<O>(
   if (memoizedValue) {
     return memoizedValue;
   } else {
-    const callResponseOrPromise = (execute.bind(this) as typeof execute)(blockFunction.bind(this) as typeof blockFunction, inputs);
+    const callResponseOrPromise = (execute.bind(this) as typeof execute)(
+      blockFunction.bind(this) as typeof blockFunction,
+      inputs,
+      [],
+      (res) => res,
+      (error) => {
+        throw error;
+      }
+    );
     memoizationStore.set(inputsHash, callResponseOrPromise);
     this[memoizationKey].set(options.functionId, memoizationStore);
 
     if (callResponseOrPromise instanceof Promise) {
-      callResponseOrPromise.finally(() => setTimeout(() => memoizationStore.delete(inputsHash), expirationMs));
+      return callResponseOrPromise.finally(() => setTimeout(() => memoizationStore.delete(inputsHash), expirationMs));
     } else {
       setTimeout(() => memoizationStore.delete(inputsHash), expirationMs);
+      return callResponseOrPromise;
     }
-    return callResponseOrPromise;
   }
 }

src/execution/memoizeDecorator.spec.ts

@@ -70,4 +70,78 @@ describe('memoize decorator', () => {
     expect(memoizationCheckCount).toBe(1); // 1memoized
     expect(fib5).toBe(8);
   });
+
+  it('should memoize async function results and prevent redundant calls', async () => {
+    let memoizationCheckCount = 0;
+    let memoizedCalls = 0;
+    let totalFunctionCalls = 0;
+
+    class DataService {
+      @memoize(async (memoContext: MemoizationContext<string>) => {
+        memoizationCheckCount++;
+        if (memoContext.isMemoized) {
+          memoizedCalls++;
+        }
+      })
+      async fetchData(id: number): Promise<string> {
+        totalFunctionCalls++;
+        return new Promise((resolve) =>
+          setTimeout(() => resolve(`Data for ID: ${id}`), 100)
+        );
+      }
+
+      @memoize(async (memoContext: MemoizationContext<string>) => {
+        memoizationCheckCount++;
+        if (memoContext.isMemoized) {
+          memoizedCalls++;
+        }
+      }) async throwData(name: string): Promise<string> {
+        totalFunctionCalls++;
+        throw new Error(`hello ${name} but I throw!`);
+      }
+    }
+
+    const service = new DataService();
+
+    memoizationCheckCount = 0;
+    memoizedCalls = 0;
+    totalFunctionCalls = 0;
+
+    const result1 = await service.fetchData(1);
+    expect(result1).toBe('Data for ID: 1');
+    expect(memoizedCalls).toBe(0);
+    expect(totalFunctionCalls).toBe(1);
+    expect(memoizationCheckCount).toBe(1); // Called once
+
+    const result2 = await service.fetchData(1);
+    expect(result2).toBe('Data for ID: 1');
+    expect(memoizedCalls).toBe(1); // Now it should be memoized
+    expect(totalFunctionCalls).toBe(1); // No new calls
+    expect(memoizationCheckCount).toBe(2); // Checked twice
+
+    const result3 = await service.fetchData(2);
+    expect(result3).toBe('Data for ID: 2');
+    expect(memoizedCalls).toBe(1); // No extra memoized calls yet
+    expect(totalFunctionCalls).toBe(2); // New call for different ID
+    expect(memoizationCheckCount).toBe(3); // Three checks (1st, 2nd for ID 1, and 3rd for ID 2)
+
+    const result4 = await service.fetchData(2);
+    expect(result4).toBe('Data for ID: 2');
+    expect(memoizedCalls).toBe(2); // ID 2 result is now memoized
+    expect(totalFunctionCalls).toBe(2); // No extra new calls
+    expect(memoizationCheckCount).toBe(4); // 4 checks in total
+
+    // test memoize a throwing async method
+    memoizationCheckCount = 0;
+    memoizedCalls = 0;
+    totalFunctionCalls = 0;
+    await Promise.all([
+      expect(service.throwData('akram')).rejects.toThrow('hello akram but I throw!'),
+      expect(service.throwData('akram')).rejects.toThrow('hello akram but I throw!'),
+      expect(service.throwData('akram')).rejects.toThrow('hello akram but I throw!')
+    ]);
+    expect(memoizationCheckCount).toEqual(totalFunctionCalls + memoizedCalls);
+    expect(memoizedCalls).toEqual(2);
+    expect(totalFunctionCalls).toBe(1); // No extra new calls
+  });
 });