Add ReenterAfter override (#1915)

* Add ReenterAfter override that takes in an action with a Task with a result

* Add a test for ReenterAfter with result

Author: woksin

src/Proto.Actor/Context/ActorContext.cs

@@ -280,6 +280,23 @@ public void ReenterAfter(Task target, Action<Task> action)
         ScheduleContinuation(target, cont);
     }
 
+    public void ReenterAfter<T>(Task<T> target, Action<Task<T>> action)
+    {
+        var msg = _messageOrEnvelope;
+
+        var cont = new Continuation(
+            () =>
+            {
+                action(target);
+
+                return Task.CompletedTask;
+            },
+            msg,
+            Actor);
+
+        ScheduleContinuation(target, cont);
+    }
+
     public void ReenterAfter(Task target, Func<Task, Task> action)
     {
         var msg = _messageOrEnvelope;

src/Proto.Actor/Context/ActorContextDecorator.cs

@@ -67,6 +67,8 @@ public virtual void ReenterAfter<T>(Task<T> target, Func<Task<T>, Task> action)
 
     public void ReenterAfter(Task target, Action<Task> action) => _context.ReenterAfter(target, action);
 
+    public void ReenterAfter<T>(Task<T> target, Action<Task<T>> action) => _context.ReenterAfter(target, action);
+
     public void ReenterAfter(Task target, Func<Task, Task> action) => _context.ReenterAfter(target, action);
 
     public CapturedContext Capture() => _context.Capture();

src/Proto.Actor/Context/IContext.cs

@@ -111,6 +111,16 @@ public interface IContext : ISenderContext, IReceiverContext, ISpawnerContext, I
     /// <param name="action">The continuation to call once the task is completed. The awaited task is passed in as a parameter.</param>
     void ReenterAfter(Task target, Action<Task> action);
 
+    /// <summary>
+    ///     Awaits the given target task and once completed, the given action is then completed within the actors concurrency
+    ///     constraint.
+    ///     The concept is called Reentrancy, where an actor can continue to process messages while also awaiting that some
+    ///     asynchronous operation completes.
+    /// </summary>
+    /// <param name="target">The Task to await</param>
+    /// <param name="action">The continuation to call once the task is completed. The awaited task is passed in as a parameter.</param>
+    void ReenterAfter<T>(Task<T> target, Action<Task<T>> action);
+
     /// <summary>
     ///     Awaits the given target task and once completed, the given action is then completed within the actors concurrency
     ///     constraint.

tests/Proto.Actor.Tests/ReenterTests.cs

@@ -69,6 +69,29 @@ public async Task ReenterAfterCompletedTask()
         Assert.Equal("response", res);
     }
 
+    [Fact]
+    public async Task ReenterTaskWithResult()
+    {
+        const int expectedResult = 2;
+
+        var props = Props.FromFunc(ctx =>
+            {
+                if (ctx.Message is "reenter")
+                {
+                    var task = Task.FromResult(expectedResult);
+                    ctx.ReenterAfter(task, completedTask => { ctx.Respond(completedTask.Result); });
+                }
+
+                return Task.CompletedTask;
+            }
+        );
+
+        var pid = Context.Spawn(props);
+
+        var res = await Context.RequestAsync<int>(pid, "reenter", TimeSpan.FromSeconds(5));
+        Assert.Equal(expectedResult, res);
+    }
+
     [Fact]
     public async Task ReenterAfterTimerCancelledToken()
     {