fix: never hand out a disposed shared data-source instance

docs/docs/writing-tests/class-data-source.md

@@ -170,6 +170,35 @@ public class ProductTests
 
 Use `IAsyncInitializer`, not `IAsyncDiscoveryInitializer`, for this work. `IAsyncDiscoveryInitializer` runs while tests are built. See the [test lifecycle](lifecycle.md) for the full order.
 
+## When Shared Instances Are Disposed
+
+TUnit counts the tests that hold a shared instance. When the last of them ends, TUnit disposes the instance. A test that asks for the same instance after that, such as a dynamic test added later in the run, gets a new instance.
+
+A test can also fetch a shared instance itself with `SharedDataSources.GetOrCreate`. It gets the same instance that `ClassDataSource` injects with the same `SharedType` (and key or class). While the test runs, it holds the instance, so TUnit does not dispose it before the test ends. This covers the test class constructor, the test method, and `[Before(Test)]` and `[After(Test)]` hooks: an instance fetched in any of them is held by that test, and is disposed after the test ends if no other test holds it:
+
+<!-- doc-test-declaration -->
+```csharp
+public class Database : IAsyncDisposable
+{
+    public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+}
+
+public class ReportTests
+{
+    [Test]
+    public async Task Reads_From_Shared_Database()
+    {
+        var database = SharedDataSources.GetOrCreate<Database>(SharedType.PerTestSession);
+
+        await Assert.That(database).IsNotNull();
+    }
+}
+```
+
+`GetOrCreate` does not call `IAsyncInitializer.InitializeAsync`. Only instances that TUnit injects are initialized.
+
+Outside a running test, nothing holds a fetched instance. This includes `[Before(Class)]`, `[After(Class)]`, `[Before(Assembly)]`, `[After(Assembly)]`, `[Before(TestSession)]` and `[After(TestSession)]` hooks, and background code. TUnit disposes it when the last test that holds it ends, which can be while you still use it. If no test holds it, TUnit does not dispose it. To keep a shared instance for the whole session, inject it into a static property with `ClassDataSource` instead.
+
 # Class Data Source Overloads
 
 If you are using an overload that supports injecting multiple classes at once (e.g. `ClassDataSource<T1, T2, T3>`) then you should specify multiple SharedTypes in an array and keys where applicable.

src/TUnit.Core/SharedDataSources.cs

@@ -1,6 +1,7 @@
 using System.Diagnostics.CodeAnalysis;
 using TUnit.Core.Helpers;
 using TUnit.Core.Interfaces;
+using TUnit.Core.Tracking;
 
 namespace TUnit.Core;
 
@@ -64,15 +65,12 @@ public static T GetOrCreate<T>(SharedType sharedType, Type? testClassType, strin
     {
         _ = factory ?? throw new ArgumentNullException(nameof(factory));
 
-        return sharedType switch
+        if (sharedType == SharedType.None)
         {
-            SharedType.None => factory(),
-            SharedType.PerTestSession => (T)TestDataContainer.GetGlobalInstance(typeof(T), _ => factory()!)!,
-            SharedType.PerClass => GetForClass<T>(testClassType, factory),
-            SharedType.Keyed => GetForKey<T>(key, factory),
-            SharedType.PerAssembly => GetForAssembly<T>(testClassType, factory),
-            _ => throw new ArgumentOutOfRangeException(nameof(sharedType), sharedType, "Invalid SharedType value.")
-        };
+            return factory();
+        }
+
+        return (T)GetShared(sharedType, typeof(T), testClassType, key, () => factory())!;
     }
 
     /// <summary>
@@ -144,9 +142,42 @@ public static T GetOrCreate<T>(SharedType sharedType, DataGeneratorMetadata data
         _ = type ?? throw new ArgumentNullException(nameof(type));
         _ = factory ?? throw new ArgumentNullException(nameof(factory));
 
+        if (sharedType == SharedType.None)
+        {
+            return factory();
+        }
+
+        return GetShared(sharedType, type, testClassType, key, factory);
+    }
+
+    /// <summary>
+    /// Fetches a shared instance. Inside a running test, the test holds a reference to it until the
+    /// test ends, so it is not disposed while the test uses it, and is disposed afterwards if no
+    /// other test holds it. If the instance was disposed before the test could hold it, a new one
+    /// is created.
+    /// </summary>
+    private static object? GetShared(SharedType sharedType, Type type, Type? testClassType, string? key, Func<object?> factory)
+    {
+        var instance = Resolve(sharedType, type, testClassType, key, factory);
+
+        if (instance is null || TestContext.Current is not { } testContext)
+        {
+            return instance;
+        }
+
+        // Fails only when the instance has just been retired, which also evicts it from the cache.
+        while (!ObjectTracker.TryHoldForRunningTest(testContext, instance))
+        {
+            instance = Resolve(sharedType, type, testClassType, key, factory)!;
+        }
+
+        return instance;
+    }
+
+    private static object? Resolve(SharedType sharedType, Type type, Type? testClassType, string? key, Func<object?> factory)
+    {
         return sharedType switch
         {
-            SharedType.None => factory(),
             SharedType.PerTestSession => TestDataContainer.GetGlobalInstance(type, _ => factory()!),
             SharedType.PerClass => GetForClass(type, testClassType, factory),
             SharedType.Keyed => GetForKey(type, key, factory),
@@ -155,16 +186,6 @@ public static T GetOrCreate<T>(SharedType sharedType, DataGeneratorMetadata data
         };
     }
 
-    private static T GetForClass<T>(Type? testClassType, Func<T> factory)
-    {
-        if (testClassType is null)
-        {
-            throw new ArgumentNullException(nameof(testClassType), "testClassType is required when SharedType is PerClass.");
-        }
-
-        return (T)TestDataContainer.GetInstanceForClass(testClassType, typeof(T), _ => factory()!)!;
-    }
-
     private static object? GetForClass(Type type, Type? testClassType, Func<object?> factory)
     {
         if (testClassType is null)
@@ -175,26 +196,6 @@ private static T GetForClass<T>(Type? testClassType, Func<T> factory)
         return TestDataContainer.GetInstanceForClass(testClassType, type, _ => factory()!);
     }
 
-    private static T GetForKey<T>(string? key, Func<T> factory)
-    {
-        if (string.IsNullOrEmpty(key))
-        {
-            throw new ArgumentNullException(nameof(key), "key is required when SharedType is Keyed.");
-        }
-
-        return (T)TestDataContainer.GetInstanceForKey(key!, typeof(T), _ =>
-        {
-            var instance = factory()!;
-
-            if (instance is IKeyedDataSource keyed)
-            {
-                keyed.Key = key!;
-            }
-
-            return instance;
-        })!;
-    }
-
     private static object? GetForKey(Type type, string? key, Func<object?> factory)
     {
         if (string.IsNullOrEmpty(key))
@@ -215,16 +216,6 @@ private static T GetForKey<T>(string? key, Func<T> factory)
         });
     }
 
-    private static T GetForAssembly<T>(Type? testClassType, Func<T> factory)
-    {
-        if (testClassType is null)
-        {
-            throw new ArgumentNullException(nameof(testClassType), "testClassType is required when SharedType is PerAssembly.");
-        }
-
-        return (T)TestDataContainer.GetInstanceForAssembly(testClassType.Assembly, typeof(T), _ => factory()!)!;
-    }
-
     private static object? GetForAssembly(Type type, Type? testClassType, Func<object?> factory)
     {
         if (testClassType is null)

src/TUnit.Core/TestContext.cs

@@ -495,6 +495,62 @@ internal SortedList<int, HashSet<object>> TrackedObjects
 
     internal SortedList<int, HashSet<object>>? TrackedObjectsIfCreated => Volatile.Read(ref _trackedObjects);
 
+    private int _holdsSharedReferences;
+
+    /// <summary>
+    /// True while the test's own lifecycle runs (from instance creation until its tracked objects
+    /// are untracked). Shared instances that the test fetches through <see cref="SharedDataSources"/>
+    /// during this window are held by the test (<see cref="Tracking.ObjectTracker.TryHoldForRunningTest"/>).
+    /// False during discovery and registration, where data-source objects are counted by the
+    /// object-graph walk instead.
+    /// </summary>
+    internal bool HoldsSharedReferences => Volatile.Read(ref _holdsSharedReferences) == 1;
+
+    internal void StartHoldingSharedReferences() => Volatile.Write(ref _holdsSharedReferences, 1);
+
+    private HashSet<object>? _heldSharedInstances;
+
+    /// <summary>
+    /// Shared instances the test fetched itself through <see cref="SharedDataSources"/> while
+    /// <see cref="HoldsSharedReferences"/> was true. Each holds one reference, released when the
+    /// test's tracked objects are untracked. Kept apart from <see cref="TrackedObjects"/> so the
+    /// engine does not initialize them and its initialization walk never sees the set change.
+    /// Lock on the set itself to read or change it.
+    /// </summary>
+    internal HashSet<object> HeldSharedInstances
+    {
+        get
+        {
+            if (Volatile.Read(ref _heldSharedInstances) is { } existing)
+            {
+                return existing;
+            }
+
+            var created = new HashSet<object>(Helpers.ReferenceEqualityComparer.Instance);
+            return Interlocked.CompareExchange(ref _heldSharedInstances, created, null) ?? created;
+        }
+    }
+
+    internal HashSet<object>? HeldSharedInstancesIfCreated => Volatile.Read(ref _heldSharedInstances);
+
+    /// <summary>
+    /// Stops new shared references being added, and waits for any add in progress, so the held
+    /// instances can be untracked safely.
+    /// </summary>
+    internal void StopHoldingSharedReferences()
+    {
+        // Full fence: either an adder that creates the held set after this sees the flag
+        // cleared, or this sees the set and waits on its lock.
+        Interlocked.Exchange(ref _holdsSharedReferences, 0);
+
+        if (HeldSharedInstancesIfCreated is { } heldSharedInstances)
+        {
+            lock (heldSharedInstances)
+            {
+            }
+        }
+    }
+
     /// <summary>
     /// Tracked objects that are reachable only from data rows the test never receives
     /// (<see cref="TestDetails.UnreceivedDataRowObjects"/>). They are tracked for reference-counted

src/TUnit.Core/TestDataContainer.cs

@@ -1,33 +1,60 @@
-using System.Reflection;
-using TUnit.Core.Data;
+using System.Collections.Concurrent;
+using System.Reflection;
 
 namespace TUnit.Core;
 
 internal static class TestDataContainer
 {
-    private static readonly ScopedDictionary<string> _globalContainer = new();
-    private static readonly ScopedDictionary<Type> _classContainer = new();
-    private static readonly ScopedDictionary<Assembly> _assemblyContainer = new();
-    private static readonly ScopedDictionary<string> _keyContainer = new();
+    private static readonly object GlobalScope = new();
+
+    private static readonly ConcurrentDictionary<InstanceKey, Lazy<object?>> _instances = new();
+
+    // Disposable cached instances only, by identity, so that a disposed instance can be evicted.
+    // Without eviction, a test registered after the last holder has disposed a shared instance
+    // (e.g. a dynamic test) would be handed the disposed instance (#7204). A factory can return
+    // the same instance for several keys or types, so every entry holding it is recorded.
+    private static readonly ConcurrentDictionary<object, KeyValuePair<InstanceKey, Lazy<object?>>[]> _entriesByInstance =
+        new(Helpers.ReferenceEqualityComparer.Instance);
 
     public static object? GetInstanceForClass(Type testClass, Type type, Func<Type, object> func)
     {
-        return _classContainer.GetOrCreate(testClass, type, func);
+        return GetOrCreate(new InstanceKey(SharedType.PerClass, testClass, type), func);
     }
 
     public static object? GetInstanceForAssembly(Assembly assembly, Type type, Func<Type, object> func)
     {
-        return _assemblyContainer.GetOrCreate(assembly, type, func);
+        return GetOrCreate(new InstanceKey(SharedType.PerAssembly, assembly, type), func);
     }
 
     public static object? GetGlobalInstance(Type type, Func<Type, object> func)
     {
-        return _globalContainer.GetOrCreate(typeof(object).FullName!, type, func);
+        return GetOrCreate(new InstanceKey(SharedType.PerTestSession, GlobalScope, type), func);
     }
 
     public static object? GetInstanceForKey(string key, Type type, Func<Type, object> func)
     {
-        return _keyContainer.GetOrCreate(key, type, func);
+        return GetOrCreate(new InstanceKey(SharedType.Keyed, key, type), func);
+    }
+
+    /// <summary>
+    /// Whether <paramref name="instance"/> is a disposable instance that is still cached.
+    /// </summary>
+    public static bool IsCached(object instance) => _entriesByInstance.ContainsKey(instance);
+
+    /// <summary>
+    /// Removes <paramref name="instance"/> from the cache, so the next request for any scope and
+    /// type it is cached under creates a new instance. Does nothing if it is not cached.
+    /// </summary>
+    public static void Evict(object instance)
+    {
+        if (_entriesByInstance.TryRemove(instance, out var entries))
+        {
+            foreach (var entry in entries)
+            {
+                // Removes the entry only while it still holds this instance's Lazy (reference equality).
+                ((ICollection<KeyValuePair<InstanceKey, Lazy<object?>>>) _instances).Remove(entry);
+            }
+        }
     }
 
     /// <summary>
@@ -37,9 +64,58 @@ internal static class TestDataContainer
     /// </summary>
     public static void Reset()
     {
-        _globalContainer.Clear();
-        _classContainer.Clear();
-        _assemblyContainer.Clear();
-        _keyContainer.Clear();
+        _instances.Clear();
+        _entriesByInstance.Clear();
+    }
+
+    private static object? GetOrCreate(InstanceKey key, Func<Type, object> func)
+    {
+        if (_instances.TryGetValue(key, out var existing))
+        {
+            return existing.Value;
+        }
+
+        return AddSlow(key, func);
+    }
+
+    // Kept separate from GetOrCreate so the fast path does not allocate the closure.
+    private static object? AddSlow(InstanceKey key, Func<Type, object> func)
+    {
+        Lazy<object?> lazy = null!;
+        lazy = new Lazy<object?>(() =>
+        {
+            var instance = func(key.Type);
+
+            if (instance is IDisposable or IAsyncDisposable)
+            {
+                var entry = new KeyValuePair<InstanceKey, Lazy<object?>>(key, lazy);
+                _entriesByInstance.AddOrUpdate(instance, _ => [entry], (_, entries) => [..entries, entry]);
+            }
+
+            return instance;
+        }, LazyThreadSafetyMode.ExecutionAndPublication);
+
+        return _instances.GetOrAdd(key, lazy).Value;
+    }
+
+    private readonly struct InstanceKey(SharedType kind, object scope, Type type) : IEquatable<InstanceKey>
+    {
+        public SharedType Kind { get; } = kind;
+
+        public object Scope { get; } = scope;
+
+        public Type Type { get; } = type;
+
+        public bool Equals(InstanceKey other) => Kind == other.Kind && Type == other.Type && Scope.Equals(other.Scope);
+
+        public override bool Equals(object? obj) => obj is InstanceKey other && Equals(other);
+
+        public override int GetHashCode()
+        {
+            unchecked
+            {
+                return ((int) Kind * 397 ^ Scope.GetHashCode()) * 397 ^ Type.GetHashCode();
+            }
+        }
     }
 }