Preserve DI scope ownership and synchronize concurrent property registration

src/TUnit.Core/AsyncEvent.cs

@@ -49,6 +49,20 @@ public class AsyncEvent<TEventArgs>
 {
     private List<Invocation>? _handlers;
 
+    internal AsyncEvent<TEventArgs> Clone() => new()
+    {
+        _handlers = _handlers is null ? null : [.. _handlers]
+    };
+
+    internal void AddMissingFrom(AsyncEvent<TEventArgs> source)
+    {
+        foreach (var invocation in source.InvocationList)
+        {
+            if (_handlers?.Contains(invocation) != true)
+                (_handlers ??= []).Insert(FindInsertionIndex(invocation.Order), invocation);
+        }
+    }
+
     /// <summary>
     /// Represents a single invocation of an async event handler.
     /// </summary>

src/TUnit.Core/Attributes/TestData/DependencyInjectionDataSourceSourceAttribute.cs

@@ -16,18 +16,9 @@ public abstract class DependencyInjectionDataSourceAttribute<TScope> : UntypedDa
             // disposing the returned objects, as it does for any other data source.
             var scopeOwnsCreatedObjects = scope is IDisposable or IAsyncDisposable;
 
-            // Set up disposal for this specific scope in the current test context
-            dataGeneratorMetadata.TestBuilderContext.Current.Events.OnDispose += async (_, _) =>
-            {
-                if (scope is IAsyncDisposable asyncDisposable)
-                {
-                    await asyncDisposable.DisposeAsync().ConfigureAwait(false);
-                }
-                else if (scope is IDisposable disposable)
-                {
-                    disposable.Dispose();
-                }
-            };
+            // Register before resolving so a failed factory still releases its scope.
+            var cleanup = dataGeneratorMetadata.TestBuilderContext.Current.Events.GetOrCreateDataRowScopes();
+            var owner = cleanup.Register(scope);
 
             return dataGeneratorMetadata.MembersToGenerate
                 .Select(m => m switch
@@ -41,6 +32,7 @@ public abstract class DependencyInjectionDataSourceAttribute<TScope> : UntypedDa
                 .Select(x =>
                 {
                     var created = Create(scope, x);
+                    cleanup.RegisterObject(owner, created);
 
                     if (scopeOwnsCreatedObjects)
                     {

src/TUnit.Core/DataRowScopeCollection.cs

@@ -0,0 +1,119 @@
+namespace TUnit.Core;
+
+/// <summary>
+/// Owns a row's scopes. Object lookup releases each scope at its first resolved object's graph
+/// depth; completion releases scopes left by failed factories or untracked values.
+/// </summary>
+internal sealed class DataRowScopeCollection
+{
+    private readonly object _lock = new();
+    private readonly List<ScopeOwner> _owners = [];
+    private readonly Dictionary<object, List<ScopeOwner>> _objects = new(Helpers.ReferenceEqualityComparer.Instance);
+
+    internal ScopeOwner Register(object? scope)
+    {
+        var owner = new ScopeOwner(scope);
+        lock (_lock)
+            _owners.Add(owner);
+        return owner;
+    }
+
+    internal void RegisterObject(ScopeOwner owner, object? instance)
+    {
+        if (instance is null)
+        {
+            return;
+        }
+
+        lock (_lock)
+        {
+            if (!_objects.TryGetValue(instance, out var owners))
+                _objects.Add(instance, owners = []);
+            if (!owners.Contains(owner))
+                owners.Add(owner);
+        }
+    }
+
+    internal DataRowScopeCollection Clone()
+    {
+        var clone = new DataRowScopeCollection();
+        clone.AddRange(this);
+        return clone;
+    }
+
+    internal void AddRange(DataRowScopeCollection scopes)
+    {
+        // Snapshot the source before locking the destination to avoid lock-order inversions.
+        ScopeOwner[] owners;
+        KeyValuePair<object, ScopeOwner[]>[] objects;
+        lock (scopes._lock)
+        {
+            owners = scopes._owners.ToArray();
+            objects = scopes._objects.Select(pair =>
+                new KeyValuePair<object, ScopeOwner[]>(pair.Key, pair.Value.ToArray())).ToArray();
+        }
+        lock (_lock)
+        {
+            _owners.AddRange(owners);
+            foreach (var pair in objects)
+                foreach (var owner in pair.Value)
+                    RegisterObject(owner, pair.Key);
+        }
+    }
+
+    internal ValueTask<List<Exception>?> ReleaseAsync(object instance, List<Exception>? exceptions = null)
+    {
+        ScopeOwner[] owners;
+        lock (_lock)
+            owners = _objects.TryGetValue(instance, out var registered) ? registered.ToArray() : [];
+        return DisposeAsync(owners, exceptions);
+    }
+
+    internal ValueTask<List<Exception>?> CompleteAsync(List<Exception>? exceptions = null)
+    {
+        ScopeOwner[] owners;
+        lock (_lock)
+            owners = _owners.ToArray();
+        return DisposeAsync(owners, exceptions);
+    }
+
+    private static async ValueTask<List<Exception>?> DisposeAsync(ScopeOwner[] owners, List<Exception>? exceptions)
+    {
+        foreach (var owner in owners)
+        {
+            try
+            {
+                await owner.DisposeAsync().ConfigureAwait(false);
+            }
+            catch (Exception exception)
+            {
+                (exceptions ??= []).Add(exception);
+            }
+        }
+
+        return exceptions;
+    }
+
+    /// <summary>Shared by cloned rows; scope disposal is attempted once, including when it throws.</summary>
+    internal sealed class ScopeOwner(object? scope)
+    {
+        private int _disposed;
+
+        internal async ValueTask DisposeAsync()
+        {
+            if (Interlocked.Exchange(ref _disposed, 1) != 0)
+            {
+                return;
+            }
+
+            if (scope is IAsyncDisposable asyncDisposable)
+            {
+                await asyncDisposable.DisposeAsync().ConfigureAwait(false);
+            }
+            else if (scope is IDisposable disposable)
+            {
+                disposable.Dispose();
+            }
+        }
+    }
+}

src/TUnit.Core/Models/TestSessionContext.cs

@@ -110,6 +110,18 @@ internal static void DeferOnDisposeToSessionEnd(TestContextEvents? events)
     {
         lock (s_deferredOnDisposeLock)
         {
+            // Preserve session ownership for discovery rows now that DI cleanup is stored separately.
+            if (events?.DetachDataRowScopes() is { } scopes)
+            {
+                events.OnDispose += async (_, _) =>
+                {
+                    if (await scopes.CompleteAsync().ConfigureAwait(false) is { Count: > 0 } errors)
+                    {
+                        throw new AggregateException("Errors occurred while disposing discovery scopes", errors);
+                    }
+                };
+            }
+
             if (events is not { OnDispose: { } onDispose })
             {
                 return;

src/TUnit.Core/TestBuilderContext.cs

@@ -91,7 +91,7 @@ public TestContextEvents Events
     {
         // A sharing context resolves the source's container on every read rather than caching it, so a
         // later assignment to the source's Events reaches this context too.
-        get => _events ?? _eventsSource?.Events ?? (_events = new TestContextEvents());
+        get => _events ?? _eventsSource?.Events ?? LazyInitializer.EnsureInitialized(ref _events)!;
         set
         {
             _eventsSource = null;
@@ -141,6 +141,30 @@ internal void ShareEventsWith(TestBuilderContext source)
     /// </summary>
     internal Attribute[]? InitializedAttributes { get; set; }
 
+    // Child rows start with isolated state and no event subscriptions. Callers that build an
+    // executable view can explicitly share the state/events with their factory context instead.
+    internal TestBuilderContext CreateChild(IDataSourceAttribute? classDataSource, IDataSourceAttribute? dataSource = null, bool shareFactoryContext = false)
+    {
+        var child = new TestBuilderContext
+        {
+            TestMetadata = TestMetadata,
+            InitializedAttributes = InitializedAttributes,
+            ClassConstructor = ClassConstructor,
+            ClassDataSourceAttribute = classDataSource,
+            DataSourceAttribute = dataSource
+        };
+        if (shareFactoryContext)
+        {
+            child.ShareStateBagWith(this);
+            child.ShareEventsWith(this);
+        }
+        else
+        {
+            CopyStateBagTo(child);
+        }
+        return child;
+    }
+
     public void RegisterForInitialization(object? obj)
     {
         Events.OnInitialize += async (sender, args) =>

src/TUnit.Core/TestContextEvents.cs

@@ -7,6 +7,68 @@ namespace TUnit.Core;
 /// </summary>
 public class TestContextEvents : ITestEvents
 {
+    internal TestContextEvents Clone() => new()
+    {
+        DataRowScopes = DataRowScopes?.Clone(),
+        OnDispose = OnDispose?.Clone(),
+        OnTestRegistered = OnTestRegistered?.Clone(),
+        OnInitialize = OnInitialize?.Clone(),
+        OnTestStart = OnTestStart?.Clone(),
+        OnTestEnd = OnTestEnd?.Clone(),
+        OnTestSkipped = OnTestSkipped?.Clone(),
+        OnTestRetry = OnTestRetry?.Clone()
+    };
+
+    private object? _scopeLock;
+
+    internal void AddDataRowScopes(DataRowScopeCollection scopes)
+    {
+        lock (LazyInitializer.EnsureInitialized(ref _scopeLock)!)
+        {
+            GetOrCreateDataRowScopes().AddRange(scopes);
+        }
+    }
+
+    internal DataRowScopeCollection? DetachDataRowScopes()
+    {
+        lock (LazyInitializer.EnsureInitialized(ref _scopeLock)!)
+        {
+            var scopes = DataRowScopes;
+            DataRowScopes = null;
+            return scopes;
+        }
+    }
+
+    internal void AddCleanupFrom(TestContextEvents source)
+    {
+        if (source.DataRowScopes is { } scopes)
+            AddDataRowScopes(scopes);
+        if (source.OnDispose is { } onDispose)
+        {
+            lock (LazyInitializer.EnsureInitialized(ref _scopeLock)!)
+            {
+                (OnDispose ??= new()).AddMissingFrom(onDispose);
+            }
+        }
+    }
+
+    private DataRowScopeCollection? _dataRowScopes;
+
+    internal DataRowScopeCollection GetOrCreateDataRowScopes()
+    {
+        lock (LazyInitializer.EnsureInitialized(ref _scopeLock)!)
+        {
+            return _dataRowScopes ??= new();
+        }
+    }
+
+    // Clones share once-only scope owners while keeping row ownership collections isolated.
+    internal DataRowScopeCollection? DataRowScopes
+    {
+        get => Volatile.Read(ref _dataRowScopes);
+        set => Volatile.Write(ref _dataRowScopes, value);
+    }
+
     public AsyncEvent<TestContext>? OnDispose { get; set; }
     public AsyncEvent<TestContext>? OnTestRegistered { get; set; }
     public AsyncEvent<TestContext>? OnInitialize { get; set; }

src/TUnit.Core/Tracking/ObjectTracker.cs

@@ -269,6 +269,11 @@ private static void RemoveObjectsNotIn(SortedList<int, HashSet<object>>? tracked
 
         foreach (var obj in bucket)
         {
+            if (testContext?.InternalEventsIfCreated?.DataRowScopes is { } scopes)
+            {
+                cleanupExceptions = await scopes.ReleaseAsync(obj, cleanupExceptions).ConfigureAwait(false);
+            }
+
             try
             {
                 var task = UntrackObject(obj, testContext);
@@ -718,15 +723,19 @@ internal static bool IsExternallyOwned(object obj)
         => s_hasExternallyOwned && s_externallyOwned.TryGetValue(obj, out _);
 
     // Objects whose lifetime another owner controls (a dependency injection scope, for example).
-    // The tracker counts references but leaves disposal to their owner.
+    // The tracker balances references but never disposes them; their owner controls disposal.
     private static readonly ConditionalWeakTable<object, object> s_externallyOwned = new();
     private static readonly object s_externallyOwnedMarker = new();
 
+    // Ownership cannot live only on a row or counted entry: a container singleton can be
+    // supplied by another source before or after its DI row, including after its count reaches zero.
+    // A weak identity mark preserves that ownership without retaining the singleton.
+    // The flag deliberately stays set: live marks can outlast any individual row or test session.
     // Skips the table lookup for every tracked object until something is marked.
     private static volatile bool s_hasExternallyOwned;
 
     /// <summary>
-    /// Marks <paramref name="obj"/> as disposed by another owner, so the tracker counts references without disposing it.
+    /// Marks <paramref name="obj"/> as disposed by another owner, so the tracker never disposes it.
     /// Dependency injection data sources mark the services they resolve, because the scope that created
     /// them disposes them when the test ends (#7205). Existing tracker references are released
     /// without disposing the object, regardless of which data source supplied it first.