From 3a29bf58d5d5f4e60e146763e8ade9a3b46f7882 Mon Sep 17 00:00:00 2001
From: Roger Johansson <rogeralsing@gmail.com>
Date: Wed, 15 Jun 2022 22:35:51 +0200
Subject: [PATCH] add reenterafter overloads (#1671)

---
 src/Proto.Actor/Context/ActorContext.cs       | 27 +++++++++++++++++++
 .../Context/ActorContextDecorator.cs          |  6 +++++
 src/Proto.Actor/Context/IContext.cs           | 20 ++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/src/Proto.Actor/Context/ActorContext.cs b/src/Proto.Actor/Context/ActorContext.cs
index 2744d677ea..50e0259762 100644
--- a/src/Proto.Actor/Context/ActorContext.cs
+++ b/src/Proto.Actor/Context/ActorContext.cs
@@ -231,6 +231,33 @@ public void ReenterAfter(Task target, Action action)
 
         ScheduleContinuation(target, cont);
     }
+    
+    public void ReenterAfter(Task target, Action<Task> 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;
+
+        var cont = new Continuation(
+            () => action(target),
+            msg,
+            Actor);
+
+        ScheduleContinuation(target, cont);
+    }
 
     public Task Receive(MessageEnvelope envelope)
     {
diff --git a/src/Proto.Actor/Context/ActorContextDecorator.cs b/src/Proto.Actor/Context/ActorContextDecorator.cs
index 8931fed0c4..df03005229 100644
--- a/src/Proto.Actor/Context/ActorContextDecorator.cs
+++ b/src/Proto.Actor/Context/ActorContextDecorator.cs
@@ -72,6 +72,12 @@ public virtual void ReenterAfter<T>(Task<T> target, Func<Task<T>, Task> action)
     public virtual void ReenterAfter(Task target, Action action) =>
         _context.ReenterAfter(target, action);
 
+    public void ReenterAfter(Task target, Action<Task> action) => 
+        _context.ReenterAfter(target, action);
+
+    public void ReenterAfter(Task target, Func<Task, Task> action) =>
+        _context.ReenterAfter(target, action);
+
     public CapturedContext Capture() => _context.Capture();
 
     public void Apply(CapturedContext capturedContext) => _context.Apply(capturedContext);
diff --git a/src/Proto.Actor/Context/IContext.cs b/src/Proto.Actor/Context/IContext.cs
index bae94b6825..a71bda00ab 100644
--- a/src/Proto.Actor/Context/IContext.cs
+++ b/src/Proto.Actor/Context/IContext.cs
@@ -90,6 +90,26 @@ public interface IContext : ISenderContext, IReceiverContext, ISpawnerContext, I
     /// <param name="target">the Task to await</param>
     /// <param name="action">the continuation to call once the task is completed</param>
     void ReenterAfter(Task target, Action 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</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</param>
+    void ReenterAfter(Task target, Func<Task,Task> action);
 
     /// <summary>
     /// Captures the current MessageOrEnvelope for the ActorContext