diff --git a/packages/koa/src/application.ts b/packages/koa/src/application.ts index 8d746b83b2..4c748439a5 100644 --- a/packages/koa/src/application.ts +++ b/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[]; - ctxStorage: AsyncLocalStorage; + ctxStorage: AsyncLocalStorage | null; silent: boolean; ContextClass: ProtoImplClass; 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; 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(); } /** diff --git a/packages/koa/test/application/snapshot.test.ts b/packages/koa/test/application/snapshot.test.ts new file mode 100644 index 0000000000..f8f2a436d6 --- /dev/null +++ b/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).startupSnapshot = originalStartupSnapshot; + }); + + it('should defer AsyncLocalStorage creation when building snapshot', () => { + let deserializeCallback: { cb: (data: unknown) => void; data: unknown } | undefined; + (v8 as Record).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).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).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).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).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'); + }); +});