From d27d6d5d9d8ddb8b6b3469e47e4cec1c53bac622 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Mon, 25 May 2020 19:33:54 +0900
Subject: [PATCH] UniTask.Yield in .NET Core becomes zero allocation

---
 src/UniTask.NetCore/NetCore/UniTask.Yield.cs | 41 +++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/UniTask.NetCore/NetCore/UniTask.Yield.cs b/src/UniTask.NetCore/NetCore/UniTask.Yield.cs
index 1931f4e..511f706 100644
--- a/src/UniTask.NetCore/NetCore/UniTask.Yield.cs
+++ b/src/UniTask.NetCore/NetCore/UniTask.Yield.cs
@@ -1,4 +1,6 @@
-using System;
+using Cysharp.Threading.Tasks.Internal;
+using System;
+using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
 using System.Threading;
 
@@ -41,7 +43,11 @@ namespace Cysharp.Threading.Tasks
                     }
                     else
                     {
+#if NETCOREAPP3_1
+                        ThreadPool.UnsafeQueueUserWorkItem(ThreadPoolWorkItem.Create(continuation), false);
+#else
                         ThreadPool.UnsafeQueueUserWorkItem(WaitCallbackDelegate, continuation);
+#endif
                     }
                 }
 
@@ -50,6 +56,39 @@ namespace Cysharp.Threading.Tasks
                     ((Action)state).Invoke();
                 }
             }
+
+#if NETCOREAPP3_1
+
+            public sealed class ThreadPoolWorkItem : IThreadPoolWorkItem
+            {
+                static readonly ConcurrentQueue<ThreadPoolWorkItem> pool = new ConcurrentQueue<ThreadPoolWorkItem>();
+
+                Action continuation;
+
+                [MethodImpl(MethodImplOptions.AggressiveInlining)]
+                public static ThreadPoolWorkItem Create(Action continuation)
+                {
+                    if (!pool.TryDequeue(out var item))
+                    {
+                        item = new ThreadPoolWorkItem();
+                    }
+
+                    item.continuation = continuation;
+                    return item;
+                }
+
+                [MethodImpl(MethodImplOptions.AggressiveInlining)]
+                public void Execute()
+                {
+                    var call = continuation;
+                    continuation = null;
+                    pool.Enqueue(this);
+
+                    call.Invoke();
+                }
+            }
+
+#endif
         }
     }
 }
\ No newline at end of file