Skip to content

Commit

Permalink
fix: pass waitUntil from NextRequestHint to wrapped NextWebServer
Browse files Browse the repository at this point in the history
  • Loading branch information
lubieowoce committed Aug 8, 2024
1 parent 3137c25 commit dcab3e5
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 3 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
import { createAsyncLocalStorage } from './async-local-storage'
import type { LifecycleAsyncStorage } from './lifecycle-async-storage.external'

export const _lifecycleAsyncStorage: LifecycleAsyncStorage =
createAsyncLocalStorage()
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import type { AsyncLocalStorage } from 'async_hooks'
// Share the instance module in the next-shared layer
import { _lifecycleAsyncStorage as lifecycleAsyncStorage } from './lifecycle-async-storage-instance' with { 'turbopack-transition': 'next-shared' }

export interface LifecycleStore {
readonly waitUntil: ((promise: Promise<any>) => void) | undefined
}

export type LifecycleAsyncStorage = AsyncLocalStorage<LifecycleStore>

export { lifecycleAsyncStorage }
6 changes: 6 additions & 0 deletions packages/next/src/server/base-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,7 @@ import {
} from './after/builtin-request-context'
import { ENCODED_TAGS } from './stream-utils/encodedTags'
import { NextRequestHint } from './web/adapter'
import { lifecycleAsyncStorage } from '../client/components/lifecycle-async-storage.external'

export type FindComponentsResult = {
components: LoadComponentsReturnType
Expand Down Expand Up @@ -1736,6 +1737,11 @@ export default abstract class Server<
}

protected getWaitUntil(): WaitUntil | undefined {
const lifecycleStore = lifecycleAsyncStorage.getStore()
if (lifecycleStore) {
return lifecycleStore.waitUntil
}

const builtinRequestContext = getBuiltinRequestContext()
if (builtinRequestContext) {
// the platform provided a request context.
Expand Down
27 changes: 24 additions & 3 deletions packages/next/src/server/web/adapter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import type { TextMapGetter } from 'next/dist/compiled/@opentelemetry/api'
import { MiddlewareSpan } from '../lib/trace/constants'
import { CloseController } from './web-on-close'
import { getEdgePreviewProps } from './get-edge-preview-props'
import { lifecycleAsyncStorage } from '../../client/components/lifecycle-async-storage.external'

export class NextRequestHint extends NextRequest {
sourcePage: string
Expand Down Expand Up @@ -207,13 +208,14 @@ export async function adapter(
const isMiddleware =
params.page === '/middleware' || params.page === '/src/middleware'

const isAfterEnabled =
params.request.nextConfig?.experimental?.after ??
!!process.env.__NEXT_AFTER

if (isMiddleware) {
// 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 closeController: CloseController | undefined = undefined
Expand Down Expand Up @@ -271,6 +273,25 @@ export async function adapter(
}
)
}

if (isAfterEnabled) {
// NOTE:
// Currently, `adapter` is expected to return promises passed to `waitUntil`
// as part of its result (i.e. a FetchEventResult).
// Because of this, we override any outer contexts that might provide a real `waitUntil`,
// and provide the `waitUntil` from the NextFetchEvent instead so that we can collect those promises.
// This is not ideal, but until we change this calling convention, it's the least surprising thing to do.
//
// Notably, the only case that currently cares about this ALS is Edge SSR
// (i.e. a handler created via `build/webpack/loaders/next-edge-ssr-loader/render.ts`)
// Other types of handlers will grab the waitUntil from the passed FetchEvent,
// but NextWebServer currently has no interface that'd allow for that.
return lifecycleAsyncStorage.run(
{ waitUntil: event.waitUntil.bind(event) },
() => params.handler(request, event)
)
}

return params.handler(request, event)
})

Expand Down

0 comments on commit dcab3e5

Please sign in to comment.