From 6f43c90b92ee46f907c74fd155501ca82ff7a10c Mon Sep 17 00:00:00 2001
From: Damien Simonin Feugas <damien@vercel.com>
Date: Fri, 21 Oct 2022 14:17:58 +0200
Subject: [PATCH] feat(edge): adds AsyncLocalStorage support to the edge
 function sandbox (#41622)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 📖  Feature

Adds `AsyncLocalStorage` as a global variable to any edge function (middleware, Edge API routes).
Falls back to Node.js' implementation.

## 🧪 How to test

1. `pnpm build`
2. `pnpm testheadless --testPathPattern async-local`
---
 packages/next/server/web/sandbox/context.ts   |   3 +
 .../edge-async-local-storage/index.test.ts    | 126 ++++++++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 test/e2e/edge-async-local-storage/index.test.ts

diff --git a/packages/next/server/web/sandbox/context.ts b/packages/next/server/web/sandbox/context.ts
index abdfb4252bd3..943b8e2076ff 100644
--- a/packages/next/server/web/sandbox/context.ts
+++ b/packages/next/server/web/sandbox/context.ts
@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from 'async_hooks'
 import type { AssetBinding } from '../../../build/webpack/loaders/get-module-build-info'
 import {
   decorateServerError,
@@ -286,6 +287,8 @@ Learn More: https://nextjs.org/docs/messages/edge-dynamic-code-evaluation`),
 
       Object.assign(context, wasm)
 
+      context.AsyncLocalStorage = AsyncLocalStorage
+
       return context
     },
   })
diff --git a/test/e2e/edge-async-local-storage/index.test.ts b/test/e2e/edge-async-local-storage/index.test.ts
new file mode 100644
index 000000000000..39c2798f1822
--- /dev/null
+++ b/test/e2e/edge-async-local-storage/index.test.ts
@@ -0,0 +1,126 @@
+/* eslint-disable jest/valid-expect-in-promise */
+import { createNext } from 'e2e-utils'
+import { NextInstance } from 'test/lib/next-modes/base'
+import { fetchViaHTTP } from 'next-test-utils'
+
+describe('edge api can use async local storage', () => {
+  let next: NextInstance
+
+  const cases = [
+    {
+      title: 'a single instance',
+      code: `
+        export const config = { runtime: 'experimental-edge' }
+        const storage = new AsyncLocalStorage()
+  
+        export default async function handler(request) {
+          const id = request.headers.get('req-id')
+          return storage.run({ id }, async () => {
+            await getSomeData()
+            return Response.json(storage.getStore())
+          })
+        }
+  
+        async function getSomeData() {
+          try {
+            const response = await fetch('https://example.vercel.sh')
+            await response.text()
+          } finally {
+            return true
+          }
+        }
+      `,
+      expectResponse: (response, id) =>
+        expect(response).toMatchObject({ status: 200, json: { id } }),
+    },
+    {
+      title: 'multiple instances',
+      code: `
+        export const config = { runtime: 'experimental-edge' }
+        const topStorage = new AsyncLocalStorage()
+  
+        export default async function handler(request) {
+          const id = request.headers.get('req-id')
+          return topStorage.run({ id }, async () => {
+            const nested = await getSomeData(id)
+            return Response.json({ ...nested, ...topStorage.getStore() })
+          })
+        }
+  
+        async function getSomeData(id) {
+          const nestedStorage = new AsyncLocalStorage()
+          return nestedStorage.run('nested-' + id, async () => {
+            try {
+              const response = await fetch('https://example.vercel.sh')
+              await response.text()
+            } finally {
+              return { nestedId: nestedStorage.getStore() }
+            }
+          })
+        }
+      `,
+      expectResponse: (response, id) =>
+        expect(response).toMatchObject({
+          status: 200,
+          json: { id: id, nestedId: `nested-${id}` },
+        }),
+    },
+  ]
+
+  afterEach(() => next.destroy())
+
+  it.each(cases)(
+    'cans use $title per request',
+    async ({ code, expectResponse }) => {
+      next = await createNext({
+        files: {
+          'pages/index.js': `
+            export default function () { return <div>Hello, world!</div> }
+          `,
+          'pages/api/async.js': code,
+        },
+      })
+      const ids = Array.from({ length: 100 }, (_, i) => `req-${i}`)
+
+      const responses = await Promise.all(
+        ids.map((id) =>
+          fetchViaHTTP(
+            next.url,
+            '/api/async',
+            {},
+            { headers: { 'req-id': id } }
+          ).then((response) =>
+            response.headers.get('content-type')?.startsWith('application/json')
+              ? response.json().then((json) => ({
+                  status: response.status,
+                  json,
+                  text: null,
+                }))
+              : response.text().then((text) => ({
+                  status: response.status,
+                  json: null,
+                  text,
+                }))
+          )
+        )
+      )
+      const rankById = new Map(ids.map((id, rank) => [id, rank]))
+
+      const errors: Error[] = []
+      for (const [rank, response] of responses.entries()) {
+        try {
+          expectResponse(response, ids[rank])
+        } catch (error) {
+          const received = response.json?.id
+          console.log(
+            `response #${rank} has id from request #${rankById.get(received)}`
+          )
+          errors.push(error as Error)
+        }
+      }
+      if (errors.length) {
+        throw errors[0]
+      }
+    }
+  )
+})