test: stabilize concurrent cancellation test

test/Grains/TestGrainInterfaces/IGenericInterfaces.cs

@@ -209,6 +209,8 @@ public interface ILongRunningTaskGrain<T> : IGrainWithGuidKey
         Task LongWaitWithStartNotification(TimeSpan delay, Guid callId, ILongRunningTaskObserver observer, CancellationToken cancellationToken);
         [AlwaysInterleave]
         Task LongWaitInterleaving(CancellationToken tc, TimeSpan delay, Guid callId);
+        [AlwaysInterleave]
+        Task LongWaitInterleavingWithStartNotification(TimeSpan delay, Guid callId, ILongRunningTaskObserver observer, CancellationToken cancellationToken);
         Task CallOtherLongRunningTask(ILongRunningTaskGrain<T> target, CancellationToken tc, TimeSpan delay, Guid callId);
         Task CallOtherLongRunningTaskGrainCancellation(ILongRunningTaskGrain<T> target, GrainCancellationToken tc, TimeSpan delay, Guid callId);
         Task CallOtherLongRunningTaskWithLocalGrainCancellationToken(ILongRunningTaskGrain<T> target, TimeSpan delay, TimeSpan delayBeforeCancel, Guid callId);

test/Grains/TestGrains/GenericGrains.cs

@@ -793,6 +793,12 @@ public Task LongWaitWithStartNotification(TimeSpan delay, Guid callId, ILongRunn
             return LongWait(cancellationToken, delay, callId);
         }
 
+        public Task LongWaitInterleavingWithStartNotification(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

test/Orleans.Runtime.Tests/CancellationTests/CancellationTokenTests.cs

@@ -1,3 +1,4 @@
+using System.Collections.Concurrent;
 using Orleans.Configuration;
 using Orleans.Runtime.Placement;
 using Orleans.TestingHost;
@@ -125,19 +126,48 @@ public async Task GrainTaskMultipleCancellations(int delay)
     {
         var grain = fixture.GrainFactory.GetGrain<ILongRunningTaskGrain<bool>>(Guid.NewGuid());
         var callIds = Enumerable.Range(0, 5).Select(_ => Guid.NewGuid()).ToArray();
-        var grainTasks = callIds
-            .Select(async callId =>
+        var cancellationSources = callIds.Select(_ => new CancellationTokenSource()).ToArray();
+        var observer = new LongRunningTaskObserver();
+        var observerReference = fixture.GrainFactory.CreateObjectReference<ILongRunningTaskObserver>(observer);
+        try
+        {
+            var grainTasks = callIds
+                .Select((callId, index) => delay > 0
+                    ? grain.LongWaitInterleavingWithStartNotification(
+                        TimeSpan.FromSeconds(10),
+                        callId,
+                        observerReference,
+                        cancellationSources[index].Token)
+                    : grain.LongWaitInterleaving(
+                        cancellationSources[index].Token,
+                        TimeSpan.FromSeconds(10),
+                        callId))
+                .ToArray();
+            if (delay > 0)
             {
-                using var cts = new CancellationTokenSource();
-                var task = grain.LongWaitInterleaving(cts.Token, TimeSpan.FromSeconds(10), callId);
-                cts.CancelAfter(delay);
-                await Assert.ThrowsAnyAsync<OperationCanceledException>(() => task);
-            })
-            .ToList();
-        await Task.WhenAll(grainTasks);
-        if (delay > 0)
+                await Task.WhenAll(callIds.Select(observer.WaitForCallToStart));
+            }
+
+            foreach (var cancellationSource in cancellationSources)
+            {
+                cancellationSource.CancelAfter(delay);
+            }
+
+            await Task.WhenAll(grainTasks.Select(task =>
+                Assert.ThrowsAnyAsync<OperationCanceledException>(() => task)));
+            if (delay > 0)
+            {
+                await WaitForCallCancellation(grain, callIds);
+            }
+        }
+        finally
         {
-            await WaitForCallCancellation(grain, callIds);
+            foreach (var cancellationSource in cancellationSources)
+            {
+                cancellationSource.Dispose();
+            }
+
+            fixture.GrainFactory.DeleteObjectReference<ILongRunningTaskObserver>(observerReference);
         }
     }
 
@@ -370,15 +400,14 @@ private async Task WaitForCallCancellation<T>(ILongRunningTaskGrain<T> grain, Gu
 
     private sealed class LongRunningTaskObserver : ILongRunningTaskObserver
     {
-        private readonly TaskCompletionSource<Guid> _callStarted = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly ConcurrentDictionary<Guid, TaskCompletionSource> _startedCalls = new();
 
-        public void OnCallStarted(Guid callId) => _callStarted.TrySetResult(callId);
+        public void OnCallStarted(Guid callId) => GetCallStarted(callId).TrySetResult();
 
-        public async Task WaitForCallToStart(Guid expectedCallId)
-        {
-            var callId = await _callStarted.Task.WaitAsync(TimeSpan.FromSeconds(30));
-            Assert.Equal(expectedCallId, callId);
-        }
+        public Task WaitForCallToStart(Guid callId) => GetCallStarted(callId).Task.WaitAsync(TimeSpan.FromSeconds(30));
+
+        private TaskCompletionSource GetCallStarted(Guid callId) =>
+            _startedCalls.GetOrAdd(callId, static _ => new(TaskCreationOptions.RunContinuationsAsynchronously));
     }
 
     /// <summary>