From 54eef664c51b3c1397300219412b39a09f4bcf83 Mon Sep 17 00:00:00 2001
From: Ben Wilde <benbenwilde@gmail.com>
Date: Thu, 23 Jan 2025 14:58:25 -0600
Subject: [PATCH] After an actor is stopped, block continuations from running
 (#2146)

* After an actor is stopped, block continuations from running

* Additionally check the actors cancellation token before sending the Continuation
---
 src/Proto.Actor/Context/ActorContext.cs |  4 +--
 tests/Proto.Actor.Tests/ReenterTests.cs | 43 +++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/Proto.Actor/Context/ActorContext.cs b/src/Proto.Actor/Context/ActorContext.cs
index 0ac638e382..f99edcb2da 100644
--- a/src/Proto.Actor/Context/ActorContext.cs
+++ b/src/Proto.Actor/Context/ActorContext.cs
@@ -596,7 +596,7 @@ static async Task Continue(Task target, Continuation cont, IContext ctx)
                 // ignored
             }
 
-            if (!ctx.System.Shutdown.IsCancellationRequested)
+            if (!ctx.CancellationToken.IsCancellationRequested && !ctx.System.Shutdown.IsCancellationRequested)
             {
                 ctx.Self.SendSystemMessage(ctx.System, cont);
             }
@@ -621,7 +621,7 @@ private async Task HandleContinuation(Continuation cont)
         // Don't execute the continuation if the actor instance changed.
         // Without this, Continuation's Action closure would execute with
         // an older Actor instance.
-        if (cont.Actor != Actor && cont is not { Actor: null })
+        if (_state == ContextState.Stopped || System.Shutdown.IsCancellationRequested || (cont.Actor != Actor && cont is not { Actor: null }))
         {
                 Logger.DroppingContinuation(Self, MessageEnvelope.UnwrapMessage(cont.Message));
 
diff --git a/tests/Proto.Actor.Tests/ReenterTests.cs b/tests/Proto.Actor.Tests/ReenterTests.cs
index 4352363413..70d61cdd9b 100644
--- a/tests/Proto.Actor.Tests/ReenterTests.cs
+++ b/tests/Proto.Actor.Tests/ReenterTests.cs
@@ -301,6 +301,49 @@ public async Task DropReenterContinuationAfterRestart()
         var res = await Context.RequestAsync<bool>(pid, "waitstate", TimeSpan.FromSeconds(5));
         Assert.True(res);
     }
+    
+    [Fact]
+    public async Task DropReenterContinuationAfterStop()
+    {
+        var stopped = false;
+        var completionExecuted = false;
+        CancellationTokenSource cts = new();
+
+        var props = Props.FromFunc(async ctx =>
+            {
+                switch (ctx.Message)
+                {
+                    case "start":
+
+                        ctx.ReenterAfter(
+                            Task.Delay(-1, cts.Token),
+                            () => { completionExecuted = true; });
+                        
+                        ctx.Stop(ctx.Self);
+                        
+                        // Release the cancellation token after stop gets processed.
+                        cts.Cancel();
+
+                        ctx.Respond(true);
+
+                        break;
+                    case Stopped:
+                        stopped = true;
+                        
+                        break;
+                }
+            }
+        );
+
+        var pid = Context.Spawn(props);
+
+        await Context.RequestAsync<bool>(pid, "start", TimeSpan.FromSeconds(5));
+        
+        await Task.Delay(500);
+        await Task.Yield();
+        
+        Assert.True(!completionExecuted);
+    }
 
     private class ReenterAfterCancellationActor : IActor
     {