Create server module map proxy only once (#86750)

Small follow-up to #86591.

We don't need to recreate the server module map proxy every time we set
the manifests singleton. Instead, we can just update the underlying
server actions manifest, and the server module map can read the current
manifest from the singleton.

Author: unstubbable

packages/next/src/server/app-render/action-handler.ts

@@ -49,7 +49,10 @@ import { warn } from '../../build/output/log'
 import { RequestCookies, ResponseCookies } from '../web/spec-extension/cookies'
 import { HeadersAdapter } from '../web/spec-extension/adapters/headers'
 import { fromNodeOutgoingHttpHeaders } from '../web/utils'
-import { selectWorkerForForwarding, type ServerModuleMap } from './action-utils'
+import {
+  selectWorkerForForwarding,
+  type ServerModuleMap,
+} from './manifests-singleton'
 import { isNodeNextRequest, isWebNextRequest } from '../base-http/helpers'
 import { RedirectStatusCode } from '../../client/components/redirect-status-code'
 import { synchronizeMutableCookies } from '../async-storage/request-store'

packages/next/src/server/app-render/action-utils.ts

@@ -3,9 +3,19 @@ import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 import { InvariantError } from '../../shared/lib/invariant-error'
 import { normalizeAppPath } from '../../shared/lib/router/utils/app-paths'
-import { createServerModuleMap, type ServerModuleMap } from './action-utils'
+import { pathHasPrefix } from '../../shared/lib/router/utils/path-has-prefix'
+import { removePathPrefix } from '../../shared/lib/router/utils/remove-path-prefix'
 import { workAsyncStorage } from './work-async-storage.external'
 
+export interface ServerModuleMap {
+  readonly [name: string]: {
+    readonly id: string | number
+    readonly name: string
+    readonly chunks: Readonly<Array<string>> // currently not used
+    readonly async?: boolean
+  }
+}
+
 // This is a global singleton that is, among other things, also used to
 // encode/decode bound args of server function closures. This can't be using a
 // AsyncLocalStorage as it might happen at the module level.
@@ -162,6 +172,105 @@ function createProxiedClientReferenceManifest(
   ) as DeepReadonly<ClientReferenceManifest>
 }
 
+/**
+ * This function creates a Flight-acceptable server module map proxy from our
+ * Server Reference Manifest similar to our client module map. This is because
+ * our manifest contains a lot of internal Next.js data that are relevant to the
+ * runtime, workers, etc. that React doesn't need to know.
+ */
+function createServerModuleMap(): ServerModuleMap {
+  return new Proxy(
+    {},
+    {
+      get: (_, id: string) => {
+        const workers =
+          getServerActionsManifest()[
+            process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
+          ]?.[id]?.workers
+
+        if (!workers) {
+          return undefined
+        }
+
+        const workStore = workAsyncStorage.getStore()
+
+        let workerEntry:
+          | { moduleId: string | number; async: boolean }
+          | undefined
+
+        if (workStore) {
+          workerEntry = workers[normalizeWorkerPageName(workStore.page)]
+        } else {
+          // If there's no work store defined, we can assume that a server
+          // module map is needed during module evaluation, e.g. to create a
+          // server action using a higher-order function. Therefore it should be
+          // safe to return any entry from the manifest that matches the action
+          // ID. They all refer to the same module ID, which must also exist in
+          // the current page bundle. TODO: This is currently not guaranteed in
+          // Turbopack, and needs to be fixed.
+          workerEntry = Object.values(workers).at(0)
+        }
+
+        if (!workerEntry) {
+          return undefined
+        }
+
+        const { moduleId, async } = workerEntry
+
+        return { id: moduleId, name: id, chunks: [], async }
+      },
+    }
+  )
+}
+
+/**
+ * The flight entry loader keys actions by bundlePath. bundlePath corresponds
+ * with the relative path (including 'app') to the page entrypoint.
+ */
+function normalizeWorkerPageName(pageName: string) {
+  if (pathHasPrefix(pageName, 'app')) {
+    return pageName
+  }
+
+  return 'app' + pageName
+}
+
+/**
+ * Converts a bundlePath (relative path to the entrypoint) to a routable page
+ * name.
+ */
+function denormalizeWorkerPageName(bundlePath: string) {
+  return normalizeAppPath(removePathPrefix(bundlePath, 'app'))
+}
+
+/**
+ * Checks if the requested action has a worker for the current page.
+ * If not, it returns the first worker that has a handler for the action.
+ */
+export function selectWorkerForForwarding(
+  actionId: string,
+  pageName: string
+): string | undefined {
+  const serverActionsManifest = getServerActionsManifest()
+  const workers =
+    serverActionsManifest[
+      process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
+    ][actionId]?.workers
+
+  // There are no workers to handle this action, nothing to forward to.
+  if (!workers) {
+    return
+  }
+
+  // If there is an entry for the current page, we don't need to forward.
+  if (workers[normalizeWorkerPageName(pageName)]) {
+    return
+  }
+
+  // Otherwise, grab the first worker that has a handler for this action id.
+  return denormalizeWorkerPageName(Object.keys(workers)[0])
+}
+
 export function setManifestsSingleton({
   page,
   clientReferenceManifest,
@@ -180,10 +289,6 @@ export function setManifestsSingleton({
     )
 
     existingSingleton.serverActionsManifest = serverActionsManifest
-
-    existingSingleton.serverModuleMap = createServerModuleMap({
-      serverActionsManifest,
-    })
   } else {
     const clientReferenceManifestsPerRoute = new Map<
       string,
@@ -194,15 +299,11 @@ export function setManifestsSingleton({
       clientReferenceManifestsPerRoute
     )
 
-    const serverModuleMap = createServerModuleMap({
-      serverActionsManifest,
-    })
-
     globalThisWithManifests[MANIFESTS_SINGLETON] = {
       clientReferenceManifestsPerRoute,
       proxiedClientReferenceManifest,
       serverActionsManifest,
-      serverModuleMap,
+      serverModuleMap: createServerModuleMap(),
     }
   }
 }