From e90fa80bc2d2b2b71eccba43f75096b968c280fc Mon Sep 17 00:00:00 2001
From: Taku Amano <taku@taaas.jp>
Date: Sun, 5 May 2024 05:45:06 +0900
Subject: [PATCH] feat: use AsyncLocalStorageFallback when AsyncLocalStorage is
 not available

---
 src/server/components/has-islands.tsx |  4 +--
 src/server/context-storage.ts         | 40 +++++++++++++++++++++++++--
 src/server/server.ts                  |  5 ++--
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/src/server/components/has-islands.tsx b/src/server/components/has-islands.tsx
index dbfb21a..65d8d69 100644
--- a/src/server/components/has-islands.tsx
+++ b/src/server/components/has-islands.tsx
@@ -1,9 +1,9 @@
 import { IMPORTING_ISLANDS_ID } from '../../constants.js'
-import { contextStorage } from '../context-storage.js'
+import { getStorageSync } from '../context-storage.js'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export const HasIslands = ({ children }: { children: any }): any => {
-  const c = contextStorage.getStore()
+  const c = getStorageSync().getStore()
   if (!c) {
     throw new Error('No context found')
   }
diff --git a/src/server/context-storage.ts b/src/server/context-storage.ts
index 4e958cf..81b7bf7 100644
--- a/src/server/context-storage.ts
+++ b/src/server/context-storage.ts
@@ -1,3 +1,39 @@
-import { AsyncLocalStorage } from 'node:async_hooks'
+import type { AsyncLocalStorage } from 'node:async_hooks'
 import type { Context } from 'hono'
-export const contextStorage = new AsyncLocalStorage<Context>()
+
+class AsyncLocalStorageFallback<T> {
+  store: T | undefined
+  getStore() {
+    return this.store
+  }
+  async run(store: T, callback: () => Promise<void>) {
+    if (this.store) {
+      throw new Error('Already running')
+    }
+
+    this.store = store
+    await callback()
+    this.store = undefined
+  }
+}
+
+let contextStorage: AsyncLocalStorage<Context> | undefined
+export const getStorage = async () => {
+  if (!contextStorage) {
+    try {
+      const { AsyncLocalStorage } = await import('node:async_hooks')
+      contextStorage = new AsyncLocalStorage<Context>()
+    } catch (e) {
+      contextStorage =
+        new AsyncLocalStorageFallback<Context>() as unknown as AsyncLocalStorage<Context>
+    }
+  }
+  return contextStorage
+}
+
+export const getStorageSync = () => {
+  if (!contextStorage) {
+    throw new Error('No context found')
+  }
+  return contextStorage
+}
diff --git a/src/server/server.ts b/src/server/server.ts
index 480ba57..3ecb433 100644
--- a/src/server/server.ts
+++ b/src/server/server.ts
@@ -10,7 +10,7 @@ import {
   listByDirectory,
   sortDirectoriesByDepth,
 } from '../utils/file.js'
-import { contextStorage } from './context-storage.js'
+import { getStorage } from './context-storage.js'
 
 const NOTFOUND_FILENAME = '_404.tsx'
 const ERROR_FILENAME = '_error.tsx'
@@ -64,8 +64,9 @@ export const createApp = <E extends Env>(options: BaseServerOptions<E>): Hono<E>
   const trailingSlash = options.trailingSlash ?? false
 
   // Share context by AsyncLocalStorage
+  const getStoragePromise = getStorage()
   app.use(async function ShareContext(c, next) {
-    await contextStorage.run(c, () => next())
+    await (await getStoragePromise).run(c, () => next())
   })
 
   if (options.init) {