From 979ac407bdc13acb3b0fedbd9201baa2f896ed03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= Date: Tue, 8 Oct 2024 10:23:25 -0400 Subject: [PATCH] Run after() callbacks in the same context as the outer function (#70810) Stacked on #70806. We want `after` callbacks to run in the same ALS context that `after` was called in. This can be achieved via https://nodejs.org/api/async_context.html#static-method-asynclocalstoragebindfn We want this be cause these should be as close as possible in every way (except timing): ``` after(() => x()) after(x()) await x() ``` We'll need a different approach for disabling setting cookies (and revalidate). This will be done in a follow up. --------- Co-authored-by: Janka Uryga --- .../client/components/async-local-storage.ts | 11 +++ .../src/server/after/after-context.test.ts | 46 +++++++++++ .../next/src/server/after/after-context.ts | 77 ++++--------------- 3 files changed, 72 insertions(+), 62 deletions(-) diff --git a/packages/next/src/client/components/async-local-storage.ts b/packages/next/src/client/components/async-local-storage.ts index 166f34d3f14c1..b3d8d247eeab4 100644 --- a/packages/next/src/client/components/async-local-storage.ts +++ b/packages/next/src/client/components/async-local-storage.ts @@ -27,6 +27,10 @@ class FakeAsyncLocalStorage enterWith(): void { throw sharedAsyncLocalStorageNotAvailableError } + + static bind(fn: T): T { + return fn + } } const maybeGlobalAsyncLocalStorage = @@ -41,6 +45,13 @@ export function createAsyncLocalStorage< return new FakeAsyncLocalStorage() } +export function bindSnapshot(fn: T): T { + if (maybeGlobalAsyncLocalStorage) { + return maybeGlobalAsyncLocalStorage.bind(fn) + } + return FakeAsyncLocalStorage.bind(fn) +} + export function createSnapshot(): ( fn: (...args: TArgs) => R, ...args: TArgs diff --git a/packages/next/src/server/after/after-context.test.ts b/packages/next/src/server/after/after-context.test.ts index 19b378870c250..da6d024b0f405 100644 --- a/packages/next/src/server/after/after-context.test.ts +++ b/packages/next/src/server/after/after-context.test.ts @@ -488,6 +488,52 @@ describe('AfterContext', () => { expect(store1 === workStore).toBe(true) expect(store2 === store1).toBe(true) }) + + it('preserves the ALS context the callback was created in', async () => { + type TestStore = string + const testStorage = new AsyncLocalStorage() + + const waitUntil = jest.fn() + + let onCloseCallback: (() => void) | undefined = undefined + const onClose = jest.fn((cb) => { + onCloseCallback = cb + }) + + const afterContext = new AfterContext({ + waitUntil, + onClose, + }) + + const workStore = createMockWorkStore(afterContext) + const run = createRun(afterContext, workStore) + + // ================================== + + const stores = new DetachedPromise< + [TestStore | undefined, TestStore | undefined] + >() + + await testStorage.run('value', () => + run(async () => { + const store1 = testStorage.getStore() + after(() => { + const store2 = testStorage.getStore() + stores.resolve([store1, store2]) + }) + }) + ) + + // the response is done. + onCloseCallback!() + + const [store1, store2] = await stores.promise + // if we use .toBe, the proxy from createMockWorkStore throws because jest checks '$$typeof' + expect(store1).toBeDefined() + expect(store2).toBeDefined() + expect(store1 === 'value').toBe(true) + expect(store2 === store1).toBe(true) + }) }) const createMockWorkStore = (afterContext: AfterContext): WorkStore => { diff --git a/packages/next/src/server/after/after-context.ts b/packages/next/src/server/after/after-context.ts index d23d518dbf4ee..39cc107b06a31 100644 --- a/packages/next/src/server/after/after-context.ts +++ b/packages/next/src/server/after/after-context.ts @@ -1,16 +1,11 @@ import PromiseQueue from 'next/dist/compiled/p-queue' -import { - workUnitAsyncStorage, - type RequestStore, - type WorkUnitStore, -} from '../../server/app-render/work-unit-async-storage.external' -import { ResponseCookies } from '../web/spec-extension/cookies' import type { RequestLifecycleOpts } from '../base-server' import type { AfterCallback, AfterTask } from './after' import { InvariantError } from '../../shared/lib/invariant-error' import { isThenable } from '../../shared/lib/is-thenable' import { workAsyncStorage } from '../../client/components/work-async-storage.external' import { withExecuteRevalidates } from './revalidation-utils' +import { bindSnapshot } from '../../client/components/async-local-storage' export type AfterContextOpts = { waitUntil: RequestLifecycleOpts['waitUntil'] | undefined @@ -21,8 +16,6 @@ export class AfterContext { private waitUntil: RequestLifecycleOpts['waitUntil'] | undefined private onClose: RequestLifecycleOpts['onClose'] | undefined - private workUnitStore: WorkUnitStore | undefined - private runCallbacksOnClosePromise: Promise | undefined private callbackQueue: PromiseQueue @@ -56,12 +49,6 @@ export class AfterContext { if (!this.waitUntil) { errorWaitUntilNotAvailable() } - if (!this.workUnitStore) { - // We just stash the first request store we have but this is not sufficient. - // TODO: We should store a request store per callback since each callback might - // be inside a different store. E.g. inside different batched actions, prerenders or caches. - this.workUnitStore = workUnitAsyncStorage.getStore() - } if (!this.onClose) { throw new InvariantError( 'unstable_after: Missing `onClose` implementation' @@ -70,15 +57,16 @@ export class AfterContext { // this should only happen once. if (!this.runCallbacksOnClosePromise) { - // NOTE: We're creating a promise here, which means that - // we will propagate any AsyncLocalStorage contexts we're currently in - // to the callbacks that'll execute later. - // This includes e.g. `workUnitAsyncStorage` and React's `requestStorage` (which backs `React.cache()`). this.runCallbacksOnClosePromise = this.runCallbacksOnClose() this.waitUntil(this.runCallbacksOnClosePromise) } - const wrappedCallback = async () => { + // Bind the callback to the current execution context (i.e. preserve all currently available ALS-es). + // We do this because we want all of these to be equivalent in every regard except timing: + // after(() => x()) + // after(x()) + // await x() + const wrappedCallback = bindSnapshot(async () => { try { await callback() } catch (err) { @@ -88,38 +76,26 @@ export class AfterContext { err ) } - } + }) this.callbackQueue.add(wrappedCallback) } private async runCallbacksOnClose() { await new Promise((resolve) => this.onClose!(resolve)) - return this.runCallbacks(this.workUnitStore) + return this.runCallbacks() } - private async runCallbacks( - workUnitStore: undefined | WorkUnitStore - ): Promise { + private async runCallbacks(): Promise { if (this.callbackQueue.size === 0) return - const readonlyRequestStore: undefined | WorkUnitStore = - workUnitStore === undefined || workUnitStore.type !== 'request' - ? undefined - : // TODO: This is not sufficient. It should just be the same store that mutates. - wrapRequestStoreForAfterCallbacks(workUnitStore) - const workStore = workAsyncStorage.getStore() - return withExecuteRevalidates(workStore, () => - // Clearing it out or running the first request store. - // TODO: This needs to be the request store that was active at the time the - // callback was scheduled but p-queue makes this hard so need further refactoring. - workUnitAsyncStorage.run(readonlyRequestStore as any, async () => { - this.callbackQueue.start() - await this.callbackQueue.onIdle() - }) - ) + // TODO(after): Change phase in workUnitStore to disable e.g. `cookies().set()` + return withExecuteRevalidates(workStore, () => { + this.callbackQueue.start() + return this.callbackQueue.onIdle() + }) } } @@ -128,26 +104,3 @@ function errorWaitUntilNotAvailable(): never { '`unstable_after()` will not work correctly, because `waitUntil` is not available in the current environment.' ) } - -/** Disable mutations of `requestStore` within `after()` and disallow nested after calls. */ -function wrapRequestStoreForAfterCallbacks( - requestStore: RequestStore -): RequestStore { - return { - type: 'request', - url: requestStore.url, - get headers() { - return requestStore.headers - }, - get cookies() { - return requestStore.cookies - }, - get draftMode() { - return requestStore.draftMode - }, - // TODO(after): calling a `cookies.set()` in an after() that's in an action doesn't currently error. - mutableCookies: new ResponseCookies(new Headers()), - isHmrRefresh: requestStore.isHmrRefresh, - serverComponentsHmrCache: requestStore.serverComponentsHmrCache, - } -}