Fix memoize with Promise rejections

README.md

@@ -260,7 +260,7 @@ memoizedFunc(1, 2)
   with some improvements for variadic performance and additional support for a TTL based cache.
 </details>
 
-<sup>[Source](./src/memoize/memoize.ts) • [Benchmark](./src/memoize/BENCHMARK.md) • Minify: 820 B • Minify & GZIP: 393 B<sup>
+<sup>[Source](./src/memoize/memoize.ts) • [Benchmark](./src/memoize/BENCHMARK.md) • Minify: 962 B • Minify & GZIP: 441 B<sup>
 
 ### min(array)
 

src/memoize/BENCHMARK.md

@@ -2,10 +2,10 @@
 
 [Source for this benchmark](./benchmark.ts)
 
-|                      | lodash                      | fast-memoize                 | flocky                            |
-| -------------------- | --------------------------- | ---------------------------- | --------------------------------- |
-| monadic (primitive)  | 70,899,764 ops/sec (34.95%) | 201,953,383 ops/sec (99.54%) | **202,886,184 ops/sec (100.00%)** |
-| monadic (serialized) | 2,826,332 ops/sec (37.76%)  | 2,272,833 ops/sec (30.36%)   | **7,485,491 ops/sec (100.00%)**   |
-| variadic             | 2,893,648 ops/sec (71.46%)  | 1,414,302 ops/sec (34.93%)   | **4,049,054 ops/sec (100.00%)**   |
+|                      | lodash                      | fast-memoize                      | flocky                          |
+| -------------------- | --------------------------- | --------------------------------- | ------------------------------- |
+| monadic (primitive)  | 70,207,892 ops/sec (34.59%) | **202,998,947 ops/sec (100.00%)** | 199,242,881 ops/sec (98.15%)    |
+| monadic (serialized) | 2,884,723 ops/sec (40.50%)  | 2,409,543 ops/sec (33.83%)        | **7,122,748 ops/sec (100.00%)** |
+| variadic             | 2,864,947 ops/sec (76.45%)  | 1,486,428 ops/sec (39.67%)        | **3,747,245 ops/sec (100.00%)** |
 
-<sup>Generated at 2021-07-05 with Node.JS v16.4.1</sup>
+<sup>Generated at 2023-11-27 with Node.JS v18.18.2</sup>

src/memoize/memoize.spec.ts

@@ -1,5 +1,5 @@
 import { sleep } from '../sleep/sleep'
-import { memoize } from './memoize'
+import { MemoizeOptions, memoize } from './memoize'
 
 type NumberObject = { n: number }
 
@@ -223,6 +223,10 @@ describe('memoize', () => {
     expect(await a).toEqual('A')
     expect(await b).toEqual('A')
     expect(calls).toEqual(1)
+
+    const c = memoizedFunc()
+    expect(await c).toEqual('A')
+    expect(calls).toEqual(1)
   })
 
   test('memoizes function calls with a maximum TTL', async () => {
@@ -254,6 +258,45 @@ describe('memoize', () => {
     ])
   })
 
+  test.each([
+    ['monadic', { strategy: 'monadic' }],
+    ['variadic', { strategy: 'variadic' }],
+    ['ttl', { ttl: 50 }],
+  ] as Array<[string, MemoizeOptions]>)(
+    'smartly memoizes function calls that return promise rejections (%s)',
+    async (_: string, options: MemoizeOptions) => {
+      let calls = 0
+      const func = async () => {
+        calls++
+        await sleep(10)
+
+        if (calls === 1) {
+          throw new Error('Uh oh')
+        }
+
+        return 'A'
+      }
+
+      const memoizedFunc = memoize(func, options)
+
+      // We want to memoize the second call because the Promise is still pending
+      const a = memoizedFunc()
+      expect(a).toBeInstanceOf(Promise)
+      const b = memoizedFunc()
+      expect(b).toBeInstanceOf(Promise)
+      expect(calls).toEqual(1)
+
+      await expect(a).rejects.toThrow('Uh oh')
+      await expect(b).rejects.toThrow('Uh oh')
+      expect(calls).toEqual(1)
+
+      // We don't want to memoize the third call because the Promise has rejected
+      const c = memoizedFunc()
+      expect(await c).toEqual('A')
+      expect(calls).toEqual(2)
+    }
+  )
+
   test('has the correct type', async () => {
     const func = (a: number, b: number): number => {
       return a + b

src/memoize/memoize.ts

@@ -1,6 +1,6 @@
 import { TAnyFunction } from '../typeHelpers'
 
-interface MemoizeOptions {
+export interface MemoizeOptions {
   strategy?: MemoizeStrategy
   serializer?: MemoizeSerializer
   ttl?: number
@@ -63,6 +63,11 @@ function monadic<TThis, TReturn, TFunc extends TAnyFunction<TReturn>>(
   let value = cache.get(cacheKey)
   if (typeof value === 'undefined') {
     value = func.call(this, arg)
+
+    if (value instanceof Promise) {
+      value.catch(() => cache.remove(cacheKey))
+    }
+
     cache.set(cacheKey, value)
   }
 
@@ -81,6 +86,11 @@ function variadic<TThis, TReturn, TFunc extends TAnyFunction<TReturn>>(
   let value = cache.get(cacheKey)
   if (typeof value === 'undefined') {
     value = func.apply(this, args)
+
+    if (value instanceof Promise) {
+      value.catch(() => cache.remove(cacheKey))
+    }
+
     cache.set(cacheKey, value)
   }
 
@@ -94,6 +104,7 @@ function defaultSerializer(data: unknown) {
 interface MemoizeCache<TReturn> {
   get: (key: string) => TReturn | undefined
   set: (key: string, value: TReturn) => void
+  remove: (key: string) => void
 }
 
 function defaultCache<TReturn>(): MemoizeCache<TReturn> {
@@ -104,6 +115,7 @@ function defaultCache<TReturn>(): MemoizeCache<TReturn> {
     set: (key, value) => {
       cache[key] = value
     },
+    remove: (key) => delete cache[key],
   }
 }
 
@@ -121,5 +133,6 @@ function ttlCache<TReturn>(ttl: number): MemoizeCache<TReturn> {
         delete cache[key]
       }, ttl)
     },
+    remove: (key) => delete cache[key],
   }
 }