Skip to content

Commit 8554929

Browse files
SyMindSyMind
committed
perf: JsHookUsageTracker
1 parent 2fa739a commit 8554929

File tree

3 files changed

+170
-67
lines changed

3 files changed

+170
-67
lines changed

‎crates/node_binding/napi-binding.d.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -645,6 +645,8 @@ export const EXPECTED_RSPACK_CORE_VERSION: string
645645

646646
export declare function formatDiagnostic(diagnostic: JsDiagnostic): ExternalObject<'Diagnostic'>
647647

648+
export declare function getRegisterJsTapScopeKinds(): RegisterJsTapScopeKinds
649+
648650
export interface JsAddingRuntimeModule {
649651
name: string
650652
generator: () => String
@@ -3190,6 +3192,11 @@ export interface RegisterJsTaps {
31903192
registerRsdoctorPluginAssetsTaps: (stages: Array<number>) => Array<{ function: ((arg: JsRsdoctorAssetPatch) => Promise<boolean | undefined>); stage: number; }>
31913193
}
31923194

3195+
export interface RegisterJsTapScopeKinds {
3196+
compiler: Array<RegisterJsTapKind>
3197+
compilation: Array<RegisterJsTapKind>
3198+
}
3199+
31933200
export interface ResolveResult {
31943201
path?: string
31953202
error?: string

‎crates/rspack_binding_api/src/plugins/interceptor.rs‎

Lines changed: 92 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -196,42 +196,43 @@ pub(super) struct HookUsageChecker {
196196
struct HookUsageCheckerInner {
197197
_buffer: Buffer,
198198
bytes: NonNull<AtomicU8>,
199-
len: usize,
200199
}
201200

202201
unsafe impl Send for HookUsageCheckerInner {}
203202
unsafe impl Sync for HookUsageCheckerInner {}
204203

205204
impl HookUsageChecker {
206205
pub(super) fn new(buffer: Buffer) -> Self {
206+
assert_eq!(
207+
buffer.len(),
208+
HOOK_USAGE_BUFFER_BYTE_LENGTH,
209+
"hook_usage_buffer length mismatch",
210+
);
207211
Self {
208212
inner: {
209-
let len = buffer.len();
210213
#[allow(clippy::unwrap_used)]
211214
let bytes = NonNull::new(buffer.as_ref().as_ptr() as *mut AtomicU8).unwrap();
212215
Arc::new(HookUsageCheckerInner {
213216
_buffer: buffer,
214217
bytes,
215-
len,
216218
})
217219
},
218220
}
219221
}
220222

221-
pub(super) fn is_used(&self, kind: &RegisterJsTapKind) -> bool {
223+
#[inline(always)]
224+
pub(super) fn is_used(&self, kind: RegisterJsTapKind) -> bool {
222225
let inner = &self.inner;
223-
let bit_index = *kind as usize;
226+
let bit_index = kind as usize;
224227
let byte_index = bit_index >> 3;
225228
let bit_mask = 1 << (bit_index & 7);
226-
if byte_index >= inner.len {
227-
return false;
228-
}
229229
let byte = unsafe { &*inner.bytes.as_ptr().add(byte_index) }.load(Ordering::Acquire);
230230
byte & bit_mask != 0
231231
}
232232
}
233233

234-
fn should_skip_register(hook_usage_checker: &HookUsageChecker, kind: &RegisterJsTapKind) -> bool {
234+
#[inline(always)]
235+
fn should_skip_register(hook_usage_checker: &HookUsageChecker, kind: RegisterJsTapKind) -> bool {
235236
!hook_usage_checker.is_used(kind)
236237
}
237238

@@ -399,7 +400,7 @@ macro_rules! define_register {
399400
&self,
400401
hook: &$tap_hook,
401402
) -> rspack_error::Result<Vec<<$tap_hook as Hook>::Tap>> {
402-
if should_skip_register(&self.inner.hook_usage_checker, &$kind) {
403+
if should_skip_register(&self.inner.hook_usage_checker, $kind) {
403404
return Ok(Vec::new());
404405
}
405406
let js_taps = self.inner.call_register(hook).await?;
@@ -414,6 +415,7 @@ macro_rules! define_register {
414415
}
415416

416417
#[napi]
418+
#[repr(u8)]
417419
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
418420
pub enum RegisterJsTapKind {
419421
CompilerThisCompilation,
@@ -469,6 +471,80 @@ pub enum RegisterJsTapKind {
469471
RsdoctorPluginAssets,
470472
}
471473

474+
const HOOK_USAGE_BUFFER_BYTE_LENGTH: usize =
475+
((RegisterJsTapKind::RsdoctorPluginAssets as usize) >> 3) + 1;
476+
477+
const COMPILER_HOOKS: &[RegisterJsTapKind] = &[
478+
RegisterJsTapKind::CompilerThisCompilation,
479+
RegisterJsTapKind::CompilerCompilation,
480+
RegisterJsTapKind::CompilerMake,
481+
RegisterJsTapKind::CompilerFinishMake,
482+
RegisterJsTapKind::CompilerShouldEmit,
483+
RegisterJsTapKind::CompilerEmit,
484+
RegisterJsTapKind::CompilerAfterEmit,
485+
RegisterJsTapKind::CompilerAssetEmitted,
486+
];
487+
488+
const COMPILATION_HOOKS: &[RegisterJsTapKind] = &[
489+
RegisterJsTapKind::CompilationBuildModule,
490+
RegisterJsTapKind::CompilationStillValidModule,
491+
RegisterJsTapKind::CompilationSucceedModule,
492+
RegisterJsTapKind::CompilationExecuteModule,
493+
RegisterJsTapKind::CompilationFinishModules,
494+
RegisterJsTapKind::CompilationOptimizeModules,
495+
RegisterJsTapKind::CompilationAfterOptimizeModules,
496+
RegisterJsTapKind::CompilationOptimizeTree,
497+
RegisterJsTapKind::CompilationOptimizeChunkModules,
498+
RegisterJsTapKind::CompilationBeforeModuleIds,
499+
RegisterJsTapKind::CompilationAdditionalTreeRuntimeRequirements,
500+
RegisterJsTapKind::CompilationRuntimeRequirementInTree,
501+
RegisterJsTapKind::CompilationRuntimeModule,
502+
RegisterJsTapKind::CompilationChunkHash,
503+
RegisterJsTapKind::CompilationChunkAsset,
504+
RegisterJsTapKind::CompilationProcessAssets,
505+
RegisterJsTapKind::CompilationAfterProcessAssets,
506+
RegisterJsTapKind::CompilationSeal,
507+
RegisterJsTapKind::CompilationAfterSeal,
508+
RegisterJsTapKind::NormalModuleFactoryBeforeResolve,
509+
RegisterJsTapKind::NormalModuleFactoryFactorize,
510+
RegisterJsTapKind::NormalModuleFactoryResolve,
511+
RegisterJsTapKind::NormalModuleFactoryAfterResolve,
512+
RegisterJsTapKind::NormalModuleFactoryCreateModule,
513+
RegisterJsTapKind::NormalModuleFactoryResolveForScheme,
514+
RegisterJsTapKind::ContextModuleFactoryBeforeResolve,
515+
RegisterJsTapKind::ContextModuleFactoryAfterResolve,
516+
RegisterJsTapKind::JavascriptModulesChunkHash,
517+
RegisterJsTapKind::HtmlPluginBeforeAssetTagGeneration,
518+
RegisterJsTapKind::HtmlPluginAlterAssetTags,
519+
RegisterJsTapKind::HtmlPluginAlterAssetTagGroups,
520+
RegisterJsTapKind::HtmlPluginAfterTemplateExecution,
521+
RegisterJsTapKind::HtmlPluginBeforeEmit,
522+
RegisterJsTapKind::HtmlPluginAfterEmit,
523+
RegisterJsTapKind::RuntimePluginCreateScript,
524+
RegisterJsTapKind::RuntimePluginCreateLink,
525+
RegisterJsTapKind::RuntimePluginLinkPreload,
526+
RegisterJsTapKind::RuntimePluginLinkPrefetch,
527+
RegisterJsTapKind::RsdoctorPluginModuleGraph,
528+
RegisterJsTapKind::RsdoctorPluginChunkGraph,
529+
RegisterJsTapKind::RsdoctorPluginModuleIds,
530+
RegisterJsTapKind::RsdoctorPluginModuleSources,
531+
RegisterJsTapKind::RsdoctorPluginAssets,
532+
];
533+
534+
#[napi(object)]
535+
pub struct RegisterJsTapScopeKinds {
536+
pub compiler: Vec<RegisterJsTapKind>,
537+
pub compilation: Vec<RegisterJsTapKind>,
538+
}
539+
540+
#[napi]
541+
pub fn get_register_js_tap_scope_kinds() -> RegisterJsTapScopeKinds {
542+
RegisterJsTapScopeKinds {
543+
compiler: COMPILER_HOOKS.to_vec(),
544+
compilation: COMPILATION_HOOKS.to_vec(),
545+
}
546+
}
547+
472548
#[derive(Clone)]
473549
#[napi(object, object_to_js = false)]
474550
pub struct RegisterJsTaps {
@@ -2033,8 +2109,7 @@ mod tests {
20332109
use super::*;
20342110

20352111
fn checker_with_used_kinds(kinds: &[RegisterJsTapKind]) -> HookUsageChecker {
2036-
let byte_len = ((RegisterJsTapKind::RsdoctorPluginAssets as usize) >> 3) + 1;
2037-
let mut bytes = vec![0; byte_len];
2112+
let mut bytes = vec![0; HOOK_USAGE_BUFFER_BYTE_LENGTH];
20382113
for kind in kinds {
20392114
let bit_index = *kind as usize;
20402115
bytes[bit_index >> 3] |= 1 << (bit_index & 7);
@@ -2049,9 +2124,9 @@ mod tests {
20492124
RegisterJsTapKind::CompilationProcessAssets,
20502125
]);
20512126

2052-
assert!(checker.is_used(&RegisterJsTapKind::CompilerCompilation));
2053-
assert!(checker.is_used(&RegisterJsTapKind::CompilationProcessAssets));
2054-
assert!(!checker.is_used(&RegisterJsTapKind::CompilerMake));
2127+
assert!(checker.is_used(RegisterJsTapKind::CompilerCompilation));
2128+
assert!(checker.is_used(RegisterJsTapKind::CompilationProcessAssets));
2129+
assert!(!checker.is_used(RegisterJsTapKind::CompilerMake));
20552130
}
20562131

20572132
#[test]
@@ -2061,11 +2136,11 @@ mod tests {
20612136

20622137
assert!(should_skip_register(
20632138
&empty_checker,
2064-
&RegisterJsTapKind::CompilationBuildModule,
2139+
RegisterJsTapKind::CompilationBuildModule,
20652140
));
20662141
assert!(!should_skip_register(
20672142
&used_checker,
2068-
&RegisterJsTapKind::CompilerCompilation,
2143+
RegisterJsTapKind::CompilerCompilation,
20692144
));
20702145
}
20712146
}

‎packages/rspack/src/HookUsageTracker.ts‎

Lines changed: 71 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,27 @@
11
import { Buffer } from 'node:buffer';
22
import binding from '@rspack/binding';
3-
import type * as liteTapable from '@rspack/lite-tapable';
3+
import * as liteTapable from '@rspack/lite-tapable';
44
import type { Compiler } from './Compiler';
55

66
const HOOK_USAGE_TRACKED = Symbol('rspack.hookUsageTracked');
7+
const HOOK_USAGE_MARK_INTERCEPT = Symbol('rspack.hookUsageMarkIntercept');
78
const HOOK_MAP_USAGE_TRACKED = Symbol('rspack.hookMapUsageTracked');
89
const INTERNAL_HOOK_INTERCEPTOR = Symbol('rspack.internalHookInterceptor');
910
export const COMPILER_HOOK_USAGE_TRACKERS = new WeakMap<
1011
Compiler,
1112
JsHookUsageTracker
1213
>();
14+
let hookInterceptUsagePatched = false;
1315

14-
const COMPILER_SCOPED_HOOK_KINDS = new Set<binding.RegisterJsTapKind>([
15-
binding.RegisterJsTapKind.CompilerThisCompilation,
16-
binding.RegisterJsTapKind.CompilerCompilation,
17-
binding.RegisterJsTapKind.CompilerMake,
18-
binding.RegisterJsTapKind.CompilerFinishMake,
19-
binding.RegisterJsTapKind.CompilerShouldEmit,
20-
binding.RegisterJsTapKind.CompilerEmit,
21-
binding.RegisterJsTapKind.CompilerAfterEmit,
22-
binding.RegisterJsTapKind.CompilerAssetEmitted,
23-
]);
24-
25-
const HOOK_USAGE_BUFFER_BYTE_LENGTH = 8;
26-
const COMPILER_SCOPED_BYTE_MASKS = Buffer.alloc(HOOK_USAGE_BUFFER_BYTE_LENGTH);
27-
28-
for (const kind of COMPILER_SCOPED_HOOK_KINDS) {
29-
COMPILER_SCOPED_BYTE_MASKS[kind >> 3] |= 1 << (kind & 7);
30-
}
16+
const { compiler: COMPILER_HOOKS, compilation: COMPILATION_HOOKS } =
17+
binding.getRegisterJsTapScopeKinds();
18+
const HOOK_USAGE_BUFFER_BYTE_LENGTH =
19+
(Math.max(...COMPILER_HOOKS, ...COMPILATION_HOOKS) >> 3) + 1;
3120

3221
type InterceptableHook = liteTapable.Hook<any, any, any> & {
3322
[HOOK_USAGE_TRACKED]?: boolean;
34-
_tap?: (type: string, options: liteTapable.Options, fn: Function) => void;
35-
intercept?: (interceptor: object) => void;
23+
[HOOK_USAGE_MARK_INTERCEPT]?: () => void;
24+
interceptors?: object[];
3625
};
3726

3827
type InterceptableHookMap<H extends liteTapable.Hook<any, any, any>> =
@@ -45,31 +34,30 @@ export class JsHookUsageTracker {
4534
#buffer: Buffer;
4635

4736
constructor() {
48-
const storage =
49-
typeof SharedArrayBuffer === 'function'
50-
? new SharedArrayBuffer(HOOK_USAGE_BUFFER_BYTE_LENGTH)
51-
: new ArrayBuffer(HOOK_USAGE_BUFFER_BYTE_LENGTH);
52-
this.#buffer = Buffer.from(storage);
37+
this.#buffer = Buffer.alloc(HOOK_USAGE_BUFFER_BYTE_LENGTH);
5338
}
5439

5540
getBuffer() {
5641
return this.#buffer;
5742
}
5843

5944
mark(kind: binding.RegisterJsTapKind) {
60-
this.#buffer[kind >> 3] |= 1 << (kind & 7);
45+
const byteIndex = kind >> 3;
46+
const bitMask = 1 << (kind & 7);
47+
if (this.#buffer[byteIndex] & bitMask) {
48+
return;
49+
}
50+
this.#buffer[byteIndex] |= bitMask;
6151
}
6252

6353
resetCompilationScopedBits() {
64-
for (let i = 0; i < this.#buffer.length; i++) {
65-
this.#buffer[i] &= COMPILER_SCOPED_BYTE_MASKS[i];
54+
for (const kind of COMPILATION_HOOKS) {
55+
this.#buffer[kind >> 3] &= ~(1 << (kind & 7));
6656
}
6757
}
6858

6959
resetAllBits() {
70-
for (let i = 0; i < this.#buffer.length; i++) {
71-
this.#buffer[i] = 0;
72-
}
60+
this.#buffer.fill(0);
7361
}
7462
}
7563

@@ -78,6 +66,28 @@ const isInternalHookInterceptor = (interceptor: unknown) =>
7866
interceptor !== null &&
7967
INTERNAL_HOOK_INTERCEPTOR in interceptor;
8068

69+
const patchHookInterceptUsage = () => {
70+
if (hookInterceptUsagePatched) {
71+
return;
72+
}
73+
hookInterceptUsagePatched = true;
74+
75+
const hookPrototype = liteTapable.HookBase.prototype as liteTapable.Hook<
76+
any,
77+
any,
78+
any
79+
> & {
80+
intercept: (interceptor: object) => void;
81+
};
82+
const rawIntercept = hookPrototype.intercept;
83+
hookPrototype.intercept = function (interceptor: object) {
84+
rawIntercept.call(this, interceptor);
85+
if (!isInternalHookInterceptor(interceptor)) {
86+
(this as InterceptableHook)[HOOK_USAGE_MARK_INTERCEPT]?.();
87+
}
88+
};
89+
};
90+
8191
export const markHookInterceptorAsInternal = <T extends object>(
8292
interceptor: T,
8393
): T => {
@@ -90,30 +100,38 @@ export const trackHookUsage = <T extends liteTapable.Hook<any, any, any>>(
90100
tracker: JsHookUsageTracker,
91101
kind: binding.RegisterJsTapKind,
92102
): T => {
103+
patchHookInterceptUsage();
93104
const trackableHook = hook as InterceptableHook;
94105
if (trackableHook[HOOK_USAGE_TRACKED]) {
95106
return hook;
96107
}
97108
trackableHook[HOOK_USAGE_TRACKED] = true;
98-
99-
if (trackableHook._tap) {
100-
const tap = trackableHook._tap.bind(hook);
101-
trackableHook._tap = ((type, options, fn) => {
102-
tracker.mark(kind);
103-
return tap(type, options, fn);
104-
}) as typeof trackableHook._tap;
105-
}
106-
107-
if (trackableHook.intercept) {
108-
const intercept = trackableHook.intercept.bind(hook);
109-
trackableHook.intercept = ((interceptor: object) => {
110-
if (!isInternalHookInterceptor(interceptor)) {
111-
tracker.mark(kind);
112-
}
113-
return intercept(interceptor);
114-
}) as typeof trackableHook.intercept;
109+
let marked = false;
110+
const markUsed = () => {
111+
if (marked) {
112+
return;
113+
}
114+
marked = true;
115+
trackableHook[HOOK_USAGE_MARK_INTERCEPT] = undefined;
116+
tracker.mark(kind);
117+
};
118+
trackableHook[HOOK_USAGE_MARK_INTERCEPT] = markUsed;
119+
120+
hook.intercept(
121+
markHookInterceptorAsInternal({
122+
register: (tapInfo) => {
123+
markUsed();
124+
return tapInfo;
125+
},
126+
}),
127+
);
128+
if (
129+
trackableHook.interceptors?.some(
130+
(interceptor) => !isInternalHookInterceptor(interceptor),
131+
)
132+
) {
133+
markUsed();
115134
}
116-
117135
return hook;
118136
};
119137

@@ -130,7 +148,10 @@ export const trackHookMapUsage = <H extends liteTapable.Hook<any, any, any>>(
130148

131149
hookMap.intercept({
132150
factory: (_key, hook) => {
133-
return trackHookUsage<H>(hook!, tracker, kind);
151+
if (hook === undefined) {
152+
throw new Error('HookMap factory interceptor must receive a hook');
153+
}
154+
return trackHookUsage<H>(hook, tracker, kind);
134155
},
135156
});
136157
for (const hook of trackableHookMap._map?.values() ?? []) {

0 commit comments

Comments
 (0)