Optimized object pool by checking for early-return condition when the…

…re is contention.
timcassell · Dec 3, 2024 · fb994da · fb994da
1 parent 36dd532
commit fb994da
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/Package/Core/InternalShared/SpinLocker.cs b/Package/Core/InternalShared/SpinLocker.cs
@@ -39,7 +39,7 @@ internal void Enter()
             [MethodImpl(MethodImplOptions.NoInlining)]
             private void EnterCore()
             {
-                // Spin until we successfully get lock.
+                // Spin until we acquire the lock.
                 var spinner = new SpinWait();
                 do
                 {
@@ -64,7 +64,7 @@ internal void EnterWithoutSleep1()
             [MethodImpl(MethodImplOptions.NoInlining)]
             private void EnterWithoutSleep1Core()
             {
-                // Spin until we successfully get lock.
+                // Spin until we acquire the lock.
                 var spinner = new SpinWait();
                 do
                 {

diff --git a/Package/Core/InternalShared/ValueCollectionsInternal.cs b/Package/Core/InternalShared/ValueCollectionsInternal.cs
@@ -215,7 +215,7 @@ internal ValueLinkedStackSafe(HandleablePromiseBase tailSentinel)
             [MethodImpl(InlineOption)]
             internal void ClearUnsafe()
             {
-                // Worst case scenario, ClearUnsafe() is called concurrently with Push() and/or TryPop() and the objects are re-pooled.
+                // Worst case scenario, ClearUnsafe() is called concurrently with Push() and/or PopOrInvalid() and the objects are re-pooled.
                 // Very low probability, probably not a big deal, not worth adding an extra lock.
                 _head = PromiseRefBase.InvalidAwaitSentinel.s_instance;
             }
@@ -233,9 +233,14 @@ internal void Push(HandleablePromiseBase item)
 
             [MethodImpl(InlineOption)]
             internal HandleablePromiseBase PopOrInvalid()
+                => _locker.TryEnter()
+                ? PopOrInvalidLocked()
+                : PopOrInvalidSlow();
+
+            [MethodImpl(InlineOption)]
+            private HandleablePromiseBase PopOrInvalidLocked()
             {
                 // We use InvalidAwaitSentinel as the sentinel, so we don't need a branch here to check the bottom of the stack, because it references itself.
-                _locker.EnterWithoutSleep1();
                 var head = _head;
                 _head = head._next;
                 _locker.Exit();
@@ -245,6 +250,27 @@ internal HandleablePromiseBase PopOrInvalid()
 #endif
                 return head;
             }
+
+            [MethodImpl(MethodImplOptions.NoInlining)]
+            private HandleablePromiseBase PopOrInvalidSlow()
+            {
+                // Spin until we acquire the lock.
+                var spinner = new SpinWait();
+                do
+                {
+                    // We're in the slow path, we can spend some extra time to check to see if we can exit early without taking the lock.
+                    var head = _head;
+                    if (head == PromiseRefBase.InvalidAwaitSentinel.s_instance)
+                    {
+                        // The stack is empty, exit early.
+                        return head;
+                    }
+                    spinner.SpinOnce(sleep1Threshold: -1);
+                }
+                while (!_locker.TryEnter());
+
+                return PopOrInvalidLocked();
+            }
         }
 
         /// <summary>