From bbe5322474868edae87bace0212c3d5b06445958 Mon Sep 17 00:00:00 2001
From: Vladyslav Taranov <wvlad.dev@gmail.com>
Date: Tue, 2 Apr 2024 11:15:13 +0300
Subject: [PATCH] feat: 2110 memory_leak_bugfix

---
 src/Proto.Remote/Endpoints/Endpoint.cs | 44 +++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/src/Proto.Remote/Endpoints/Endpoint.cs b/src/Proto.Remote/Endpoints/Endpoint.cs
index cbe25f00ff..c932a07a5c 100644
--- a/src/Proto.Remote/Endpoints/Endpoint.cs
+++ b/src/Proto.Remote/Endpoints/Endpoint.cs
@@ -375,17 +375,19 @@ private void RejectRemoteDeliver(RemoteDeliver env)
 
     private async Task RunAsync()
     {
+        var waiter = new MultiTaskReuseWaiter<bool>(
+            () => _remoteDelivers.Reader.WaitToReadAsync(CancellationToken),
+            () => _remotePriorityDelivers.Reader.WaitToReadAsync(CancellationToken));
+        
         while (!CancellationToken.IsCancellationRequested)
         {
             try
             {
                 var messages = new List<RemoteDeliver>(RemoteConfig.EndpointWriterOptions.EndpointWriterBatchSize);
-                
                 while (true)
                 {
-                    var t1 = _remoteDelivers.Reader.WaitToReadAsync(CancellationToken).AsTask();
-                    var t2 = _remotePriorityDelivers.Reader.WaitToReadAsync(CancellationToken).AsTask();
-                    await Task.WhenAny(t1, t2);
+                    await waiter.WaitAnyAsync();
+                    
                     var i = 0;
                     while (true)
                     {
@@ -571,4 +573,38 @@ private MessageBatch CreateBatch(IReadOnlyCollection<RemoteDeliver> m)
         //     Logger.LogTrace("[{SystemAddress}] Sending {Count} envelopes for {Address}", System.Address, envelopes.Count, Address);
         return batch;
     }
+    
+    /// <summary>
+    /// Preserves non completed Tasks between <see cref="WaitAnyAsync"/> calls.
+    /// This is necessary to prevent memory leak https://github.com/asynkron/protoactor-dotnet/issues/2110 
+    /// </summary>
+    class MultiTaskReuseWaiter<T>
+    {
+        private readonly Func<ValueTask<T>>[] _taskFactories;
+        private readonly Task<T>?[] _tasks;
+
+        public MultiTaskReuseWaiter(params Func<ValueTask<T>>[] taskFactories)
+        {
+            _taskFactories = taskFactories;
+            _tasks = new Task<T>?[taskFactories.Length];
+        }
+
+        public async ValueTask<T> WaitAnyAsync()
+        {
+            for (var i = 0; i < _taskFactories.Length; i++)
+            {
+                if (_tasks[i]?.IsCompleted == false)
+                    continue;
+                
+                var vt = _taskFactories[i].Invoke();
+                if (vt.IsCompleted)
+                    return await vt;
+                
+                _tasks[i] = vt.AsTask();
+            }
+
+            return await await Task.WhenAny(_tasks!);
+        }
+
+    }
 }
\ No newline at end of file