Skip to content

Commit ece7b7e

Browse files
killaguclaude
andauthored
feat(koa): defer AsyncLocalStorage creation for V8 startup snapshot support (#5851)
## Summary - Defer `AsyncLocalStorage` creation during V8 snapshot building via `v8.startupSnapshot.isBuildingSnapshot()` - Make `ctxStorage` nullable (`AsyncLocalStorage<Context> | null`) with null-safe access in `callback()` and `currentContext` - Register deserialize callback to restore ALS when snapshot is restored This is **PR1 of 6** in the V8 startup snapshot series. Independent, no dependencies. ## Changes - `packages/koa/src/application.ts` — Defer ALS creation, null-safe ctxStorage access - `packages/koa/test/application/snapshot.test.ts` — 5 new tests covering defer, restore, null safety ## Test plan - [x] All 70 koa test files pass (381 tests, 1 skipped) - [x] 5 new snapshot-specific tests - [x] oxlint --type-aware clean 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * App initialization now safely handles absent context storage during special runtime startup modes, preventing crashes and improving stability. * Context accessor now returns undefined when storage is not available, avoiding errors in early startup. * **Tests** * Added tests verifying initialization, context exposure, and request handling both during startup-snapshot mode and normal startup. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 5c459e4 commit ece7b7e

File tree

2 files changed

+111
-6
lines changed

2 files changed

+111
-6
lines changed

packages/koa/src/application.ts

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import Emitter from 'node:events';
33
import http, { type IncomingMessage, type ServerResponse } from 'node:http';
44
import Stream from 'node:stream';
55
import util, { debuglog } from 'node:util';
6+
import v8 from 'node:v8';
67

78
import { getAsyncLocalStorage } from 'gals';
89
import { HttpError } from 'http-errors';
@@ -48,7 +49,7 @@ export class Application extends Emitter {
4849
maxIpsCount: number;
4950
protected _keys?: string[];
5051
middleware: MiddlewareFunc<Context>[];
51-
ctxStorage: AsyncLocalStorage<Context>;
52+
ctxStorage: AsyncLocalStorage<Context> | null;
5253
silent: boolean;
5354
ContextClass: ProtoImplClass<Context>;
5455
context: AnyProto;
@@ -88,7 +89,14 @@ export class Application extends Emitter {
8889
this._keys = options.keys;
8990
}
9091
this.middleware = [];
91-
this.ctxStorage = getAsyncLocalStorage();
92+
if (v8.startupSnapshot?.isBuildingSnapshot?.()) {
93+
this.ctxStorage = null;
94+
v8.startupSnapshot.addDeserializeCallback((app: Application) => {
95+
app.ctxStorage = getAsyncLocalStorage();
96+
}, this);
97+
} else {
98+
this.ctxStorage = getAsyncLocalStorage();
99+
}
92100
this.silent = false;
93101
this.ContextClass = class ApplicationContext extends Context {} as ProtoImplClass<Context>;
94102
this.context = this.ContextClass.prototype;
@@ -184,9 +192,12 @@ export class Application extends Emitter {
184192

185193
const handleRequest = (req: IncomingMessage, res: ServerResponse) => {
186194
const ctx = this.createContext(req, res);
187-
return this.ctxStorage.run(ctx, async () => {
188-
return await this.handleRequest(ctx, fn);
189-
});
195+
if (this.ctxStorage) {
196+
return this.ctxStorage.run(ctx, async () => {
197+
return await this.handleRequest(ctx, fn);
198+
});
199+
}
200+
return this.handleRequest(ctx, fn);
190201
};
191202

192203
return handleRequest;
@@ -196,7 +207,7 @@ export class Application extends Emitter {
196207
* return current context from async local storage
197208
*/
198209
get currentContext(): Context | undefined {
199-
return this.ctxStorage.getStore();
210+
return this.ctxStorage?.getStore();
200211
}
201212

202213
/**
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
import assert from 'node:assert/strict';
2+
import { AsyncLocalStorage } from 'node:async_hooks';
3+
import v8 from 'node:v8';
4+
5+
import { request } from '@eggjs/supertest';
6+
import { afterEach, describe, it } from 'vitest';
7+
8+
import Koa from '../../src/index.ts';
9+
10+
describe('v8 startup snapshot', () => {
11+
const originalStartupSnapshot = v8.startupSnapshot;
12+
13+
afterEach(() => {
14+
(v8 as Record<string, unknown>).startupSnapshot = originalStartupSnapshot;
15+
});
16+
17+
it('should defer AsyncLocalStorage creation when building snapshot', () => {
18+
let deserializeCallback: { cb: (data: unknown) => void; data: unknown } | undefined;
19+
(v8 as Record<string, unknown>).startupSnapshot = {
20+
isBuildingSnapshot: () => true,
21+
addDeserializeCallback: (cb: (data: unknown) => void, data: unknown) => {
22+
deserializeCallback = { cb, data };
23+
},
24+
};
25+
26+
const app = new Koa();
27+
assert.strictEqual(app.ctxStorage, null);
28+
assert.ok(deserializeCallback, 'deserialize callback should be registered');
29+
30+
// simulate snapshot deserialization
31+
deserializeCallback.cb(deserializeCallback.data);
32+
assert.ok(app.ctxStorage! instanceof AsyncLocalStorage);
33+
});
34+
35+
it('should return undefined for currentContext when ctxStorage is null', () => {
36+
(v8 as Record<string, unknown>).startupSnapshot = {
37+
isBuildingSnapshot: () => true,
38+
addDeserializeCallback: () => {},
39+
};
40+
41+
const app = new Koa();
42+
assert.strictEqual(app.ctxStorage, null);
43+
assert.strictEqual(app.currentContext, undefined);
44+
});
45+
46+
it('should work normally after deserialization', async () => {
47+
let deserializeCallback: { cb: (data: unknown) => void; data: unknown } | undefined;
48+
(v8 as Record<string, unknown>).startupSnapshot = {
49+
isBuildingSnapshot: () => true,
50+
addDeserializeCallback: (cb: (data: unknown) => void, data: unknown) => {
51+
deserializeCallback = { cb, data };
52+
},
53+
};
54+
55+
const app = new Koa();
56+
57+
// simulate snapshot deserialization
58+
deserializeCallback!.cb(deserializeCallback!.data);
59+
60+
app.use(async (ctx) => {
61+
assert.equal(ctx, app.currentContext);
62+
ctx.body = 'ok';
63+
});
64+
65+
await request(app.callback()).get('/').expect('ok');
66+
assert.strictEqual(app.currentContext, undefined);
67+
});
68+
69+
it('should not defer when not building snapshot', () => {
70+
(v8 as Record<string, unknown>).startupSnapshot = {
71+
isBuildingSnapshot: () => false,
72+
};
73+
74+
const app = new Koa();
75+
assert.ok(app.ctxStorage! instanceof AsyncLocalStorage);
76+
});
77+
78+
it('should handle callback without ctxStorage during snapshot build', async () => {
79+
(v8 as Record<string, unknown>).startupSnapshot = {
80+
isBuildingSnapshot: () => true,
81+
addDeserializeCallback: () => {},
82+
};
83+
84+
const app = new Koa();
85+
assert.strictEqual(app.ctxStorage, null);
86+
87+
app.use(async (ctx) => {
88+
assert.strictEqual(app.currentContext, undefined);
89+
ctx.body = 'ok';
90+
});
91+
92+
await request(app.callback()).get('/').expect('ok');
93+
});
94+
});

0 commit comments

Comments
 (0)