From 1119298391c2dc4a1bf0561999ee2d7f03386866 Mon Sep 17 00:00:00 2001
From: Reuben Bond <rebond@microsoft.com>
Date: Fri, 31 Jul 2026 23:31:08 -0700
Subject: [PATCH] test: wait for grain call before checking cancellation

Ensure delayed cancellation cases observe that the grain method has started before expecting its catch block to record cancellation.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 44628696-d4d6-43f1-802d-dcf7a0d8952a
---
 .../TestGrainInterfaces/IGenericInterfaces.cs |  6 +++
 test/Grains/TestGrains/GenericGrains.cs       |  7 ++++
 .../CancellationTokenTests.cs                 | 42 +++++++++++++++----
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/test/Grains/TestGrainInterfaces/IGenericInterfaces.cs b/test/Grains/TestGrainInterfaces/IGenericInterfaces.cs
index 9506b448b3c..b7d033ad5f8 100644
--- a/test/Grains/TestGrainInterfaces/IGenericInterfaces.cs
+++ b/test/Grains/TestGrainInterfaces/IGenericInterfaces.cs
@@ -189,6 +189,11 @@ public interface IGenericPingSelf<T> : IGrainWithGuidKey
         Task ScheduleDelayedPingToSelfAndDeactivate(IGenericPingSelf<T> target, T t, TimeSpan delay);
     }
 
+    public interface ILongRunningTaskObserver : IGrainObserver
+    {
+        void OnCallStarted(Guid callId);
+    }
+
     public interface ILongRunningTaskGrain<T> : IGrainWithGuidKey
     {
         Task<string> GetRuntimeInstanceId();
@@ -201,6 +206,7 @@ public interface ILongRunningTaskGrain<T> : IGrainWithGuidKey
         [AlwaysInterleave]
         Task LongWaitGrainCancellationInterleaving(GrainCancellationToken tc, TimeSpan delay, Guid callId);
         Task LongWait(CancellationToken tc, TimeSpan delay, Guid callId);
+        Task LongWaitWithStartNotification(TimeSpan delay, Guid callId, ILongRunningTaskObserver observer, CancellationToken cancellationToken);
         [AlwaysInterleave]
         Task LongWaitInterleaving(CancellationToken tc, TimeSpan delay, Guid callId);
         Task CallOtherLongRunningTask(ILongRunningTaskGrain<T> target, CancellationToken tc, TimeSpan delay, Guid callId);
diff --git a/test/Grains/TestGrains/GenericGrains.cs b/test/Grains/TestGrains/GenericGrains.cs
index 5cd33341b9b..4ea36868824 100644
--- a/test/Grains/TestGrains/GenericGrains.cs
+++ b/test/Grains/TestGrains/GenericGrains.cs
@@ -785,6 +785,13 @@ public async Task LongWaitGrainCancellation(GrainCancellationToken ct, TimeSpan
         }
 
         public Task LongWaitInterleaving(CancellationToken ct, TimeSpan delay, Guid callId) => LongWait(ct, delay, callId);
+
+        public Task LongWaitWithStartNotification(TimeSpan delay, Guid callId, ILongRunningTaskObserver observer, CancellationToken cancellationToken)
+        {
+            observer.OnCallStarted(callId);
+            return LongWait(cancellationToken, delay, callId);
+        }
+
         public async Task LongWait(CancellationToken ct, TimeSpan delay, Guid callId)
         {
             try
diff --git a/test/Orleans.Runtime.Tests/CancellationTests/CancellationTokenTests.cs b/test/Orleans.Runtime.Tests/CancellationTests/CancellationTokenTests.cs
index ce1e74827a7..e5076828d87 100644
--- a/test/Orleans.Runtime.Tests/CancellationTests/CancellationTokenTests.cs
+++ b/test/Orleans.Runtime.Tests/CancellationTests/CancellationTokenTests.cs
@@ -67,14 +67,29 @@ protected override void ConfigureTestCluster(InProcessTestClusterBuilder builder
     public async Task GrainTaskCancellation(int delay)
     {
         var grain = fixture.GrainFactory.GetGrain<ILongRunningTaskGrain<bool>>(Guid.NewGuid());
-        using var cts = new CancellationTokenSource();
-        var callId = Guid.NewGuid();
-        var grainTask = grain.LongWait(cts.Token, TimeSpan.FromSeconds(10), callId);
-        cts.CancelAfter(delay);
-        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => grainTask);
-        if (delay > 0)
+        var observer = new LongRunningTaskObserver();
+        var observerReference = fixture.GrainFactory.CreateObjectReference<ILongRunningTaskObserver>(observer);
+        try
         {
-            await WaitForCallCancellation(grain, callId);
+            using var cts = new CancellationTokenSource();
+            var callId = Guid.NewGuid();
+            var grainTask = grain.LongWaitWithStartNotification(TimeSpan.FromSeconds(10), callId, observerReference, cts.Token);
+            if (delay > 0)
+            {
+                // A timer does not guarantee that a new activation has begun executing the request.
+                await observer.WaitForCallToStart(callId);
+            }
+
+            cts.CancelAfter(delay);
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => grainTask);
+            if (delay > 0)
+            {
+                await WaitForCallCancellation(grain, callId);
+            }
+        }
+        finally
+        {
+            fixture.GrainFactory.DeleteObjectReference<ILongRunningTaskObserver>(observerReference);
         }
     }
 
@@ -353,6 +368,19 @@ private async Task WaitForCallCancellation<T>(ILongRunningTaskGrain<T> grain, Gu
         Assert.Fail("Did not encounter the expected call id");
     }
 
+    private sealed class LongRunningTaskObserver : ILongRunningTaskObserver
+    {
+        private readonly TaskCompletionSource<Guid> _callStarted = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public void OnCallStarted(Guid callId) => _callStarted.TrySetResult(callId);
+
+        public async Task WaitForCallToStart(Guid expectedCallId)
+        {
+            var callId = await _callStarted.Task.WaitAsync(TimeSpan.FromSeconds(30));
+            Assert.Equal(expectedCallId, callId);
+        }
+    }
+
     /// <summary>
     /// Tests that a running interleaving grain operation can be cancelled via CancellationToken.
     /// Interleaving requests run concurrently without queueing and should also be cancellable.