Fixed node removal logic

Author: sakno

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

@@ -37,20 +37,25 @@ private protected WaitQueueVisitor GetWaitQueue(ref LinkedValueTaskCompletionSou
         return detachedQueue.First;
     }
 
-    private void ReturnNode(WaitNode node)
+    private void ReturnNode<TStrategy>(TStrategy strategy)
+        where TStrategy : struct, IRemovalStrategy
     {
+        var returnToPool = strategy.Node.TryReset(out _);
         LinkedValueTaskCompletionSource<bool>? suspendedCallers;
         lock (SyncRoot)
         {
-            suspendedCallers = node.NeedsRemoval && waitQueue.Remove(node) && node.DrainOnReturn
+            suspendedCallers = strategy.Remove(ref waitQueue)
                 ? DrainWaitQueue()
                 : null;
 
-            pool.Return(node);
+            if (returnToPool)
+                pool.Return(strategy.Node);
         }
 
         suspendedCallers?.Unwind();
     }
+    
+    
 
     private protected TNode? Acquire<T, TBuilder, TNode>(ref TBuilder builder, bool acquired)
         where T : struct, IEquatable<T>
@@ -268,4 +273,11 @@ public void Add(LinkedValueTaskCompletionSource<bool> node)
             return result;
         }
     }
+    
+    private interface IRemovalStrategy
+    {
+        bool Remove(ref WaitQueue queue);
+        
+        LinkedValueTaskCompletionSource<bool> Node { get; }
+    }
 }

src/DotNext.Threading/Threading/QueuedSynchronizer.WaitNode.cs

@@ -1,4 +1,5 @@
 using System.Diagnostics;
+using System.Runtime.InteropServices;
 using DotNext.Patterns;
 
 namespace DotNext.Threading;
@@ -41,12 +42,55 @@ internal bool DrainOnReturn
 
         protected sealed override void AfterConsumed()
         {
-            if (owner is { } ownerCopy && TryReset(out _))
-                ownerCopy.ReturnNode(this);
+            if (owner is null)
+            {
+                // nothing to do
+            }
+            else if (!NeedsRemoval)
+            {
+                owner.ReturnNode<DoNotRemoveStrategy>(new(this));
+            }
+            else if (DrainOnReturn)
+            {
+                owner.ReturnNode<RemoveAndDrainStrategy>(new(this));
+            }
+            else
+            {
+                owner.ReturnNode<RemoveStrategy>(new(this));
+            }
         }
 
         protected sealed override Result<bool> OnTimeout()
             => (flags & WaitNodeFlags.ThrowOnTimeout) is not 0 ? base.OnTimeout() : false;
+
+        [StructLayout(LayoutKind.Auto)]
+        private readonly struct DoNotRemoveStrategy(WaitNode node) : IRemovalStrategy
+        {
+            bool IRemovalStrategy.Remove(ref WaitQueue queue) => false;
+
+            LinkedValueTaskCompletionSource<bool> IRemovalStrategy.Node => node;
+        }
+
+        [StructLayout(LayoutKind.Auto)]
+        private readonly struct RemoveStrategy(WaitNode node) : IRemovalStrategy
+        {
+            bool IRemovalStrategy.Remove(ref WaitQueue queue)
+            {
+                queue.Remove(node);
+                return false;
+            }
+
+            LinkedValueTaskCompletionSource<bool> IRemovalStrategy.Node => node;
+        }
+
+        [StructLayout(LayoutKind.Auto)]
+        private readonly struct RemoveAndDrainStrategy(WaitNode node) : IRemovalStrategy
+        {
+            bool IRemovalStrategy.Remove(ref WaitQueue queue)
+                => queue.Remove(node);
+
+            LinkedValueTaskCompletionSource<bool> IRemovalStrategy.Node => node;
+        }
     }
 
     [Flags]