Run feature factories outside the OperationFeatureCollection write lock

src/HotChocolate/Core/src/Types/Execution/Processing/OperationFeatureCollection.Selections.cs

@@ -1,3 +1,4 @@
+using System.Collections.Concurrent;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 
@@ -6,6 +7,10 @@ namespace HotChocolate.Execution.Processing;
 [SuppressMessage("ReSharper", "NonAtomicCompoundOperator")]
 public sealed partial class OperationFeatureCollection
 {
+    // One creation gate per feature, so a feature is still built exactly once while the build itself
+    // runs outside the write lock. Lives as long as the operation, like the features themselves.
+    private readonly ConcurrentDictionary<(int, Type), object> _selectionFactories = new();
+
     internal object? this[int selectionId, Type featureType]
     {
         get
@@ -52,19 +57,47 @@ internal TFeature GetOrSetSafe<TFeature>(int selectionId, Func<TFeature> factory
     {
         ArgumentNullException.ThrowIfNull(factory);
 
-        if (!TryGet<TFeature>(selectionId, out var feature))
+        if (TryGet<TFeature>(selectionId, out var feature))
         {
-            lock (_writeLock)
+            return feature;
+        }
+
+        // The factory must not run while the write lock is held. Feature factories reach back into the
+        // operation (building a projection selector calls Operation.GetSelectionSet, which takes the
+        // operation lock), while compiling a selection set holds the operation lock and writes selection
+        // features here. Invoking the factory under this lock lets the two be taken in opposite orders,
+        // which deadlocks the executor.
+        //
+        // A single lock guards every feature of the operation, so those two are only ever in each
+        // other's way because they share it. Gating creation per feature keeps a costly value built
+        // exactly once without holding anything while it is built.
+        var created = GetOrCreate(selectionId, factory);
+
+        lock (_writeLock)
+        {
+            if (TryGet<TFeature>(selectionId, out var existing))
             {
-                if (!TryGet(selectionId, out feature))
-                {
-                    feature = factory();
-                    this[selectionId, typeof(TFeature)] = feature;
-                }
+                return existing;
             }
+
+            this[selectionId, typeof(TFeature)] = created;
+            return created;
         }
+    }
 
-        return feature;
+    private TFeature GetOrCreate<TFeature>(int selectionId, Func<TFeature> factory)
+    {
+        // Lazy<object> rather than Lazy<TFeature>: the value is stored as object anyway, and the
+        // trimmer requires a parameterless constructor on the type argument, which a feature need not
+        // have.
+        var lazy = (Lazy<object>)_selectionFactories.GetOrAdd(
+            (selectionId, typeof(TFeature)),
+            static (_, f) => new Lazy<object>(() => f()!, LazyThreadSafetyMode.ExecutionAndPublication),
+            factory);
+
+        // Racing callers may each build a Lazy, but only the one that won the slot is ever evaluated,
+        // so the factory runs once per feature.
+        return (TFeature)lazy.Value;
     }
 
     internal TFeature GetOrSetSafe<TFeature, TContext>(
@@ -74,19 +107,24 @@ internal TFeature GetOrSetSafe<TFeature, TContext>(
     {
         ArgumentNullException.ThrowIfNull(factory);
 
-        if (!TryGet<TFeature>(selectionId, out var feature))
+        if (TryGet<TFeature>(selectionId, out var feature))
         {
-            lock (_writeLock)
+            return feature;
+        }
+
+        // Built outside the lock and once per feature, see the overload above.
+        var created = GetOrCreate(selectionId, () => factory(context));
+
+        lock (_writeLock)
+        {
+            if (TryGet<TFeature>(selectionId, out var existing))
             {
-                if (!TryGet(selectionId, out feature))
-                {
-                    feature = factory(context);
-                    this[selectionId, typeof(TFeature)] = feature;
-                }
+                return existing;
             }
-        }
 
-        return feature;
+            this[selectionId, typeof(TFeature)] = created;
+            return created;
+        }
     }
 
     internal IEnumerable<KeyValuePair<Type, object>> GetFeatures(int selectionId)

src/HotChocolate/Core/src/Types/Execution/Processing/OperationFeatureCollection.cs

@@ -1,4 +1,5 @@
 using System.Collections;
+using System.Collections.Concurrent;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using HotChocolate.Features;
@@ -22,6 +23,10 @@ public sealed partial class OperationFeatureCollection : IFeatureCollection
 #endif
     private volatile int _containerRevision;
 
+    // One creation gate per feature, so a feature is still built exactly once while the build itself
+    // runs outside the write lock.
+    private readonly ConcurrentDictionary<Type, object> _factories = new();
+
     /// <summary>
     /// Initializes a new instance of <see cref="FeatureCollection"/>.
     /// </summary>
@@ -94,19 +99,26 @@ public TFeature GetOrSetSafe<TFeature>(Func<TFeature> factory)
     {
         ArgumentNullException.ThrowIfNull(factory);
 
-        if (!TryGet<TFeature>(out var feature))
+        if (TryGet<TFeature>(out var feature))
         {
-            lock (_writeLock)
+            return feature;
+        }
+
+        // Built outside the lock, so a factory that reaches back into the operation cannot deadlock
+        // against a thread holding the operation lock and writing features here, and gated per feature
+        // so a costly value is still built only once.
+        var created = GetOrCreate(factory);
+
+        lock (_writeLock)
+        {
+            if (TryGet<TFeature>(out var existing))
             {
-                if (!TryGet(out feature))
-                {
-                    feature = factory();
-                    this[typeof(TFeature)] = feature;
-                }
+                return existing;
             }
-        }
 
-        return feature;
+            this[typeof(TFeature)] = created;
+            return created;
+        }
     }
 
     internal TFeature GetOrSetSafe<TFeature, TContext>(
@@ -115,19 +127,37 @@ internal TFeature GetOrSetSafe<TFeature, TContext>(
     {
         ArgumentNullException.ThrowIfNull(factory);
 
-        if (!TryGet<TFeature>(out var feature))
+        if (TryGet<TFeature>(out var feature))
         {
-            lock (_writeLock)
+            return feature;
+        }
+
+        // Built outside the lock and once per feature, see the overload above.
+        var created = GetOrCreate(() => factory(context));
+
+        lock (_writeLock)
+        {
+            if (TryGet<TFeature>(out var existing))
             {
-                if (!TryGet(out feature))
-                {
-                    feature = factory(context);
-                    this[typeof(TFeature)] = feature;
-                }
+                return existing;
             }
+
+            this[typeof(TFeature)] = created;
+            return created;
         }
+    }
 
-        return feature;
+    private TFeature GetOrCreate<TFeature>(Func<TFeature> factory)
+    {
+        // Lazy<object> rather than Lazy<TFeature>: the value is stored as object anyway, and the
+        // trimmer requires a parameterless constructor on the type argument, which a feature need not
+        // have.
+        var lazy = (Lazy<object>)_factories.GetOrAdd(
+            typeof(TFeature),
+            static (_, f) => new Lazy<object>(() => f()!, LazyThreadSafetyMode.ExecutionAndPublication),
+            factory);
+
+        return (TFeature)lazy.Value;
     }
 
     /// <inheritdoc />

src/HotChocolate/Core/test/Execution.Tests/Processing/OperationFeatureCollectionTests.cs

@@ -0,0 +1,137 @@
+namespace HotChocolate.Execution.Processing;
+
+/// <summary>
+/// Feature factories reach back into the operation: building a projection selector expression calls
+/// <c>Operation.GetSelectionSet</c>, which takes the operation lock. Compiling a selection set runs the
+/// other way round, holding the operation lock while it writes selection features here. Invoking a
+/// factory while this collection's write lock is held lets the two locks be taken in opposite orders,
+/// which deadlocks the executor.
+/// </summary>
+public class OperationFeatureCollectionTests
+{
+    private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10);
+
+    [Fact]
+    public void GetOrSetSafe_Should_NotHoldWriteLock_When_SelectionFactoryRuns()
+    {
+        // arrange
+        var features = new OperationFeatureCollection();
+        var concurrentWriteFinished = new ManualResetEventSlim(false);
+
+        // act
+        // the write from the other thread takes the collection lock, so it only completes if the
+        // factory is not holding it
+        var feature = features.GetOrSetSafe(
+            selectionId: 1,
+            () =>
+            {
+                var writer = Task.Run(() =>
+                {
+                    features[2, typeof(OtherFeature)] = new OtherFeature();
+                    concurrentWriteFinished.Set();
+                });
+
+                // Asserted here, not after the factory returns: by then the lock is released and the
+                // writer completes regardless, so a late check passes even when the lock was held.
+                Assert.True(
+                    concurrentWriteFinished.Wait(s_timeout),
+                    "the concurrent write blocked, so the factory ran while holding the write lock");
+                Assert.True(writer.Wait(s_timeout), "the concurrent write did not finish");
+                return new Feature();
+            });
+
+        // assert
+        Assert.True(features.TryGet<Feature>(1, out var stored));
+        Assert.Same(feature, stored);
+    }
+
+    [Fact]
+    public void GetOrSetSafe_Should_NotHoldWriteLock_When_OperationFactoryRuns()
+    {
+        // arrange
+        var features = new OperationFeatureCollection();
+        var concurrentWriteFinished = new ManualResetEventSlim(false);
+
+        // act
+        var feature = features.GetOrSetSafe(() =>
+        {
+            var writer = Task.Run(() =>
+            {
+                features[typeof(OtherFeature)] = new OtherFeature();
+                concurrentWriteFinished.Set();
+            });
+
+            // Asserted here, not after the factory returns: by then the lock is released and the
+            // writer completes regardless, so a late check passes even when the lock was held.
+            Assert.True(
+                concurrentWriteFinished.Wait(s_timeout),
+                "the concurrent write blocked, so the factory ran while holding the write lock");
+            Assert.True(writer.Wait(s_timeout), "the concurrent write did not finish");
+            return new Feature();
+        });
+
+        // assert
+        Assert.Same(feature, features.Get<Feature>());
+    }
+
+    [Fact]
+    public async Task GetOrSetSafe_Should_ReturnTheSameInstanceToEveryCaller_When_CallersRace()
+    {
+        // arrange
+        // creation is gated per feature, so racing callers share one build
+        var features = new OperationFeatureCollection();
+        var start = new ManualResetEventSlim(false);
+        var callers = Enumerable.Range(0, 8)
+            .Select(_ => Task.Run(() =>
+            {
+                start.Wait(s_timeout);
+                return features.GetOrSetSafe(selectionId: 1, () => new Feature());
+            }))
+            .ToArray();
+
+        // act
+        start.Set();
+        var features1 = await Task.WhenAll(callers)
+            .WaitAsync(s_timeout, TestContext.Current.CancellationToken);
+
+        // assert
+        Assert.True(features.TryGet<Feature>(1, out var stored));
+        Assert.All(features1, f => Assert.Same(stored, f));
+    }
+
+    [Fact]
+    public async Task GetOrSetSafe_Should_BuildTheFeatureOnce_When_CallersRace()
+    {
+        // arrange
+        // Feature creation is costly where this is used, so racing callers must not each pay for it.
+        var features = new OperationFeatureCollection();
+        var start = new ManualResetEventSlim(false);
+        var builds = 0;
+        var callers = Enumerable.Range(0, 16)
+            .Select(_ => Task.Run(() =>
+            {
+                start.Wait(s_timeout);
+                return features.GetOrSetSafe(
+                    selectionId: 1,
+                    () =>
+                    {
+                        Interlocked.Increment(ref builds);
+                        return new Feature();
+                    });
+            }))
+            .ToArray();
+
+        // act
+        start.Set();
+        var results = await Task.WhenAll(callers)
+            .WaitAsync(s_timeout, TestContext.Current.CancellationToken);
+
+        // assert
+        Assert.Equal(1, Volatile.Read(ref builds));
+        Assert.All(results, f => Assert.Same(results[0], f));
+    }
+
+    private sealed class Feature;
+
+    private sealed class OtherFeature;
+}