perf(engine): lazy hook metadata registration (#5448)

TUnit.Core.SourceGenerator.Tests/HooksTests.DisposableFieldTests.verified.txt

@@ -1,4 +1,4 @@
-// <auto-generated/>
+// <auto-generated/>
 #pragma warning disable
 
 #nullable enable
@@ -30,7 +30,8 @@ namespace TUnit.Generated
 {
     internal static partial class TUnit_HookRegistration
     {
-        static readonly int _h_TUnit_TestProject_DisposableFieldTests_Setup_Before_Test = global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.BeforeTestHooks, typeof(global::TUnit.TestProject.DisposableFieldTests),
+        static readonly int _h_TUnit_TestProject_DisposableFieldTests_Setup_Before_Test = global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.BeforeTestHooks, typeof(global::TUnit.TestProject.DisposableFieldTests), global::TUnit.Core.HookRegistrationIndices.GetNextBeforeTestHookIndex(),
+    static __registrationIndex =>
     new InstanceHookMethod
     {
         InitClassType = typeof(global::TUnit.TestProject.DisposableFieldTests),
@@ -57,7 +58,7 @@ namespace TUnit.Generated
         },
         HookExecutor = DefaultExecutor.Instance,
         Order = 0,
-        RegistrationIndex = global::TUnit.Core.HookRegistrationIndices.GetNextBeforeTestHookIndex(),
+        RegistrationIndex = __registrationIndex,
         Body = global::TUnit.Generated.Hooks.TUnit_TestProject_DisposableFieldTests_Setup_Before_Test.TUnit_TestProject_DisposableFieldTests_Setup_Before_TestInitializer.global_TUnit_TestProject_DisposableFieldTests_Setup_0Params_Body
     }
 );
@@ -99,7 +100,8 @@ namespace TUnit.Generated
 {
     internal static partial class TUnit_HookRegistration
     {
-        static readonly int _h_TUnit_TestProject_DisposableFieldTests_Blah_After_Test = global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.AfterTestHooks, typeof(global::TUnit.TestProject.DisposableFieldTests),
+        static readonly int _h_TUnit_TestProject_DisposableFieldTests_Blah_After_Test = global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.AfterTestHooks, typeof(global::TUnit.TestProject.DisposableFieldTests), global::TUnit.Core.HookRegistrationIndices.GetNextAfterTestHookIndex(),
+    static __registrationIndex =>
     new InstanceHookMethod
     {
         InitClassType = typeof(global::TUnit.TestProject.DisposableFieldTests),
@@ -126,7 +128,7 @@ namespace TUnit.Generated
         },
         HookExecutor = DefaultExecutor.Instance,
         Order = 0,
-        RegistrationIndex = global::TUnit.Core.HookRegistrationIndices.GetNextAfterTestHookIndex(),
+        RegistrationIndex = __registrationIndex,
         Body = global::TUnit.Generated.Hooks.TUnit_TestProject_DisposableFieldTests_Blah_After_Test.TUnit_TestProject_DisposableFieldTests_Blah_After_TestInitializer.global_TUnit_TestProject_DisposableFieldTests_Blah_0Params_Body
     }
 );

TUnit.Core.SourceGenerator/Generators/HookMetadataGenerator.cs

@@ -459,7 +459,9 @@ private static string GetSimpleTypeName(string fullyQualifiedTypeName)
 
     /// <summary>
     /// Emits a single SourceRegistrar.RegisterHook(...) expression for use as a field initializer.
-    /// The hook object is constructed inline, so no per-hook method needs JIT.
+    /// Module init pays only for one HookRegistrationIndices.GetNext*() increment + a static
+    /// lambda reference; the heavy MethodMetadata/ClassMetadata graph is constructed lazily
+    /// on first hook execution via LazyHookEntry.Materialize().
     /// </summary>
     private static void GenerateInlineHookRegistration(CodeWriter writer, HookModel hook, string safeFileName)
     {
@@ -469,17 +471,19 @@ private static void GenerateInlineHookRegistration(CodeWriter writer, HookModel
 
         // Determine collection name and key expression
         var (collectionName, keyExpr) = GetHookCollectionAndKey(hook, typeDisplay, isInstance);
+        var indexExpr = $"global::TUnit.Core.HookRegistrationIndices.GetNext{GetHookIndexMethodName(hook)}";
 
         if (keyExpr != null)
         {
-            writer.AppendLine($"global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.{collectionName}, {keyExpr},");
+            writer.AppendLine($"global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.{collectionName}, {keyExpr}, {indexExpr},");
         }
         else
         {
-            writer.AppendLine($"global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.{collectionName},");
+            writer.AppendLine($"global::TUnit.Core.SourceRegistrar.RegisterHook(global::TUnit.Core.Sources.{collectionName}, {indexExpr},");
         }
 
         writer.Indent();
+        writer.AppendLine("static __registrationIndex =>");
         GenerateHookObject(writer, hook, isInstance, delegatePrefix);
         writer.Unindent();
         writer.Append(")");
@@ -715,7 +719,7 @@ private static void GenerateHookObject(CodeWriter writer, HookModel hook, bool i
 
         writer.AppendLine($"HookExecutor = {HookExecutorHelper.GetHookExecutor(hook.HookExecutorTypeName)},");
         writer.AppendLine($"Order = {hook.Order},");
-        writer.AppendLine($"RegistrationIndex = global::TUnit.Core.HookRegistrationIndices.GetNext{GetHookIndexMethodName(hook)},");
+        writer.AppendLine("RegistrationIndex = __registrationIndex,");
         writer.AppendLine($"Body = {bodyRef}" + (isInstance ? "" : ","));
 
         if (!isInstance)

TUnit.Core/Hooks/LazyHookEntry.cs

@@ -0,0 +1,79 @@
+using System.ComponentModel;
+
+namespace TUnit.Core.Hooks;
+
+/// <summary>
+/// Wraps a hook registration as a factory plus eagerly-computed registration index.
+/// The full <see cref="HookMethod"/> graph (MethodMetadata, ClassMetadata, parameter arrays,
+/// delegate, etc.) is constructed only on first call to <see cref="Materialize"/>.
+/// This avoids paying O(N) construction cost per hook at module initialization.
+/// <para>
+/// The registration index is captured eagerly so that hook ordering is preserved even
+/// when materialization happens out of declaration order.
+/// </para>
+/// </summary>
+#if !DEBUG
+[EditorBrowsable(EditorBrowsableState.Never)]
+#endif
+public sealed class LazyHookEntry<T> where T : HookMethod
+{
+    private Func<int, T>? _factory;
+    // volatile is required for the double-checked-lock pattern in Materialize() — without it
+    // the JIT is allowed to cache the unlocked read in a register on weakly-ordered architectures,
+    // which would let one thread perpetually see null after another thread published the value.
+    // T is constrained to HookMethod (a reference type), so `volatile T?` is legal.
+    private volatile T? _materialized;
+    private readonly object _lock = new();
+
+    /// <summary>
+    /// The eagerly-computed registration index. Used to preserve declaration order
+    /// across hooks that share the same <see cref="HookMethod.Order"/> value.
+    /// </summary>
+    public int RegistrationIndex { get; }
+
+    /// <summary>
+    /// Creates a lazy entry from a factory that will produce the materialized hook on first
+    /// access. The factory MUST be a static lambda (no captures) to avoid per-hook closure
+    /// allocations and to remain AOT compatible. The eagerly-computed registration index is
+    /// passed back into the factory so it can be assigned to <see cref="HookMethod.RegistrationIndex"/>.
+    /// </summary>
+    public LazyHookEntry(int registrationIndex, Func<int, T> factory)
+    {
+        RegistrationIndex = registrationIndex;
+        _factory = factory;
+    }
+
+    /// <summary>
+    /// Creates a lazy entry from an already-materialized hook (used by reflection-mode discovery).
+    /// </summary>
+    public LazyHookEntry(T hook)
+    {
+        _materialized = hook;
+        RegistrationIndex = hook.RegistrationIndex;
+    }
+
+    /// <summary>
+    /// Returns the materialized hook, constructing it on first access.
+    /// Subsequent calls return the cached instance. Thread-safe.
+    /// </summary>
+    public T Materialize()
+    {
+        if (_materialized is not null)
+        {
+            return _materialized;
+        }
+
+        lock (_lock)
+        {
+            if (_materialized is not null)
+            {
+                return _materialized;
+            }
+
+            var hook = _factory!(RegistrationIndex);
+            _materialized = hook;
+            _factory = null;
+            return hook;
+        }
+    }
+}

TUnit.Core/SourceRegistrar.cs

@@ -2,6 +2,7 @@
 using System.Diagnostics;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using TUnit.Core.Hooks;
 using TUnit.Core.Interfaces.SourceGenerator;
 
 namespace TUnit.Core;
@@ -61,29 +62,43 @@ public static void RegisterProperty(IPropertySource propertySource)
     }
 
     /// <summary>
-    /// Registers a hook into a type-keyed dictionary and returns a dummy value for use as a field initializer.
+    /// Registers a hook factory into a type-keyed dictionary. The factory is not invoked
+    /// until the engine materializes the hook for execution. Use a <c>static</c> lambda
+    /// (no captures) to keep module-init cost minimal and to remain AOT compatible.
+    /// Returns a dummy value for use as a static field initializer.
     /// </summary>
-    public static int RegisterHook<T>(ConcurrentDictionary<Type, ConcurrentBag<T>> dictionary, Type key, T hook)
+    public static int RegisterHook<T>(ConcurrentDictionary<Type, ConcurrentBag<LazyHookEntry<T>>> dictionary, Type key, int registrationIndex, Func<int, T> factory)
+        where T : HookMethod
     {
-        dictionary.GetOrAdd(key, static _ => new ConcurrentBag<T>()).Add(hook);
+        dictionary.GetOrAdd(key, static _ => new ConcurrentBag<LazyHookEntry<T>>())
+            .Add(new LazyHookEntry<T>(registrationIndex, factory));
         return 0;
     }
 
     /// <summary>
-    /// Registers a hook into an assembly-keyed dictionary and returns a dummy value for use as a field initializer.
+    /// Registers a hook factory into an assembly-keyed dictionary. The factory is not invoked
+    /// until the engine materializes the hook for execution. Use a <c>static</c> lambda
+    /// (no captures) to keep module-init cost minimal and to remain AOT compatible.
+    /// Returns a dummy value for use as a static field initializer.
     /// </summary>
-    public static int RegisterHook<T>(ConcurrentDictionary<Assembly, ConcurrentBag<T>> dictionary, Assembly key, T hook)
+    public static int RegisterHook<T>(ConcurrentDictionary<Assembly, ConcurrentBag<LazyHookEntry<T>>> dictionary, Assembly key, int registrationIndex, Func<int, T> factory)
+        where T : HookMethod
     {
-        dictionary.GetOrAdd(key, static _ => new ConcurrentBag<T>()).Add(hook);
+        dictionary.GetOrAdd(key, static _ => new ConcurrentBag<LazyHookEntry<T>>())
+            .Add(new LazyHookEntry<T>(registrationIndex, factory));
         return 0;
     }
 
     /// <summary>
-    /// Registers a hook into a global bag and returns a dummy value for use as a field initializer.
+    /// Registers a hook factory into a global bag. The factory is not invoked until the engine
+    /// materializes the hook for execution. Use a <c>static</c> lambda (no captures) to keep
+    /// module-init cost minimal and to remain AOT compatible.
+    /// Returns a dummy value for use as a static field initializer.
     /// </summary>
-    public static int RegisterHook<T>(ConcurrentBag<T> bag, T hook)
+    public static int RegisterHook<T>(ConcurrentBag<LazyHookEntry<T>> bag, int registrationIndex, Func<int, T> factory)
+        where T : HookMethod
     {
-        bag.Add(hook);
+        bag.Add(new LazyHookEntry<T>(registrationIndex, factory));
         return 0;
     }
 

TUnit.Core/Sources.cs

@@ -1,5 +1,6 @@
 using System.Collections.Concurrent;
 using System.Reflection;
+using TUnit.Core.Hooks;
 using TUnit.Core.Interfaces.SourceGenerator;
 
 namespace TUnit.Core;
@@ -12,25 +13,28 @@ public static class Sources
     public static readonly ConcurrentQueue<Func<Assembly>> AssemblyLoaders = [];
     public static readonly ConcurrentQueue<IDynamicTestSource> DynamicTestSources = [];
 
-    public static readonly ConcurrentDictionary<Type, ConcurrentBag<Hooks.InstanceHookMethod>> BeforeTestHooks = new();
-    public static readonly ConcurrentDictionary<Type, ConcurrentBag<Hooks.InstanceHookMethod>> AfterTestHooks = new();
-    public static readonly ConcurrentBag<Hooks.BeforeTestHookMethod> BeforeEveryTestHooks = [];
-    public static readonly ConcurrentBag<Hooks.AfterTestHookMethod> AfterEveryTestHooks = [];
-
-    public static readonly ConcurrentDictionary<Type, ConcurrentBag<Hooks.BeforeClassHookMethod>> BeforeClassHooks = new();
-    public static readonly ConcurrentDictionary<Type, ConcurrentBag<Hooks.AfterClassHookMethod>> AfterClassHooks = new();
-    public static readonly ConcurrentBag<Hooks.BeforeClassHookMethod> BeforeEveryClassHooks = [];
-    public static readonly ConcurrentBag<Hooks.AfterClassHookMethod> AfterEveryClassHooks = [];
-
-    public static readonly ConcurrentDictionary<Assembly, ConcurrentBag<Hooks.BeforeAssemblyHookMethod>> BeforeAssemblyHooks = new();
-    public static readonly ConcurrentDictionary<Assembly, ConcurrentBag<Hooks.AfterAssemblyHookMethod>> AfterAssemblyHooks = new();
-    public static readonly ConcurrentBag<Hooks.BeforeAssemblyHookMethod> BeforeEveryAssemblyHooks = [];
-    public static readonly ConcurrentBag<Hooks.AfterAssemblyHookMethod> AfterEveryAssemblyHooks = [];
-
-    public static readonly ConcurrentBag<Hooks.BeforeTestSessionHookMethod> BeforeTestSessionHooks = [];
-    public static readonly ConcurrentBag<Hooks.AfterTestSessionHookMethod> AfterTestSessionHooks = [];
-    public static readonly ConcurrentBag<Hooks.BeforeTestDiscoveryHookMethod> BeforeTestDiscoveryHooks = [];
-    public static readonly ConcurrentBag<Hooks.AfterTestDiscoveryHookMethod> AfterTestDiscoveryHooks = [];
+    // Hook collections store LazyHookEntry wrappers — the heavy MethodMetadata/ClassMetadata
+    // construction is deferred until first Materialize() call (typically during engine
+    // discovery/execution rather than at module initialization).
+    public static readonly ConcurrentDictionary<Type, ConcurrentBag<LazyHookEntry<InstanceHookMethod>>> BeforeTestHooks = new();
+    public static readonly ConcurrentDictionary<Type, ConcurrentBag<LazyHookEntry<InstanceHookMethod>>> AfterTestHooks = new();
+    public static readonly ConcurrentBag<LazyHookEntry<BeforeTestHookMethod>> BeforeEveryTestHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<AfterTestHookMethod>> AfterEveryTestHooks = [];
+
+    public static readonly ConcurrentDictionary<Type, ConcurrentBag<LazyHookEntry<BeforeClassHookMethod>>> BeforeClassHooks = new();
+    public static readonly ConcurrentDictionary<Type, ConcurrentBag<LazyHookEntry<AfterClassHookMethod>>> AfterClassHooks = new();
+    public static readonly ConcurrentBag<LazyHookEntry<BeforeClassHookMethod>> BeforeEveryClassHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<AfterClassHookMethod>> AfterEveryClassHooks = [];
+
+    public static readonly ConcurrentDictionary<Assembly, ConcurrentBag<LazyHookEntry<BeforeAssemblyHookMethod>>> BeforeAssemblyHooks = new();
+    public static readonly ConcurrentDictionary<Assembly, ConcurrentBag<LazyHookEntry<AfterAssemblyHookMethod>>> AfterAssemblyHooks = new();
+    public static readonly ConcurrentBag<LazyHookEntry<BeforeAssemblyHookMethod>> BeforeEveryAssemblyHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<AfterAssemblyHookMethod>> AfterEveryAssemblyHooks = [];
+
+    public static readonly ConcurrentBag<LazyHookEntry<BeforeTestSessionHookMethod>> BeforeTestSessionHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<AfterTestSessionHookMethod>> AfterTestSessionHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<BeforeTestDiscoveryHookMethod>> BeforeTestDiscoveryHooks = [];
+    public static readonly ConcurrentBag<LazyHookEntry<AfterTestDiscoveryHookMethod>> AfterTestDiscoveryHooks = [];
 
     public static readonly ConcurrentQueue<Func<Task>> GlobalInitializers = [];
     public static readonly ConcurrentQueue<IPropertySource> PropertySources = [];