Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

add phase tracking to workUnitStore #71030

Merged
merged 3 commits into from
Oct 10, 2024
Merged

add phase tracking to workUnitStore #71030

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16d1fad
add phase to workUnitStore
lubieowoce Oct 8, 2024
5b067e6
change phase to "after" when running callbacks
lubieowoce Oct 8, 2024
83f4e78
some assertions about ALS in after
lubieowoce Oct 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions packages/next/src/server/after/after-context.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,19 @@ import { DetachedPromise } from '../../lib/detached-promise'
import { AsyncLocalStorage } from 'async_hooks'

import type { WorkStore } from '../app-render/work-async-storage.external'
import type { WorkUnitStore } from '../app-render/work-unit-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 WASMod = typeof import('../app-render/work-async-storage.external')
type WSMod = typeof import('../app-render/work-unit-async-storage.external')
type AfterMod = typeof import('./after')
type AfterContextMod = typeof import('./after-context')

let workAsyncStorage: WASMod['workAsyncStorage']
let workUnitAsyncStorage: WSMod['workUnitAsyncStorage']
let AfterContext: AfterContextMod['AfterContext']
let after: AfterMod['unstable_after']

Expand All @@ -22,6 +25,9 @@ describe('AfterContext', () => {
const WASMod = await import('../app-render/work-async-storage.external')
workAsyncStorage = WASMod.workAsyncStorage

const WSMod = await import('../app-render/work-unit-async-storage.external')
workUnitAsyncStorage = WSMod.workUnitAsyncStorage

const AfterContextMod = await import('./after-context')
AfterContext = AfterContextMod.AfterContext

Expand All @@ -32,7 +38,9 @@ describe('AfterContext', () => {
const createRun =
(_afterContext: AfterContext, workStore: WorkStore) =>
<T>(cb: () => T): T => {
return workAsyncStorage.run(workStore, cb)
return workAsyncStorage.run(workStore, () =>
workUnitAsyncStorage.run(createMockWorkUnitStore(), cb)
)
}

it('runs after() callbacks from a run() callback that resolves', async () => {
Expand Down Expand Up @@ -362,11 +370,13 @@ describe('AfterContext', () => {
const promise3 = new DetachedPromise<string>()
const afterCallback3 = jest.fn(() => promise3.promise)

workAsyncStorage.run(workStore, () => {
after(afterCallback1)
after(afterCallback2)
after(afterCallback3)
})
workAsyncStorage.run(workStore, () =>
workUnitAsyncStorage.run(createMockWorkUnitStore(), () => {
after(afterCallback1)
after(afterCallback2)
after(afterCallback3)
})
)

expect(afterCallback1).not.toHaveBeenCalled()
expect(afterCallback2).not.toHaveBeenCalled()
Expand Down Expand Up @@ -557,3 +567,7 @@ const createMockWorkStore = (afterContext: AfterContext): WorkStore => {
},
})
}

const createMockWorkUnitStore = () => {
return { phase: 'render' } as WorkUnitStore
}
21 changes: 20 additions & 1 deletion packages/next/src/server/after/after-context.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@ import { isThenable } from '../../shared/lib/is-thenable'
import { workAsyncStorage } from '../app-render/work-async-storage.external'
import { withExecuteRevalidates } from './revalidation-utils'
import { bindSnapshot } from '../app-render/async-local-storage'
import {
workUnitAsyncStorage,
type WorkUnitStore,
} from '../app-render/work-unit-async-storage.external'

export type AfterContextOpts = {
waitUntil: RequestLifecycleOpts['waitUntil'] | undefined
Expand All @@ -18,6 +22,7 @@ export class AfterContext {

private runCallbacksOnClosePromise: Promise<void> | undefined
private callbackQueue: PromiseQueue
private workUnitStores = new Set<WorkUnitStore>()

constructor({ waitUntil, onClose }: AfterContextOpts) {
this.waitUntil = waitUntil
Expand Down Expand Up @@ -55,6 +60,14 @@ export class AfterContext {
)
}

const workUnitStore = workUnitAsyncStorage.getStore()
if (!workUnitStore) {
throw new InvariantError(
'Missing workUnitStore in AfterContext.addCallback'
)
}
this.workUnitStores.add(workUnitStore)

// this should only happen once.
if (!this.runCallbacksOnClosePromise) {
this.runCallbacksOnClosePromise = this.runCallbacksOnClose()
Expand Down Expand Up @@ -89,9 +102,15 @@ export class AfterContext {
private async runCallbacks(): Promise<void> {
if (this.callbackQueue.size === 0) return

for (const workUnitStore of this.workUnitStores) {
workUnitStore.phase = 'after'
}

const workStore = workAsyncStorage.getStore()
if (!workStore) {
throw new InvariantError('Missing workStore in AfterContext.runCallbacks')
}

// TODO(after): Change phase in workUnitStore to disable e.g. `cookies().set()`
return withExecuteRevalidates(workStore, () => {
this.callbackQueue.start()
return this.callbackQueue.onIdle()
Expand Down
39 changes: 21 additions & 18 deletions packages/next/src/server/after/after.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,26 +14,29 @@ export function unstable_after<T>(task: AfterTask<T>): void {
const workStore = workAsyncStorage.getStore()
const workUnitStore = workUnitAsyncStorage.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.'
)
}
if (!workStore) {
// TODO(after): the linked docs page talks about *dynamic* APIs, which unstable_after soon won't be anymore
throw new Error(
'`unstable_after` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context'
)
}

// 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`
)
} else {
markCurrentScopeAsDynamic(workStore, workUnitStore, callingExpression)
}
const { afterContext } = workStore
if (!afterContext) {
throw new Error(
'`unstable_after` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.'
)
}

afterContext.after(task)
// 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`
)
} else {
// TODO: Error for pages?
markCurrentScopeAsDynamic(workStore, workUnitStore, callingExpression)
}

afterContext.after(task)
}
2 changes: 2 additions & 0 deletions packages/next/src/server/app-render/action-handler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -459,6 +459,8 @@ export async function handleAction({
)
}

requestStore.phase = 'action'

// When running actions the default is no-store, you can still `cache: 'force-cache'`
workStore.fetchCache = 'default-no-store'

Expand Down
14 changes: 14 additions & 0 deletions packages/next/src/server/app-render/app-render.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,6 +375,7 @@ async function generateDynamicRSCPayload(
skipFlight: boolean
}
): Promise<RSCPayload> {
ctx.requestStore.phase = 'render'
Copy link
Member Author

@lubieowoce lubieowoce Oct 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i'm not 100% sure if this is the way to do it -- i can see an argument that this should be set independently, before this function runs, and just be an invariant here. but this works (in particular for the case of going action -> render) and it's hard to do it differently for after (which is currently kinda in control of when its stuff runs without exposing much to the outside world), so this is the convention i've settled on

// Flight data that is going to be passed to the browser.
// Currently a single item array but in the future multiple patches might be combined in a single request.

Expand Down Expand Up @@ -1283,6 +1284,7 @@ export const renderToHTMLOrFlight: AppPageRender = (
req,
url,
res,
phase: 'render',
renderOpts,
isHmrRefresh,
serverComponentsHmrCache,
Expand Down Expand Up @@ -1722,6 +1724,8 @@ async function prerenderToStream(
workStore: WorkStore,
tree: LoaderTree
): Promise<PrerenderToStreamResult> {
ctx.requestStore.phase = 'render'

// When prerendering formState is always null. We still include it
// because some shared APIs expect a formState value and this is slightly
// more explicit than making it an optional function argument
Expand Down Expand Up @@ -1843,6 +1847,7 @@ async function prerenderToStream(
const prospectiveRenderPrerenderStore: PrerenderStore =
(prerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: flightController.signal,
cacheSignal,
Expand Down Expand Up @@ -1936,6 +1941,7 @@ async function prerenderToStream(

const finalRenderPrerenderStore: PrerenderStore = (prerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: flightController.signal,
// During the final prerender we don't need to track cache access so we omit the signal
Expand Down Expand Up @@ -1995,6 +2001,7 @@ async function prerenderToStream(
const SSRController = new AbortController()
const ssrPrerenderStore: PrerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: SSRController.signal,
// For HTML Generation we don't need to track cache reads (RSC only)
Expand Down Expand Up @@ -2210,6 +2217,7 @@ async function prerenderToStream(
const prospectiveRenderPrerenderStore: PrerenderStore =
(prerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: flightController.signal,
cacheSignal,
Expand Down Expand Up @@ -2292,6 +2300,7 @@ async function prerenderToStream(

const finalRenderPrerenderStore: PrerenderStore = (prerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: flightController.signal,
// During the final prerender we don't need to track cache access so we omit the signal
Expand All @@ -2305,6 +2314,7 @@ async function prerenderToStream(
const SSRController = new AbortController()
const ssrPrerenderStore: PrerenderStore = {
type: 'prerender',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
renderSignal: SSRController.signal,
// For HTML Generation we don't need to track cache reads (RSC only)
Expand Down Expand Up @@ -2492,6 +2502,7 @@ async function prerenderToStream(
)
const reactServerPrerenderStore: PrerenderStore = (prerenderStore = {
type: 'prerender-ppr',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
dynamicTracking,
revalidate: INFINITE_CACHE,
Expand Down Expand Up @@ -2520,6 +2531,7 @@ async function prerenderToStream(

const ssrPrerenderStore: PrerenderStore = {
type: 'prerender-ppr',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
dynamicTracking,
revalidate: INFINITE_CACHE,
Expand Down Expand Up @@ -2686,6 +2698,7 @@ async function prerenderToStream(
} else {
const prerenderLegacyStore: PrerenderStore = (prerenderStore = {
type: 'prerender-legacy',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
revalidate: INFINITE_CACHE,
tags: [...ctx.requestStore.implicitTags],
Expand Down Expand Up @@ -2842,6 +2855,7 @@ async function prerenderToStream(

const prerenderLegacyStore: PrerenderStore = (prerenderStore = {
type: 'prerender-legacy',
phase: 'render',
implicitTags: ctx.requestStore.implicitTags,
revalidate: INFINITE_CACHE,
tags: [...ctx.requestStore.implicitTags],
Expand Down
19 changes: 13 additions & 6 deletions packages/next/src/server/app-render/work-unit-async-storage.external.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,13 @@ import type { DynamicTrackingState } from './dynamic-rendering'
import { workUnitAsyncStorage } from './work-unit-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
import type { ServerComponentsHmrCache } from '../response-cache'

type WorkUnitPhase = 'action' | 'render' | 'after'

type PhasePartial = {
/** NOTE: Will be mutated as phases change */
phase: WorkUnitPhase
}

export type RequestStore = {
type: 'request'

Expand Down Expand Up @@ -41,7 +48,7 @@ export type RequestStore = {

// DEV-only
usedDynamic?: boolean
}
} & PhasePartial
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

personally not a fan of use of partials for unioning. In this case it makes it possible for prerender stores to be constructed with or turned into action phase which should be impossible. It also makes it harder to see the complete type details in the IDE

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In reviewing further I see why prerender store needs action as a valid phase. Kind of unfortunate since this is only true in the route handler case. But still I think I'd just recommend we repeat the valid phases per store type

Copy link
Member Author

@lubieowoce lubieowoce Oct 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah originally i only had action on RequestStore but in practice it seems like it's kinda hard to narrow these. we'll probably end up revisiting this


/**
* The Prerender store is for tracking information related to prerenders.
Expand Down Expand Up @@ -81,7 +88,7 @@ export type PrerenderStoreModern = {
// Collected revalidate times and tags for this document during the prerender.
revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
tags: null | string[]
}
} & PhasePartial

export type PrerenderStorePPR = {
type: 'prerender-ppr'
Expand All @@ -90,15 +97,15 @@ export type PrerenderStorePPR = {
// Collected revalidate times and tags for this document during the prerender.
revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
tags: null | string[]
}
} & PhasePartial

export type PrerenderStoreLegacy = {
type: 'prerender-legacy'
readonly implicitTags: string[]
// Collected revalidate times and tags for this document during the prerender.
revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
tags: null | string[]
}
} & PhasePartial

export type PrerenderStore =
| PrerenderStoreLegacy
Expand All @@ -112,11 +119,11 @@ export type UseCacheStore = {
revalidate: number // implicit revalidate time from inner caches / fetches
explicitRevalidate: undefined | number // explicit revalidate time from cacheLife() calls
tags: null | string[]
}
} & PhasePartial

export type UnstableCacheStore = {
type: 'unstable-cache'
}
} & PhasePartial

/**
* The Cache store is for tracking information inside a "use cache" or unstable_cache context.
Expand Down
3 changes: 3 additions & 0 deletions packages/next/src/server/async-storage/with-request-store.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ export type RequestContext = RequestResponsePair & {
*/
search?: string
}
phase: RequestStore['phase']
renderOpts?: WrapperRenderOpts
isHmrRefresh?: boolean
serverComponentsHmrCache?: ServerComponentsHmrCache
Expand Down Expand Up @@ -111,6 +112,7 @@ export const withRequestStore: WithStore<WorkUnitStore, RequestContext> = <
req,
url,
res,
phase,
renderOpts,
isHmrRefresh,
serverComponentsHmrCache,
Expand All @@ -133,6 +135,7 @@ export const withRequestStore: WithStore<WorkUnitStore, RequestContext> = <

const store: RequestStore = {
type: 'request',
phase,
implicitTags: implicitTags ?? [],
// Rather than just using the whole `url` here, we pull the parts we want
// to ensure we don't use parts of the URL that we shouldn't. This also
Expand Down
Loading
Loading