allow clearing the dedupe cache (#138)

* allow clearing the dedupe cache

* add changeset

* expose clearCacheForCurrentRequest

* reword changeset

* avoid adding prop to deduped fn

Author: dferber90

.changeset/twelve-feet-glow.md

@@ -0,0 +1,5 @@
+---
+'flags': patch
+---
+
+- expose `clearDedupeCacheForCurrentRequest` to allow clearing the cache of a deduped function for the current request

packages/flags/src/next/dedupe.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it, Mock, vitest, vi } from 'vitest';
-import { dedupe } from './dedupe';
+import { clearDedupeCacheForCurrentRequest, dedupe } from './dedupe';
 
 const mocks = vi.hoisted(() => {
   return {
@@ -190,6 +190,60 @@ describe('dedupe', () => {
     });
   });
 
+  describe('clearCurrent', () => {
+    it('should allow dedupe to be cleared within same request', async () => {
+      const fn = vitest.fn();
+      const deduped = dedupe(fn);
+      const same = new Headers();
+      const headersMock = await getHeadersMock();
+      headersMock.mockReturnValue(same);
+
+      await deduped();
+      await deduped();
+      await clearDedupeCacheForCurrentRequest(deduped);
+      await deduped();
+      await deduped();
+
+      expect(fn).toHaveBeenCalledTimes(2);
+    });
+
+    it('should not affect the cache of a different request', async () => {
+      const fn = vitest.fn();
+      const deduped = dedupe(fn);
+      const same = new Headers();
+      const other = new Headers();
+      const headersMock = await getHeadersMock();
+
+      // fill both caches and call once each, interleaved
+      headersMock.mockReturnValue(same);
+      await deduped();
+
+      headersMock.mockReturnValue(other);
+      await deduped();
+
+      headersMock.mockReturnValue(same);
+      await deduped();
+
+      headersMock.mockReturnValue(other);
+      await deduped();
+
+      // check the fn was only called twice (once for each request)
+      expect(fn).toHaveBeenCalledTimes(2);
+
+      // now clear one cache but not the other, and call once for each request
+      headersMock.mockReturnValue(same);
+      await clearDedupeCacheForCurrentRequest(deduped);
+      await deduped();
+
+      headersMock.mockReturnValue(other);
+      await deduped();
+
+      // it should go from 2 → 3 as the one request gets a cache hit, but
+      // the other request gets a cache miss
+      expect(fn).toHaveBeenCalledTimes(3);
+    });
+  });
+
   describe('promises', () => {
     it('should dedupe even when the promise has not resolved yet', async () => {
       let resolvePromise: (value: unknown) => void = () => {

packages/flags/src/next/dedupe.ts

@@ -4,6 +4,8 @@ enum Status {
   ERRORED = 2,
 }
 
+type RequestStore<T> = WeakMap<Headers, CacheNode<T>>;
+
 type CacheNode<T> = {
   s: Status;
   v: T | undefined | unknown;
@@ -23,6 +25,15 @@ function createCacheNode<T>(): CacheNode<T> {
   };
 }
 
+/**
+ * We use a registry to store the request store for each deduped function.
+ *
+ * This is necessary as we don't want to attach the request store to the deduped
+ * to retain its original type, and we need to be able to clear the cache for the
+ * current request.
+ */
+const cacheRegistry = new WeakMap<Function, RequestStore<unknown>>();
+
 /**
  * A middleware-friendly version of React.cache.
  *
@@ -40,9 +51,9 @@ function createCacheNode<T>(): CacheNode<T> {
 export function dedupe<A extends Array<unknown>, T>(
   fn: (...args: A) => T | Promise<T>,
 ): (...args: A) => Promise<T> {
-  const requestStore = new WeakMap<Headers, CacheNode<T>>();
+  const requestStore: RequestStore<T> = new WeakMap<Headers, CacheNode<T>>();
 
-  return async function (this: unknown, ...args: A): Promise<T> {
+  const dedupedFn = async function (this: unknown, ...args: A): Promise<T> {
     // async import required as turbopack errors in Pages Router
     // when next/headers is imported at the top-level
     const { headers } = await import('next/headers');
@@ -105,4 +116,29 @@ export function dedupe<A extends Array<unknown>, T>(
       throw error;
     }
   };
+
+  cacheRegistry.set(dedupedFn, requestStore);
+  return dedupedFn;
+}
+
+/**
+ * Clears the cached value of a deduped function for the current request.
+ *
+ * This function is useful for resetting the cache after making changes to
+ * the underlying cached information.
+ */
+export async function clearDedupeCacheForCurrentRequest(
+  dedupedFn: (...args: unknown[]) => unknown,
+) {
+  if (typeof dedupedFn !== 'function') {
+    throw new Error('dedupe: not a function');
+  }
+  const requestStore = cacheRegistry.get(dedupedFn);
+
+  if (!requestStore) {
+    throw new Error('dedupe: cache not found');
+  }
+  const { headers } = await import('next/headers');
+  const h = await headers();
+  return requestStore.delete(h);
 }

packages/flags/src/next/index.test.ts

@@ -1,5 +1,5 @@
 import { expect, it, describe, vi, beforeAll } from 'vitest';
-import { flag, precompute } from '.';
+import { flag, precompute, dedupe, clearDedupeCacheForCurrentRequest } from '.';
 import { IncomingMessage } from 'node:http';
 import { NextApiRequestCookies } from 'next/dist/server/api-utils';
 import { Readable } from 'node:stream';
@@ -42,6 +42,21 @@ function createRequest(cookies = {}): [
   return [request, socket];
 }
 
+describe('exports', () => {
+  it('should export flag', () => {
+    expect(typeof flag).toBe('function');
+  });
+  it('should export precompute', () => {
+    expect(typeof precompute).toBe('function');
+  });
+  it('should export dedupe', () => {
+    expect(typeof dedupe).toBe('function');
+  });
+  it('should export clearDedupeCacheForCurrentRequest', () => {
+    expect(typeof clearDedupeCacheForCurrentRequest).toBe('function');
+  });
+});
+
 describe('flag on app router', () => {
   beforeAll(() => {
     // a random secret for testing purposes

packages/flags/src/next/index.ts

@@ -483,5 +483,5 @@ export function getProviderData(
   return { definitions, hints: [] };
 }
 
-export { dedupe } from './dedupe';
+export { dedupe, clearDedupeCacheForCurrentRequest } from './dedupe';
 export { createFlagsDiscoveryEndpoint } from './create-flags-discovery-endpoint';