feat(koa): defer AsyncLocalStorage creation for V8 startup snapshot support

packages/koa/src/application.ts

@@ -3,6 +3,7 @@ import Emitter from 'node:events';
 import http, { type IncomingMessage, type ServerResponse } from 'node:http';
 import Stream from 'node:stream';
 import util, { debuglog } from 'node:util';
+import v8 from 'node:v8';
 
 import { getAsyncLocalStorage } from 'gals';
 import { HttpError } from 'http-errors';
@@ -48,7 +49,7 @@ export class Application extends Emitter {
   maxIpsCount: number;
   protected _keys?: string[];
   middleware: MiddlewareFunc<Context>[];
-  ctxStorage: AsyncLocalStorage<Context>;
+  ctxStorage: AsyncLocalStorage<Context> | null;
   silent: boolean;
   ContextClass: ProtoImplClass<Context>;
   context: AnyProto;
@@ -88,7 +89,14 @@ export class Application extends Emitter {
       this._keys = options.keys;
     }
     this.middleware = [];
-    this.ctxStorage = getAsyncLocalStorage();
+    if (v8.startupSnapshot?.isBuildingSnapshot?.()) {
+      this.ctxStorage = null;
+      v8.startupSnapshot.addDeserializeCallback((app: Application) => {
+        app.ctxStorage = getAsyncLocalStorage();
+      }, this);
+    } else {
+      this.ctxStorage = getAsyncLocalStorage();
+    }
     this.silent = false;
     this.ContextClass = class ApplicationContext extends Context {} as ProtoImplClass<Context>;
     this.context = this.ContextClass.prototype;
@@ -184,9 +192,12 @@ export class Application extends Emitter {
 
     const handleRequest = (req: IncomingMessage, res: ServerResponse) => {
       const ctx = this.createContext(req, res);
-      return this.ctxStorage.run(ctx, async () => {
-        return await this.handleRequest(ctx, fn);
-      });
+      if (this.ctxStorage) {
+        return this.ctxStorage.run(ctx, async () => {
+          return await this.handleRequest(ctx, fn);
+        });
+      }
+      return this.handleRequest(ctx, fn);
     };
 
     return handleRequest;
@@ -196,7 +207,7 @@ export class Application extends Emitter {
    * return current context from async local storage
    */
   get currentContext(): Context | undefined {
-    return this.ctxStorage.getStore();
+    return this.ctxStorage?.getStore();
   }
 
   /**

packages/koa/test/application/snapshot.test.ts

@@ -0,0 +1,94 @@
+import assert from 'node:assert/strict';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import v8 from 'node:v8';
+
+import { request } from '@eggjs/supertest';
+import { afterEach, describe, it } from 'vitest';
+
+import Koa from '../../src/index.ts';
+
+describe('v8 startup snapshot', () => {
+  const originalStartupSnapshot = v8.startupSnapshot;
+
+  afterEach(() => {
+    (v8 as Record<string, unknown>).startupSnapshot = originalStartupSnapshot;
+  });
+
+  it('should defer AsyncLocalStorage creation when building snapshot', () => {
+    let deserializeCallback: { cb: (data: unknown) => void; data: unknown } | undefined;
+    (v8 as Record<string, unknown>).startupSnapshot = {
+      isBuildingSnapshot: () => true,
+      addDeserializeCallback: (cb: (data: unknown) => void, data: unknown) => {
+        deserializeCallback = { cb, data };
+      },
+    };
+
+    const app = new Koa();
+    assert.strictEqual(app.ctxStorage, null);
+    assert.ok(deserializeCallback, 'deserialize callback should be registered');
+
+    // simulate snapshot deserialization
+    deserializeCallback.cb(deserializeCallback.data);
+    assert.ok(app.ctxStorage! instanceof AsyncLocalStorage);
+  });
+
+  it('should return undefined for currentContext when ctxStorage is null', () => {
+    (v8 as Record<string, unknown>).startupSnapshot = {
+      isBuildingSnapshot: () => true,
+      addDeserializeCallback: () => {},
+    };
+
+    const app = new Koa();
+    assert.strictEqual(app.ctxStorage, null);
+    assert.strictEqual(app.currentContext, undefined);
+  });
+
+  it('should work normally after deserialization', async () => {
+    let deserializeCallback: { cb: (data: unknown) => void; data: unknown } | undefined;
+    (v8 as Record<string, unknown>).startupSnapshot = {
+      isBuildingSnapshot: () => true,
+      addDeserializeCallback: (cb: (data: unknown) => void, data: unknown) => {
+        deserializeCallback = { cb, data };
+      },
+    };
+
+    const app = new Koa();
+
+    // simulate snapshot deserialization
+    deserializeCallback!.cb(deserializeCallback!.data);
+
+    app.use(async (ctx) => {
+      assert.equal(ctx, app.currentContext);
+      ctx.body = 'ok';
+    });
+
+    await request(app.callback()).get('/').expect('ok');
+    assert.strictEqual(app.currentContext, undefined);
+  });
+
+  it('should not defer when not building snapshot', () => {
+    (v8 as Record<string, unknown>).startupSnapshot = {
+      isBuildingSnapshot: () => false,
+    };
+
+    const app = new Koa();
+    assert.ok(app.ctxStorage! instanceof AsyncLocalStorage);
+  });
+
+  it('should handle callback without ctxStorage during snapshot build', async () => {
+    (v8 as Record<string, unknown>).startupSnapshot = {
+      isBuildingSnapshot: () => true,
+      addDeserializeCallback: () => {},
+    };
+
+    const app = new Koa();
+    assert.strictEqual(app.ctxStorage, null);
+
+    app.use(async (ctx) => {
+      assert.strictEqual(app.currentContext, undefined);
+      ctx.body = 'ok';
+    });
+
+    await request(app.callback()).get('/').expect('ok');
+  });
+});