From dabd978ec514b190541a67e72da0b9296ef5378b Mon Sep 17 00:00:00 2001 From: Sebastian Markbage Date: Fri, 4 Oct 2024 12:49:52 -0400 Subject: [PATCH] Move AfterContext to WorkStore --- .../request-async-storage.external.ts | 2 - .../components/work-async-storage.external.ts | 2 + .../src/server/after/after-context.test.ts | 100 +++++++++--------- .../next/src/server/after/after-context.ts | 31 +++--- packages/next/src/server/after/after.ts | 35 +++--- .../async-storage/with-request-store.ts | 33 +----- .../server/async-storage/with-work-store.ts | 17 +++ .../server/route-modules/app-route/module.ts | 3 - packages/next/src/server/web/adapter.ts | 24 ++--- .../web/spec-extension/unstable-cache.ts | 1 + 10 files changed, 115 insertions(+), 133 deletions(-) diff --git a/packages/next/src/client/components/request-async-storage.external.ts b/packages/next/src/client/components/request-async-storage.external.ts index d7a8852b1c315..041869f9b6f7c 100644 --- a/packages/next/src/client/components/request-async-storage.external.ts +++ b/packages/next/src/client/components/request-async-storage.external.ts @@ -6,7 +6,6 @@ import type { ReadonlyRequestCookies } from '../../server/web/spec-extension/ada // Share the instance module in the next-shared layer import { requestAsyncStorage } from './request-async-storage-instance' with { 'turbopack-transition': 'next-shared' } -import type { AfterContext } from '../../server/after/after-context' import type { ServerComponentsHmrCache } from '../../server/response-cache' import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external' @@ -34,7 +33,6 @@ export interface RequestStore { readonly cookies: ReadonlyRequestCookies readonly mutableCookies: ResponseCookies readonly draftMode: DraftModeProvider - readonly afterContext: AfterContext | undefined readonly isHmrRefresh?: boolean readonly serverComponentsHmrCache?: ServerComponentsHmrCache } diff --git a/packages/next/src/client/components/work-async-storage.external.ts b/packages/next/src/client/components/work-async-storage.external.ts index be6e213fa61c9..8f5ef3c4dfa9c 100644 --- a/packages/next/src/client/components/work-async-storage.external.ts +++ b/packages/next/src/client/components/work-async-storage.external.ts @@ -6,6 +6,7 @@ import type { Revalidate } from '../../server/lib/revalidate' import type { FallbackRouteParams } from '../../server/request/fallback-params' import type { DeepReadonly } from '../../shared/lib/deep-readonly' import type { AppSegmentConfig } from '../../build/segment-config/app/app-segment-config' +import type { AfterContext } from '../../server/after/after-context' // Share the instance module in the next-shared layer import { workAsyncStorage } from './work-async-storage-instance' with { 'turbopack-transition': 'next-shared' } @@ -43,6 +44,7 @@ export interface WorkStore { dynamicShouldError?: boolean pendingRevalidates?: Record> pendingRevalidateWrites?: Array> // This is like pendingRevalidates but isn't used for deduping. + readonly afterContext: AfterContext | undefined dynamicUsageDescription?: string dynamicUsageStack?: string diff --git a/packages/next/src/server/after/after-context.test.ts b/packages/next/src/server/after/after-context.test.ts index df4508f9307e2..19b378870c250 100644 --- a/packages/next/src/server/after/after-context.test.ts +++ b/packages/next/src/server/after/after-context.test.ts @@ -1,18 +1,18 @@ import { DetachedPromise } from '../../lib/detached-promise' import { AsyncLocalStorage } from 'async_hooks' -import type { RequestStore } from '../../client/components/request-async-storage.external' +import type { WorkStore } from '../../client/components/work-async-storage.external' import type { AfterContext } from './after-context' describe('AfterContext', () => { // 'async-local-storage.ts' needs `AsyncLocalStorage` on `globalThis` at import time, // so we have to do some contortions here to set it up before running anything else - type RASMod = - typeof import('../../client/components/request-async-storage.external') + type WASMod = + typeof import('../../client/components/work-async-storage.external') type AfterMod = typeof import('./after') type AfterContextMod = typeof import('./after-context') - let requestAsyncStorage: RASMod['requestAsyncStorage'] + let workAsyncStorage: WASMod['workAsyncStorage'] let AfterContext: AfterContextMod['AfterContext'] let after: AfterMod['unstable_after'] @@ -20,10 +20,10 @@ describe('AfterContext', () => { // @ts-expect-error globalThis.AsyncLocalStorage = AsyncLocalStorage - const RASMod = await import( - '../../client/components/request-async-storage.external' + const WASMod = await import( + '../../client/components/work-async-storage.external' ) - requestAsyncStorage = RASMod.requestAsyncStorage + workAsyncStorage = WASMod.workAsyncStorage const AfterContextMod = await import('./after-context') AfterContext = AfterContextMod.AfterContext @@ -33,11 +33,9 @@ describe('AfterContext', () => { }) const createRun = - (afterContext: AfterContext, requestStore: RequestStore) => + (_afterContext: AfterContext, workStore: WorkStore) => (cb: () => T): T => { - return afterContext.run(requestStore, () => - requestAsyncStorage.run(requestStore, cb) - ) + return workAsyncStorage.run(workStore, cb) } it('runs after() callbacks from a run() callback that resolves', async () => { @@ -54,8 +52,8 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) - const run = createRun(afterContext, requestStore) + const workStore = createMockWorkStore(afterContext) + const run = createRun(afterContext, workStore) // ================================== @@ -120,9 +118,9 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) - const run = createRun(afterContext, requestStore) + const run = createRun(afterContext, workStore) // ================================== @@ -167,9 +165,9 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) - const run = createRun(afterContext, requestStore) + const run = createRun(afterContext, workStore) // ================================== @@ -257,8 +255,8 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) - const run = createRun(afterContext, requestStore) + const workStore = createMockWorkStore(afterContext) + const run = createRun(afterContext, workStore) // ================================== @@ -316,9 +314,9 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) - const run = createRun(afterContext, requestStore) + const run = createRun(afterContext, workStore) // ================================== @@ -353,7 +351,7 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) // ================================== @@ -367,13 +365,11 @@ describe('AfterContext', () => { const promise3 = new DetachedPromise() const afterCallback3 = jest.fn(() => promise3.promise) - requestAsyncStorage.run(requestStore, () => - afterContext.run(requestStore, () => { - after(afterCallback1) - after(afterCallback2) - after(afterCallback3) - }) - ) + workAsyncStorage.run(workStore, () => { + after(afterCallback1) + after(afterCallback2) + after(afterCallback3) + }) expect(afterCallback1).not.toHaveBeenCalled() expect(afterCallback2).not.toHaveBeenCalled() @@ -405,9 +401,9 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) - const run = createRun(afterContext, requestStore) + const run = createRun(afterContext, workStore) // ================================== @@ -434,9 +430,9 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) + const workStore = createMockWorkStore(afterContext) - const run = createRun(afterContext, requestStore) + const run = createRun(afterContext, workStore) // ================================== @@ -452,7 +448,7 @@ describe('AfterContext', () => { expect(afterCallback1).not.toHaveBeenCalled() }) - it('shadows requestAsyncStorage within after callbacks', async () => { + it('does NOT shadow workAsyncStorage within after callbacks', async () => { const waitUntil = jest.fn() let onCloseCallback: (() => void) | undefined = undefined @@ -465,19 +461,19 @@ describe('AfterContext', () => { onClose, }) - const requestStore = createMockRequestStore(afterContext) - const run = createRun(afterContext, requestStore) + const workStore = createMockWorkStore(afterContext) + const run = createRun(afterContext, workStore) // ================================== const stores = new DetachedPromise< - [RequestStore | undefined, RequestStore | undefined] + [WorkStore | undefined, WorkStore | undefined] >() await run(async () => { - const store1 = requestAsyncStorage.getStore() + const store1 = workAsyncStorage.getStore() after(() => { - const store2 = requestAsyncStorage.getStore() + const store2 = workAsyncStorage.getStore() stores.resolve([store1, store2]) }) }) @@ -486,30 +482,34 @@ describe('AfterContext', () => { onCloseCallback!() const [store1, store2] = await stores.promise - // if we use .toBe, the proxy from createMockRequestStore throws because jest checks '$$typeof' + // if we use .toBe, the proxy from createMockWorkStore throws because jest checks '$$typeof' expect(store1).toBeTruthy() expect(store2).toBeTruthy() - expect(store1 === requestStore).toBe(true) - expect(store2 !== store1).toBe(true) + expect(store1 === workStore).toBe(true) + expect(store2 === store1).toBe(true) }) }) -const createMockRequestStore = (afterContext: AfterContext): RequestStore => { - const partialStore: Partial = { - url: { pathname: '/', search: '' }, +const createMockWorkStore = (afterContext: AfterContext): WorkStore => { + const partialStore: Partial = { afterContext: afterContext, - draftMode: undefined, - isHmrRefresh: false, - serverComponentsHmrCache: undefined, + forceStatic: false, + forceDynamic: false, + dynamicShouldError: false, + isStaticGeneration: false, + revalidatedTags: [], + pendingRevalidates: undefined, + pendingRevalidateWrites: undefined, + incrementalCache: undefined, } - return new Proxy(partialStore as RequestStore, { + return new Proxy(partialStore as WorkStore, { get(target, key) { if (key in target) { return target[key as keyof typeof target] } throw new Error( - `RequestStore property not mocked: '${typeof key === 'symbol' ? key.toString() : key}'` + `WorkStore property not mocked: '${typeof key === 'symbol' ? key.toString() : key}'` ) }, }) diff --git a/packages/next/src/server/after/after-context.ts b/packages/next/src/server/after/after-context.ts index d2816b03c5053..36427f180f479 100644 --- a/packages/next/src/server/after/after-context.ts +++ b/packages/next/src/server/after/after-context.ts @@ -33,11 +33,6 @@ export class AfterContext { this.callbackQueue.pause() } - public run(requestStore: RequestStore, callback: () => T): T { - this.requestStore = requestStore - return callback() - } - public after(task: AfterTask): void { if (isThenable(task)) { task.catch(() => {}) // avoid unhandled rejection crashes @@ -61,9 +56,10 @@ export class AfterContext { errorWaitUntilNotAvailable() } if (!this.requestStore) { - throw new InvariantError( - 'unstable_after: Expected `AfterContext.requestStore` to be initialized' - ) + // 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.requestStore = requestAsyncStorage.getStore() } if (!this.onClose) { throw new InvariantError( @@ -98,19 +94,27 @@ export class AfterContext { private async runCallbacksOnClose() { await new Promise((resolve) => this.onClose!(resolve)) - return this.runCallbacks(this.requestStore!) + return this.runCallbacks(this.requestStore) } - private async runCallbacks(requestStore: RequestStore): Promise { + private async runCallbacks( + requestStore: undefined | RequestStore + ): Promise { if (this.callbackQueue.size === 0) return - const readonlyRequestStore: RequestStore = - wrapRequestStoreForAfterCallbacks(requestStore) + const readonlyRequestStore: undefined | RequestStore = + requestStore === undefined + ? undefined + : // TODO: This is not sufficient. It should just be the same store that mutates. + wrapRequestStoreForAfterCallbacks(requestStore) const workStore = workAsyncStorage.getStore() return withExecuteRevalidates(workStore, () => - requestAsyncStorage.run(readonlyRequestStore, async () => { + // 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. + requestAsyncStorage.run(readonlyRequestStore as any, async () => { this.callbackQueue.start() await this.callbackQueue.onIdle() }) @@ -141,7 +145,6 @@ function wrapRequestStoreForAfterCallbacks( }, // TODO(after): calling a `cookies.set()` in an after() that's in an action doesn't currently error. mutableCookies: new ResponseCookies(new Headers()), - afterContext: requestStore.afterContext, isHmrRefresh: requestStore.isHmrRefresh, serverComponentsHmrCache: requestStore.serverComponentsHmrCache, } diff --git a/packages/next/src/server/after/after.ts b/packages/next/src/server/after/after.ts index 4e648d021f2c4..45f38c9f0e511 100644 --- a/packages/next/src/server/after/after.ts +++ b/packages/next/src/server/after/after.ts @@ -1,4 +1,3 @@ -import { requestAsyncStorage } from '../../client/components/request-async-storage.external' import { workAsyncStorage } from '../../client/components/work-async-storage.external' import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external' import { StaticGenBailoutError } from '../../client/components/static-generation-bailout' @@ -11,28 +10,20 @@ export type AfterCallback = () => T | Promise /** * This function allows you to schedule callbacks to be executed after the current request finishes. */ -export function unstable_after(task: AfterTask) { - const callingExpression = 'unstable_after' - - // TODO: This is not safe. afterContext should move to WorkStore. - const requestStore = requestAsyncStorage.getStore() - if (!requestStore) { - throw new Error( - `\`${callingExpression}\` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context` - ) - } - - const { afterContext } = requestStore - if (!afterContext) { - throw new Error( - '`unstable_after()` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.' - ) - } - +export function unstable_after(task: AfterTask): void { const workStore = workAsyncStorage.getStore() const cacheStore = cacheAsyncStorage.getStore() if (workStore) { + const { afterContext } = workStore + if (!afterContext) { + throw new Error( + '`unstable_after()` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.' + ) + } + + // TODO: After should not cause dynamic. + const callingExpression = 'unstable_after' if (workStore.forceStatic) { throw new StaticGenBailoutError( `Route ${workStore.route} with \`dynamic = "force-static"\` couldn't be rendered statically because it used \`${callingExpression}\`. See more info here: https://nextjs.org/docs/app/building-your-application/rendering/static-and-dynamic#dynamic-rendering` @@ -40,7 +31,9 @@ export function unstable_after(task: AfterTask) { } else { markCurrentScopeAsDynamic(workStore, cacheStore, callingExpression) } - } - return afterContext.after(task) + afterContext.after(task) + } else { + // TODO: Error for pages? + } } diff --git a/packages/next/src/server/async-storage/with-request-store.ts b/packages/next/src/server/async-storage/with-request-store.ts index 0dd70bad2eba1..7078eb46b2ede 100644 --- a/packages/next/src/server/async-storage/with-request-store.ts +++ b/packages/next/src/server/async-storage/with-request-store.ts @@ -20,8 +20,6 @@ import { import { ResponseCookies, RequestCookies } from '../web/spec-extension/cookies' import { DraftModeProvider } from './draft-mode-provider' import { splitCookiesString } from '../web/utils' -import { AfterContext } from '../after/after-context' -import type { RequestLifecycleOpts } from '../base-server' import type { ServerComponentsHmrCache } from '../response-cache' function getHeaders(headers: Headers | IncomingHttpHeaders): ReadonlyHeaders { @@ -41,11 +39,9 @@ function getMutableCookies( return MutableRequestCookiesAdapter.wrap(cookies, onUpdateCookies) } -export type WrapperRenderOpts = RequestLifecycleOpts & - Partial> & { - experimental: Pick - previewProps?: __ApiPreviewProps - } +export type WrapperRenderOpts = Partial> & { + previewProps?: __ApiPreviewProps +} export type RequestContext = RequestResponsePair & { /** @@ -188,34 +184,11 @@ export const withRequestStore: WithStore = < return cache.draftMode }, - afterContext: createAfterContext(renderOpts), isHmrRefresh, serverComponentsHmrCache: serverComponentsHmrCache || (globalThis as any).__serverComponentsHmrCache, } - if (store.afterContext) { - return store.afterContext.run(store, () => - storage.run(store, callback, store) - ) - } - return storage.run(store, callback, store) } - -function createAfterContext( - renderOpts: WrapperRenderOpts | undefined -): AfterContext | undefined { - if (!isAfterEnabled(renderOpts)) { - return undefined - } - const { waitUntil, onClose } = renderOpts - return new AfterContext({ waitUntil, onClose }) -} - -function isAfterEnabled( - renderOpts: WrapperRenderOpts | undefined -): renderOpts is WrapperRenderOpts { - return renderOpts?.experimental?.after ?? false -} diff --git a/packages/next/src/server/async-storage/with-work-store.ts b/packages/next/src/server/async-storage/with-work-store.ts index 8a0c70152e5bf..7ae5c9f9b28cc 100644 --- a/packages/next/src/server/async-storage/with-work-store.ts +++ b/packages/next/src/server/async-storage/with-work-store.ts @@ -8,6 +8,8 @@ import type { RequestLifecycleOpts } from '../base-server' import type { FallbackRouteParams } from '../request/fallback-params' import type { AppSegmentConfig } from '../../build/segment-config/app/app-segment-config' +import { AfterContext } from '../after/after-context' + import { normalizeAppPath } from '../../shared/lib/router/utils/app-paths' export type WorkStoreContext = { @@ -118,6 +120,8 @@ export const withWorkStore: WithStore = ( buildId: renderOpts.buildId, reactLoadableManifest: renderOpts?.reactLoadableManifest || {}, assetPrefix: renderOpts?.assetPrefix || '', + + afterContext: createAfterContext(renderOpts), } // TODO: remove this when we resolve accessing the store outside the execution context @@ -125,3 +129,16 @@ export const withWorkStore: WithStore = ( return storage.run(store, callback, store) } + +function createAfterContext( + renderOpts: Partial & { + experimental: Pick + } +): AfterContext | undefined { + const isAfterEnabled = renderOpts?.experimental?.after ?? false + if (!isAfterEnabled) { + return undefined + } + const { waitUntil, onClose } = renderOpts + return new AfterContext({ waitUntil, onClose }) +} diff --git a/packages/next/src/server/route-modules/app-route/module.ts b/packages/next/src/server/route-modules/app-route/module.ts index 90a177623110d..f97674553e2c8 100644 --- a/packages/next/src/server/route-modules/app-route/module.ts +++ b/packages/next/src/server/route-modules/app-route/module.ts @@ -554,9 +554,6 @@ export class AppRouteRouteModule extends RouteModule< url: req.nextUrl, renderOpts: { previewProps: context.prerenderManifest.preview, - waitUntil: context.renderOpts.waitUntil, - onClose: context.renderOpts.onClose, - experimental: context.renderOpts.experimental, }, } diff --git a/packages/next/src/server/web/adapter.ts b/packages/next/src/server/web/adapter.ts index d4caf504c61df..08d15985906f2 100644 --- a/packages/next/src/server/web/adapter.ts +++ b/packages/next/src/server/web/adapter.ts @@ -12,12 +12,12 @@ import { stripInternalSearchParams } from '../internal-utils' import { normalizeRscURL } from '../../shared/lib/router/utils/app-paths' import { FLIGHT_HEADERS } from '../../client/components/app-router-headers' import { ensureInstrumentationRegistered } from './globals' -import { - withRequestStore, - type WrapperRenderOpts, -} from '../async-storage/with-request-store' +import { withRequestStore } from '../async-storage/with-request-store' import { requestAsyncStorage } from '../../client/components/request-async-storage.external' -import { withWorkStore } from '../async-storage/with-work-store' +import { + withWorkStore, + type WorkStoreContext, +} from '../async-storage/with-work-store' import { workAsyncStorage } from '../../client/components/work-async-storage.external' import { NEXT_ROUTER_PREFETCH_HEADER } from '../../client/components/app-router-headers' import { getTracer } from '../lib/trace/tracer' @@ -212,11 +212,12 @@ export async function adapter( // if we're in an edge function, we only get a subset of `nextConfig` (no `experimental`), // so we have to inject it via DefinePlugin. // in `next start` this will be passed normally (see `NextNodeServer.runMiddleware`). + const isAfterEnabled = params.request.nextConfig?.experimental?.after ?? !!process.env.__NEXT_AFTER - let waitUntil: WrapperRenderOpts['waitUntil'] = undefined + let waitUntil: WorkStoreContext['renderOpts']['waitUntil'] = undefined let closeController: CloseController | undefined = undefined if (isAfterEnabled) { @@ -250,6 +251,10 @@ export async function adapter( }, buildId: buildId ?? '', supportsDynamicResponse: true, + waitUntil, + onClose: closeController + ? closeController.onClose.bind(closeController) + : undefined, }, requestEndedState: { ended: false }, isPrefetchRequest: request.headers.has( @@ -268,13 +273,6 @@ export async function adapter( cookiesFromResponse = cookies }, previewProps, - waitUntil, - onClose: closeController - ? closeController.onClose.bind(closeController) - : undefined, - experimental: { - after: isAfterEnabled, - }, }, }, () => params.handler(request, event) diff --git a/packages/next/src/server/web/spec-extension/unstable-cache.ts b/packages/next/src/server/web/spec-extension/unstable-cache.ts index 9da73106e5f81..15cde4723d1ee 100644 --- a/packages/next/src/server/web/spec-extension/unstable-cache.ts +++ b/packages/next/src/server/web/spec-extension/unstable-cache.ts @@ -356,6 +356,7 @@ export function unstable_cache( isStaticGeneration: false, fallbackRouteParams: null, buildId: '', // Since this is a fake one it can't "use cache" anyway. + afterContext: undefined, // TODO: Support Pages after() }, () => cacheAsyncStorage.run(innerCacheStore, cb, ...args) )