perf: Optimize concurrency gate with sync TryEnter fast path (#89)

Author: rian-be

src/Runtime/Internal/Execution/MutationExecutionConcurrencyGate.cs

@@ -1,30 +1,130 @@
 using System.Collections.Concurrent;
+using System.Runtime.CompilerServices;
 
 namespace ModularityKit.Mutator.Runtime.Internal.Execution;
 
 /// <summary>
-/// Coordinates core mutation execution concurrency across the engine.
+/// Coordinates concurrent mutation execution using global concurrency limit
+/// and per state serialization.
 /// </summary>
-internal sealed class MutationExecutionConcurrencyGate(int maxConcurrentMutations)
+/// <remarks>
+/// When state identifier is provided, an additional per state gate ensures that
+/// mutations targeting the same state are executed sequentially, while mutations
+/// targeting different states may execute concurrently up to the global limit.
+/// </remarks>
+/// <param name="maxConcurrentMutations">The maximum number of mutations that may execute concurrently across all states. </param>
+internal sealed class MutationExecutionConcurrencyGate(
+    int maxConcurrentMutations)
 {
-    private readonly SemaphoreSlim _globalGate = new(maxConcurrentMutations, maxConcurrentMutations);
-    private readonly ConcurrentDictionary<string, SemaphoreSlim> _stateGates = new(StringComparer.Ordinal);
+    private readonly SemaphoreSlim _globalGate =
+        new(maxConcurrentMutations, maxConcurrentMutations);
 
-    public async ValueTask<Lease> EnterAsync(string? stateId, CancellationToken cancellationToken)
+    private readonly ConcurrentDictionary<string, SemaphoreSlim> _stateGates =
+        new(StringComparer.Ordinal);
+
+    /// <summary>
+    /// Enters the concurrency gate for mutation execution.
+    /// </summary>
+    /// <param name="stateId">The identifier of the state being mutated.</param>
+    /// <param name="cancellationToken">Token that can be used to cancel waiting for the required concurrency gates. </param>
+    /// <returns>An asynchronous lease that releases the acquired concurrency gates when disposed. </returns>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public ValueTask<Lease> EnterAsync(
+        string? stateId,
+        CancellationToken cancellationToken)
     {
-        await _globalGate.WaitAsync(cancellationToken).ConfigureAwait(false);
+        if (stateId is null)
+            return EnterGlobalOnlyAsync(cancellationToken);
 
-        var stateGate = default(SemaphoreSlim);
+        return EnterWithStateAsync(stateId, cancellationToken);
+    }
 
-        try
+    /// <summary>
+    /// Attempts to acquire the global concurrency gate without asynchronous waiting
+    /// when capacity is immediately available.
+    /// </summary>
+    private ValueTask<Lease> EnterGlobalOnlyAsync(
+        CancellationToken cancellationToken)
+    {
+        if (_globalGate.Wait(0, cancellationToken))
+        {
+            return new ValueTask<Lease>(
+                new Lease(_globalGate, null));
+        }
+
+        return SlowEnterGlobalOnlyAsync(cancellationToken);
+    }
+
+    /// <summary>
+    /// Asynchronously waits for the global concurrency gate when it could not be
+    /// acquired immediately.
+    /// </summary>
+    private async ValueTask<Lease> SlowEnterGlobalOnlyAsync(
+        CancellationToken cancellationToken)
+    {
+        await _globalGate
+            .WaitAsync(cancellationToken)
+            .ConfigureAwait(false);
+
+        return new Lease(_globalGate, null);
+    }
+
+    /// <summary>
+    /// Attempts to acquire both the global concurrency gate and the gate associated
+    /// with the specified state.
+    /// </summary>
+    /// <remarks>
+    /// The global gate is always acquired before the per state gate to maintain
+    /// consistent acquisition order and avoid lock order inversion.
+    /// </remarks>
+    private ValueTask<Lease> EnterWithStateAsync(
+        string stateId,
+        CancellationToken cancellationToken)
+    {
+        if (_stateGates.TryGetValue(stateId, out var existing))
         {
-            if (!string.IsNullOrWhiteSpace(stateId))
+            if (_globalGate.Wait(0, cancellationToken))
             {
-                stateGate = _stateGates.GetOrAdd(stateId, static _ => new SemaphoreSlim(1, 1));
-                await stateGate.WaitAsync(cancellationToken).ConfigureAwait(false);
+                if (existing.Wait(0, cancellationToken))
+                {
+                    return new ValueTask<Lease>(
+                        new Lease(_globalGate, existing));
+                }
+
+                _globalGate.Release();
             }
+        }
+
+        return SlowEnterWithStateAsync(
+            stateId,
+            cancellationToken);
+    }
+
+    /// <summary>
+    /// Asynchronously acquires the global and per state gates when the fast path
+    /// could not acquire them immediately.
+    /// </summary>
+    private async ValueTask<Lease> SlowEnterWithStateAsync(
+        string stateId,
+        CancellationToken cancellationToken)
+    {
+        await _globalGate
+            .WaitAsync(cancellationToken)
+            .ConfigureAwait(false);
 
-            return new Lease(_globalGate, stateGate);
+        try
+        {
+            var stateGate = _stateGates.GetOrAdd(
+                stateId,
+                static _ => new SemaphoreSlim(1, 1));
+
+            await stateGate
+                .WaitAsync(cancellationToken)
+                .ConfigureAwait(false);
+
+            return new Lease(
+                _globalGate,
+                stateGate);
         }
         catch
         {
@@ -34,24 +134,23 @@ public async ValueTask<Lease> EnterAsync(string? stateId, CancellationToken canc
     }
 
     /// <summary>
-    /// Represents an acquired execution slot.
+    /// Represents an acquired concurrency lease that releases the associated gates
+    /// when disposed.
     /// </summary>
-    internal readonly struct Lease : IAsyncDisposable
+    internal readonly struct Lease(
+        SemaphoreSlim globalGate,
+        SemaphoreSlim? stateGate) : IAsyncDisposable
     {
-        private readonly SemaphoreSlim _globalGate;
-        private readonly SemaphoreSlim? _stateGate;
-
-        public Lease(SemaphoreSlim globalGate, SemaphoreSlim? stateGate)
-        {
-            _globalGate = globalGate;
-            _stateGate = stateGate;
-        }
-
+        /// <summary>
+        /// Releases the per state gate, when present, followed by the global gate.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public ValueTask DisposeAsync()
         {
-            _stateGate?.Release();
-            _globalGate.Release();
+            stateGate?.Release();
+            globalGate.Release();
+
             return ValueTask.CompletedTask;
         }
     }
-}
+}