test: unstable_after in a simulated invocation

vercel · Jun 11, 2024 · 4ef0690 · 4ef0690
1 parent 8bac76a
commit 4ef0690
Show file tree

Hide file tree

Showing 10 changed files with 524 additions and 25 deletions.
diff --git a/packages/next/src/server/web/sandbox/context.ts b/packages/next/src/server/web/sandbox/context.ts
@@ -452,6 +452,10 @@ Learn More: https://nextjs.org/docs/messages/edge-dynamic-code-evaluation`),
       context.clearTimeout = (timeout: number) =>
         timeoutsManager.remove(timeout)
 
+      if (process.env.__NEXT_TEST_MODE) {
+        context.__next_outer_globalThis__ = globalThis
+      }
+
       return context
     },
   })

diff --git a/test/e2e/app-dir/next-after-app/app/delay-deep/page.js b/test/e2e/app-dir/next-after-app/app/delay-deep/page.js
@@ -0,0 +1,55 @@
+import { Suspense } from 'react'
+import { unstable_after as after } from 'next/server'
+import { cliLog } from '../../utils/log'
+import { sleep } from '../../utils/sleep'
+
+// don't waste time prerendering, after() will bail out anyway
+export const dynamic = 'force-dynamic'
+
+export default async function Page() {
+  cliLog({ source: '[page] /delay-deep (Page) - render' })
+  return (
+    <Suspense fallback={'Loading...'}>
+      <Inner>Delay</Inner>
+    </Suspense>
+  )
+}
+
+async function Inner({ children }) {
+  cliLog({
+    source: '[page] /delay-deep (Inner) - render, sleeping',
+  })
+
+  await sleep(1000)
+
+  cliLog({
+    source: '[page] /delay-deep (Inner) - render, done sleeping',
+  })
+
+  return (
+    <div>
+      <Suspense fallback="Loading 2...">
+        <Inner2>{children}</Inner2>
+      </Suspense>
+    </div>
+  )
+}
+
+async function Inner2({ children }) {
+  cliLog({
+    source: '[page] /delay-deep (Inner2) - render, sleeping',
+  })
+
+  await sleep(1000)
+
+  cliLog({
+    source: '[page] /delay-deep (Inner2) - render, done sleeping',
+  })
+
+  after(async () => {
+    await sleep(1000)
+    cliLog({ source: '[page] /delay-deep (Inner2) - after' })
+  })
+
+  return <>{children}</>
+}
diff --git a/test/e2e/app-dir/next-after-app/app/layout.js b/test/e2e/app-dir/next-after-app/app/layout.js
@@ -1,7 +1,11 @@
+import { maybeInstallInvocationShutdownHook } from '../utils/simulated-invocation'
+
 // (patched in tests)
 // export const runtime = 'REPLACE_ME'
+// export const dynamic = 'REPLACE_ME'
 
 export default function AppLayout({ children }) {
+  maybeInstallInvocationShutdownHook()
   return (
     <html>
       <head>

diff --git a/test/e2e/app-dir/next-after-app/app/route-streaming/route.js b/test/e2e/app-dir/next-after-app/app/route-streaming/route.js
@@ -0,0 +1,46 @@
+import { unstable_after as after } from 'next/server'
+import { cliLog } from '../../utils/log'
+import { sleep } from '../../utils/sleep'
+import { maybeInstallInvocationShutdownHook } from '../../utils/simulated-invocation'
+
+export const dynamic = 'force-dynamic'
+
+// (patched in tests)
+// export const runtime = 'REPLACE_ME'
+
+export async function GET() {
+  maybeInstallInvocationShutdownHook()
+
+  /** @type {ReadableStream<Uint8Array>} */
+  const result = new ReadableStream({
+    async start(controller) {
+      cliLog({
+        source: '[route handler] /route-streaming - body, sleeping',
+      })
+      await sleep(500)
+      cliLog({
+        source: '[route handler] /route-streaming - body, done sleeping',
+      })
+
+      const encoder = new TextEncoder()
+      for (const chunk of ['one', 'two', 'three']) {
+        await sleep(500)
+        controller.enqueue(encoder.encode(chunk + '\r\n'))
+      }
+
+      after(async () => {
+        await sleep(1000)
+        cliLog({
+          source: '[route handler] /route-streaming - after',
+        })
+      })
+      controller.close()
+    },
+  })
+  return new Response(result, {
+    headers: {
+      'content-type': 'text/plain; charset=utf-8',
+      'transfer-encoding': 'chunked',
+    },
+  })
+}
diff --git a/test/e2e/app-dir/next-after-app/index.test.ts b/test/e2e/app-dir/next-after-app/index.test.ts
@@ -1,5 +1,5 @@
 /* eslint-env jest */
-import { nextTestSetup } from 'e2e-utils'
+import { NextInstance, nextTestSetup } from 'e2e-utils'
 import { retry } from 'next-test-utils'
 import { createProxyServer } from 'next/experimental/testmode/proxy'
 import { outdent } from 'outdent'
@@ -22,25 +22,42 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
     },
   })
 
+  const filesToPatchRuntime = [
+    'app/layout.js',
+    'app/route/route.js',
+    'app/route-streaming/route.js',
+  ]
+  const replaceRuntime = (contents: string, file: string) => {
+    const placeholder = `// export const runtime = 'REPLACE_ME'`
+
+    if (!contents.includes(placeholder)) {
+      throw new Error(`Placeholder "${placeholder}" not found in ${file}`)
+    }
+
+    return contents.replace(
+      placeholder,
+      `export const runtime = '${runtimeValue}'`
+    )
+  }
+
+  const runtimePatches = new Map<
+    string,
+    string | ((contents: string) => string)
+  >(
+    filesToPatchRuntime.map(
+      (file) =>
+        [file, (contents: string) => replaceRuntime(contents, file)] as const
+    )
+  )
+
   {
     const originalContents: Record<string, string> = {}
 
     beforeAll(async () => {
-      const placeholder = `// export const runtime = 'REPLACE_ME'`
-
-      const filesToPatch = ['app/layout.js', 'app/route/route.js']
-
-      for (const file of filesToPatch) {
+      for (const file of filesToPatchRuntime) {
         await next.patchFile(file, (contents) => {
-          if (!contents.includes(placeholder)) {
-            throw new Error(`Placeholder "${placeholder}" not found in ${file}`)
-          }
           originalContents[file] = contents
-
-          return contents.replace(
-            placeholder,
-            `export const runtime = '${runtimeValue}'`
-          )
+          return replaceRuntime(contents, file)
         })
       }
     })
@@ -282,6 +299,7 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
     const { cleanup } = await sandbox(
       next,
       new Map([
+        ...runtimePatches,
         [
           // this needs to be injected as early as possible, before the server tries to read the context
           // (which may be even before we load the page component in dev mode)
@@ -312,6 +330,97 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
     }
   })
 
+  if (!isNextDev) {
+    describe('keeps the invocation alive if after() is called late during streaming', () => {
+      const setup = async () => {
+        const { cleanup } = await patchSandbox(
+          next,
+          new Map<string, string | ((contents: string) => string)>([
+            ...runtimePatches,
+            [
+              'app/layout.js',
+              (contents) => {
+                contents = replaceRuntime(contents, 'app/layout.js')
+
+                contents = contents.replace(
+                  `// export const dynamic = 'REPLACE_ME'`,
+                  `export const dynamic = 'force-dynamic'`
+                )
+
+                return contents
+              },
+            ],
+            [
+              'utils/simulated-invocation.js',
+              (contents) => {
+                return contents.replace(
+                  `const shouldInstallShutdownHook = false`,
+                  `const shouldInstallShutdownHook = true`
+                )
+              },
+            ],
+            [
+              // this needs to be injected as early as possible, before the server tries to read the context
+              // (which may be even before we load the page component in dev mode)
+              'instrumentation.js',
+              outdent`
+                import { injectRequestContext } from './utils/simulated-invocation'
+                export function register() {
+                  injectRequestContext();
+                }
+              `,
+            ],
+          ])
+        )
+
+        return cleanup
+      }
+
+      /* eslint-disable jest/no-standalone-expect */
+      const it_failingForEdge = runtimeValue === 'edge' ? it.failing : it
+
+      it_failingForEdge('during render', async () => {
+        const cleanup = await setup()
+        try {
+          const response = await next.fetch('/delay-deep')
+          expect(response.status).toBe(200)
+          await response.text()
+          await retry(() => {
+            expect(getLogs()).toContainEqual('simulated-invocation :: end')
+          }, 10_000)
+
+          expect(getLogs()).toContainEqual({
+            source: '[page] /delay-deep (Inner2) - after',
+          })
+        } finally {
+          await cleanup()
+        }
+      })
+
+      it_failingForEdge(
+        'in a route handler that streams a response',
+        async () => {
+          const cleanup = await setup()
+          try {
+            const response = await next.fetch('/route-streaming')
+            expect(response.status).toBe(200)
+            await response.text()
+            await retry(() => {
+              expect(getLogs()).toContainEqual('simulated-invocation :: end')
+            }, 10_000)
+
+            expect(getLogs()).toContainEqual({
+              source: '[route handler] /route-streaming - after',
+            })
+          } finally {
+            await cleanup()
+          }
+        }
+      )
+      /* eslint-enable jest/no-standalone-expect */
+    })
+  }
+
   if (isNextDev) {
     // TODO: these are at the end because they destroy the next server.
     // is there a cleaner way to do this without making the tests slower?
@@ -323,12 +432,14 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
           const { session, cleanup } = await sandbox(
             next,
             new Map([
+              ...runtimePatches,
               [
                 'app/static/page.js',
-                (await next.readFile('app/static/page.js')).replace(
-                  `// export const dynamic = 'REPLACE_ME'`,
-                  `export const dynamic = '${dynamicValue}'`
-                ),
+                (contents) =>
+                  contents.replace(
+                    `// export const dynamic = 'REPLACE_ME'`,
+                    `export const dynamic = '${dynamicValue}'`
+                  ),
               ],
             ]),
             '/static'
@@ -350,12 +461,10 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
         const { session, cleanup } = await sandbox(
           next,
           new Map([
+            ...runtimePatches,
             [
               'app/invalid-in-client/page.js',
-              (await next.readFile('app/invalid-in-client/page.js')).replace(
-                `// 'use client'`,
-                `'use client'`
-              ),
+              (contents) => contents.replace(`// 'use client'`, `'use client'`),
             ],
           ]),
           '/invalid-in-client'
@@ -391,3 +500,24 @@ function timeoutPromise(duration: number, message = 'Timeout') {
     )
   )
 }
+
+async function patchSandbox(
+  next: NextInstance,
+  files: Map<string, string | ((contents: string) => string)>
+) {
+  await next.stop()
+  await next.clean()
+
+  for (const [file, newContents] of files) {
+    await next.patchFile(file, newContents)
+  }
+
+  await next.start()
+
+  const cleanup = async () => {
+    await next.stop()
+    await next.clean()
+  }
+
+  return { cleanup }
+}