Merge pull request #83 from UiPath/pool_tcs

Pool TaskCompletionSource-s

Author: lbargaoanu

src/UiPath.CoreIpc.Tests/SystemTests.cs

@@ -175,6 +175,7 @@ public virtual async Task UploadNoRead()
             (await _systemClient.UploadNoRead(new MemoryStream(Encoding.UTF8.GetBytes("Hello world")))).ShouldBeEmpty();
         }
         catch (IOException) { }
+        catch (ObjectDisposedException) { }
         await Guid();
     }
 

src/UiPath.CoreIpc.Tests/WebSocketTests.cs

@@ -11,8 +11,8 @@ protected override ServiceHostBuilder Configure(ServiceHostBuilder serviceHostBu
     }
     public override void Dispose()
     {
-        _listener?.Dispose();
         base.Dispose();
+        _listener?.Dispose();
     }
     protected override WebSocketClientBuilder<ISystemService> CreateSystemClientBuilder() => new(new("ws"+GetEndPoint()));
     [Fact]
@@ -43,7 +43,7 @@ protected override ServiceHostBuilder Configure(ServiceHostBuilder serviceHostBu
     }
     public override void Dispose()
     {
-        _listener?.Dispose();
         base.Dispose();
+        _listener?.Dispose();
     }
 }

src/UiPath.CoreIpc/CancellationTokenSourcePool.cs

@@ -18,10 +18,6 @@ public static PooledCancellationTokenSource Rent()
     }
     private static bool Return(PooledCancellationTokenSource cts)
     {
-        if (!cts.TryReset())
-        {
-            return false;
-        }
         if (Interlocked.Increment(ref Count) > MaxQueueSize)
         {
             Interlocked.Decrement(ref Count);
@@ -32,16 +28,15 @@ private static bool Return(PooledCancellationTokenSource cts)
     }
     public sealed class PooledCancellationTokenSource : CancellationTokenSource
     {
-        protected override void Dispose(bool disposing)
+        public void Return()
         {
             // If we failed to return to the pool then dispose
-            if (disposing && !Return(this))
+#if !NET461
+            if (!TryReset() || !CancellationTokenSourcePool.Return(this))
+#endif
             {
-                base.Dispose(disposing);
+                Dispose();
             }
         }
-#if NET461
-        public bool TryReset() => false;
-#endif
     }
 }

src/UiPath.CoreIpc/Connection.cs

@@ -1,10 +1,10 @@
 namespace UiPath.CoreIpc;
-
-using RequestCompletionSource = TaskCompletionSource<Response>;
+using static TaskCompletionPool<Response>;
 using static IOHelpers;
 public sealed class Connection : IDisposable
 {
-    private readonly ConcurrentDictionary<string, RequestCompletionSource> _requests = new();
+    private static readonly IOException ClosedException = new("Connection closed.");
+    private readonly ConcurrentDictionary<string, ManualResetValueTaskSource> _requests = new();
     private long _requestCounter = -1;
     private readonly int _maxMessageSize;
     private readonly Lazy<Task> _receiveLoop;
@@ -47,19 +47,20 @@ public Connection(Stream network, ISerializer serializer, ILogger logger, string
 #endif
     internal async ValueTask<Response> RemoteCall(Request request, CancellationToken token)
     {
-        var requestCompletion = new RequestCompletionSource();
+        var requestCompletion = Rent();
         _requests[request.Id] = requestCompletion;
         CancellationTokenRegistration tokenRegistration = default;
         try
         {
             await Send(request, token);
             tokenRegistration = token.UnsafeRegister(request.UploadStream == null ? _cancelRequest : _cancelUploadRequest, request.Id);
-            return await requestCompletion.Task;
+            return await requestCompletion.ValueTask();
         }
         finally
         {
-            tokenRegistration.Dispose();
             _requests.TryRemove(request.Id, out _);
+            tokenRegistration.Dispose();
+            requestCompletion.Return();
         }
     }
     internal ValueTask Send(Request request, CancellationToken token)
@@ -85,9 +86,9 @@ void CancelUploadRequest(string requestId)
     }
     private void TryCancelRequest(string requestId)
     {
-        if (_requests.TryGetValue(requestId, out var requestCompletion))
+        if (_requests.TryRemove(requestId, out var requestCompletion))
         {
-            requestCompletion.TrySetCanceled();
+            requestCompletion.SetCanceled();
         }
     }
     internal ValueTask Send(Response response, CancellationToken cancellationToken)
@@ -165,9 +166,12 @@ public void Dispose()
     }
     private void CompleteRequests()
     {
-        foreach (var completionSource in _requests.Values)
+        foreach (var requestId in _requests.Keys)
         {
-            completionSource.TrySetException(new IOException("Connection closed."));
+            if (_requests.TryRemove(requestId, out var requestCompletion))
+            {
+                requestCompletion.SetException(ClosedException);
+            }
         }
     }
 #if !NET461
@@ -287,7 +291,7 @@ private async Task EnterStreamMode()
     {
         if (!await ReadBuffer(sizeof(long)))
         {
-            throw new IOException("Connection closed.");
+            throw ClosedException;
         }
         var userStreamLength = BitConverter.ToInt64(_buffer, startIndex: 0);
         _nestedStream.Reset(userStreamLength);
@@ -340,9 +344,9 @@ private void OnResponseReceived(Response response)
             {
                 Log($"Received response for request {response.RequestId} {Name}.");
             }
-            if (_requests.TryGetValue(response.RequestId, out var completionSource))
+            if (_requests.TryRemove(response.RequestId, out var completionSource))
             {
-                completionSource.TrySetResult(response);
+                completionSource.SetResult(response);
             }
         }
         catch (Exception ex)

src/UiPath.CoreIpc/Helpers.cs

@@ -291,7 +291,7 @@ public void ThrowTimeout(Exception exception, string message)
     public void Dispose()
     {
         _linkedRegistration.Dispose();
-        _timeoutCancellationSource.Dispose();
+        _timeoutCancellationSource.Return();
     }
     public CancellationToken Token => _timeoutCancellationSource.Token;
 }

src/UiPath.CoreIpc/Server/Server.cs

@@ -1,18 +1,17 @@
 using System.Linq.Expressions;
-
 namespace UiPath.CoreIpc;
-
 using GetTaskResultFunc = Func<Task, object>;
 using MethodExecutor = Func<object, object[], CancellationToken, Task>;
 using static Expression;
+using static CancellationTokenSourcePool;
 class Server
 {
     private static readonly MethodInfo GetResultMethod = typeof(Server).GetStaticMethod(nameof(GetTaskResultImpl));
     private static readonly ConcurrentDictionary<(Type,string), Method> Methods = new();
     private static readonly ConcurrentDictionary<Type, GetTaskResultFunc> GetTaskResultByType = new();
     private readonly Connection _connection;
     private readonly IClient _client;
-    private readonly ConcurrentDictionary<string, CancellationTokenSource> _requests = new();
+    private readonly ConcurrentDictionary<string, PooledCancellationTokenSource> _requests = new();
     public Server(ListenerSettings settings, Connection connection, IClient client = null)
     {
         Settings = settings;
@@ -44,7 +43,7 @@ void CancelRequest(string requestId)
         if (_requests.TryRemove(requestId, out var cancellation))
         {
             cancellation.Cancel();
-            cancellation.Dispose();
+            cancellation.Return();
         }
     }
 #if !NET461
@@ -70,7 +69,7 @@ async ValueTask OnRequestReceived(Request request)
                 return;
             }
             Response response = null;
-            var requestCancellation = CancellationTokenSourcePool.Rent();
+            var requestCancellation = Rent();
             _requests[request.Id] = requestCancellation;
             var timeout = request.GetTimeout(Settings.RequestTimeout);
             var timeoutHelper = new TimeoutHelper(timeout, requestCancellation.Token);
@@ -99,7 +98,7 @@ async ValueTask OnRequestReceived(Request request)
         }
         if (_requests.TryRemove(request.Id, out var cancellation))
         {
-            cancellation.Dispose();
+            cancellation.Return();
         }
     }
     ValueTask OnError(Request request, Exception ex)

src/UiPath.CoreIpc/TaskCompletionPool.cs

@@ -0,0 +1,43 @@
+using System.Threading.Tasks.Sources;
+namespace UiPath.CoreIpc;
+internal static class TaskCompletionPool<T>
+{
+    private const int MaxQueueSize = 1024;
+    private static readonly ConcurrentQueue<ManualResetValueTaskSource> Cache = new();
+    private static int Count;
+    public static ManualResetValueTaskSource Rent()
+    {
+        if (Cache.TryDequeue(out var source))
+        {
+            Interlocked.Decrement(ref Count);
+            return source;
+        }
+        return new();
+    }
+    private static void Return(ManualResetValueTaskSource source)
+    {
+        if (Interlocked.Increment(ref Count) > MaxQueueSize)
+        {
+            Interlocked.Decrement(ref Count);
+            return;
+        }
+        source.Reset();
+        Cache.Enqueue(source);
+    }
+    public sealed class ManualResetValueTaskSource : IValueTaskSource<T>, IValueTaskSource
+    {
+        private ManualResetValueTaskSourceCore<T> _core; // mutable struct; do not make this readonly
+        public bool RunContinuationsAsynchronously { get => _core.RunContinuationsAsynchronously; set => _core.RunContinuationsAsynchronously = value; }
+        public short Version => _core.Version;
+        public ValueTask<T> ValueTask() => new(this, Version);
+        public void Reset() => _core.Reset();
+        public void SetResult(T result) => _core.SetResult(result);
+        public void SetException(Exception error) => _core.SetException(error);
+        public void SetCanceled() => _core.SetException(new TaskCanceledException());
+        public T GetResult(short token) => _core.GetResult(token);
+        void IValueTaskSource.GetResult(short token) => _core.GetResult(token);
+        public ValueTaskSourceStatus GetStatus(short token) => _core.GetStatus(token);
+        public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _core.OnCompleted(continuation, state, token, flags);
+        public void Return() => TaskCompletionPool<T>.Return(this);
+    }
+}