Reducd lock contention

Author: sakno

src/DotNext.Threading/Threading/AsyncExchanger.cs

@@ -29,11 +29,7 @@ internal void Initialize(AsyncExchanger<T> owner, T value)
             this.value = value;
         }
 
-        protected override void AfterConsumed()
-        {
-            if (owner is { } ownerCopy && TryReset(out _))
-                ownerCopy.OnCompleted(this);
-        }
+        protected override void AfterConsumed() => owner?.OnCompleted(this);
 
         protected override void CleanUp()
         {
@@ -80,7 +76,7 @@ private ExchangePoint RentExchangePoint(T value)
     {
         Debug.Assert(Monitor.IsEntered(SyncRoot));
 
-        var result = pool.Get<ExchangePoint>();
+        var result = pool.Rent<ExchangePoint>();
         result.Initialize(this, value);
         return result;
     }
@@ -89,10 +85,16 @@ private void OnCompleted(ExchangePoint point)
     {
         Monitor.Enter(SyncRoot);
         if (ReferenceEquals(this.point, point))
+        {
             this.point = null;
+        }
 
-        pool.Return(point);
         Monitor.Exit(SyncRoot);
+
+        if (point.TryReset(out _))
+        {
+            pool.Return(point);
+        }
     }
 
     /// <summary>

src/DotNext.Threading/Threading/QueuedSynchronizer.Queue.cs

@@ -43,34 +43,36 @@ private protected WaitQueueVisitor GetWaitQueue(ref LinkedValueTaskCompletionSou
     [MethodImpl(MethodImplOptions.NoInlining)]
     private void ReturnNode(WaitNode node)
     {
-        var syncRoot = SyncRoot;
         if (node.NeedsRemoval)
         {
-            var suspendedCallers = default(LinkedValueTaskCompletionSource<bool>);
-            Monitor.Enter(syncRoot);
-            try
-            {
-                suspendedCallers = waitQueue.Remove(node) && node.DrainOnReturn
-                    ? DrainWaitQueue()
-                    : null;
-            }
-            finally
-            {
-                Monitor.Exit(syncRoot);
-                suspendedCallers?.Unwind();
-            }
+            RemoveAndDrainIfNeeded(node);
         }
 
         // the node is removed for sure, it can be returned back to the pool
         if (node.TryReset(out _))
         {
-            lock (syncRoot)
-            {
-                pool.Return(node);
-            }
+            pool.Return(node);
         }
     }
-    
+
+    private void RemoveAndDrainIfNeeded(WaitNode node)
+    {
+        var syncRoot = SyncRoot;
+        var suspendedCallers = default(LinkedValueTaskCompletionSource<bool>);
+        Monitor.Enter(syncRoot);
+        try
+        {
+            suspendedCallers = waitQueue.Remove(node) && node.DrainOnReturn
+                ? DrainWaitQueue()
+                : null;
+        }
+        finally
+        {
+            Monitor.Exit(syncRoot);
+            suspendedCallers?.Unwind();
+        }
+    }
+
     private protected TNode? Acquire<T, TBuilder, TNode>(ref TBuilder builder, bool acquired)
         where T : struct, IEquatable<T>
         where TNode : WaitNode, new()
@@ -97,7 +99,7 @@ private void ReturnNode(WaitNode node)
         }
         else
         {
-            var node = pool.Get<TNode>();
+            var node = pool.Rent<TNode>();
             node.Initialize(this, CaptureCallerInformation(), TBuilder.ThrowOnTimeout);
             waitQueue.Add(node);
             builder.Complete(node);

src/DotNext.Threading/Threading/Tasks/LinkedValueTaskCompletionSource.cs

@@ -10,7 +10,7 @@ private protected LinkedValueTaskCompletionSource(bool runContinuationsAsynchron
     {
     }
 
-    internal LinkedValueTaskCompletionSource<T>? Next { get; private set; }
+    internal LinkedValueTaskCompletionSource<T>? Next { get; set; }
 
     internal LinkedValueTaskCompletionSource<T>? Previous { get; private set; }
 
@@ -22,14 +22,6 @@ internal void Append(LinkedValueTaskCompletionSource<T> node)
         Next = node;
     }
 
-    internal void Prepend(LinkedValueTaskCompletionSource<T> node)
-    {
-        Debug.Assert(Previous is null);
-
-        node.Next = this;
-        Previous = node;
-    }
-
     internal void Detach()
     {
         if (Previous is not null)

src/DotNext.Threading/Threading/Tasks/Pooling/ValueTaskPool.cs

@@ -7,15 +7,13 @@ namespace DotNext.Threading.Tasks.Pooling;
 /*
  * Represents a pool without any allocations. Assuming that wait queues are organized using
  * linked nodes where each node is a completion source. If so, the node pointers (next/previous)
- * can be used to keep completion sources in the pool. Moreover, access to the pool is synchronized
- * by the caller (see QueuedSynchronizer). Thus, it doesn't require explicit synchronization such as
- * Monitor lock.
+ * can be used to keep completion sources in the pool. The implementation is thread-safe.
  */
 [StructLayout(LayoutKind.Auto)]
 internal struct ValueTaskPool<T>(long maximumRetained)
 {
-    private LinkedValueTaskCompletionSource<T>? first;
-    private long count;
+    private volatile LinkedValueTaskCompletionSource<T>? first;
+    private long count; // volatile
 
     public ValueTaskPool()
         : this(long.MaxValue)
@@ -26,41 +24,57 @@ public ValueTaskPool()
 
     public void Return(LinkedValueTaskCompletionSource<T> node)
     {
-        Debug.Assert(node is not null);
-        Debug.Assert(count is 0L || first is not null);
-        Debug.Assert(node.Status is ManualResetCompletionSourceStatus.WaitForActivation);
+        Debug.Assert(node is { Status: ManualResetCompletionSourceStatus.WaitForActivation });
+        Debug.Assert(node is { Next: null, Previous: null });
 
-        if (count < maximumRetained)
+        // try to increment the counter
+        for (long current = Atomic.Read(in count), tmp; current < maximumRetained; current = tmp)
         {
-            first?.Prepend(node);
-            first = node;
-            count++;
-            Debug.Assert(count > 0L);
+            tmp = Interlocked.CompareExchange(ref count, current + 1L, current);
+            if (tmp == current)
+            {
+                ReturnCore(node);
+                break;
+            }
         }
     }
 
-    public TNode Get<TNode>()
-        where TNode : LinkedValueTaskCompletionSource<T>, new()
+    private void ReturnCore(LinkedValueTaskCompletionSource<T> node)
     {
-        TNode result;
-        if (first is null)
+        for (LinkedValueTaskCompletionSource<T>? current = first, tmp;; current = tmp)
         {
-            Debug.Assert(count == 0L);
+            node.Next = current;
 
-            result = new();
+            tmp = Interlocked.CompareExchange(ref first, node, current);
+            if (ReferenceEquals(tmp, current))
+                break;
         }
-        else
+    }
+
+    public TNode Rent<TNode>()
+        where TNode : LinkedValueTaskCompletionSource<T>, new()
+    {
+        var current = first;
+        for (LinkedValueTaskCompletionSource<T>? tmp;; current = Unsafe.As<TNode>(tmp))
         {
-            Debug.Assert(first is TNode);
-            result = Unsafe.As<TNode>(first);
-            first = result.Next;
-            result.Detach();
-            count--;
+            if (current is null)
+            {
+                current = new TNode();
+                break;
+            }
+
+            tmp = Interlocked.CompareExchange(ref first, current.Next, current);
+            if (!ReferenceEquals(tmp, current))
+                continue;
 
-            Debug.Assert(count >= 0L);
-            Debug.Assert(count is 0L || first is not null);
+            current.Next = null;
+            var actualCount = Interlocked.Decrement(ref count);
+            Debug.Assert(actualCount >= 0L);
+            break;
         }
 
-        return result;
+        Debug.Assert(current is TNode);
+        Debug.Assert(current is { Next: null, Previous: null });
+        return Unsafe.As<TNode>(current);
     }
 }

src/DotNext.Threading/Threading/Tasks/TaskCompletionPipe.Signal.cs

@@ -45,7 +45,7 @@ private LinkedValueTaskCompletionSource<bool> EnqueueNode()
     {
         Debug.Assert(Monitor.IsEntered(SyncRoot));
 
-        var result = pool.Get<Signal>();
+        var result = pool.Rent<Signal>();
         result.Initialize(this);
         waitQueue.Add(result);
         return result;

src/DotNext.Threading/Threading/Tasks/TaskCompletionPipe.cs

@@ -46,11 +46,16 @@ public bool IsCompletionTaskSupported
 
     private void OnCompleted(Signal signal)
     {
-        lock (SyncRoot)
+        if (signal.NeedsRemoval)
         {
-            if (signal.NeedsRemoval)
+            lock (SyncRoot)
+            {
                 waitQueue.Remove(signal);
+            }
+        }
 
+        if (signal.TryReset(out _))
+        {
             pool.Return(signal);
         }
     }