fix(transactions): await recovery cleanup tasks

.editorconfig

@@ -749,6 +749,10 @@ dotnet_diagnostic.CA1815.severity = warning
 [src/{Azure/Orleans.Reminders.AzureStorage,Orleans.Reminders.TestKit}/**.cs]
 dotnet_diagnostic.CA2219.severity = warning
 
+# CA2025: Do not pass IDisposable instances into unawaited tasks
+[src/Orleans.Transactions.TestKit.Base/**.cs]
+dotnet_diagnostic.CA2025.severity = warning
+
 # Production source enforces the correctness baseline. Tests, samples, documentation,
 # playgrounds, templates, and tools retain advisory diagnostics during this rollout.
 [src/**/*.{cs,vb}]

src/Orleans.Transactions.TestKit.Base/TestRunners/TransactionRecoveryTestsRunner.cs

@@ -71,26 +71,27 @@ private sealed record TransactionFailure(int Index, Guid[] GrainIds, Exception E
 
         private sealed record InFlightBatch(int Index, int PendingCount);
 
-        private sealed class ProducerCancellationScope(
+        internal sealed class ProducerCancellationScope(
             CancellationTokenSource stopProducing,
-            Task producer) : IDisposable
+            Task producer) : IAsyncDisposable
         {
-            public void Dispose()
+            public async ValueTask DisposeAsync()
             {
-                stopProducing.Cancel();
-                producer.Ignore();
-                if (producer.IsCompleted)
+                try
                 {
-                    stopProducing.Dispose();
-                    return;
+                    await stopProducing.CancelAsync();
+                }
+                finally
+                {
+                    try
+                    {
+                        await producer.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+                    }
+                    finally
+                    {
+                        stopProducing.Dispose();
+                    }
                 }
-
-                _ = producer.ContinueWith(
-                    static (_, state) => ((CancellationTokenSource)state!).Dispose(),
-                    stopProducing,
-                    CancellationToken.None,
-                    TaskContinuationOptions.ExecuteSynchronously,
-                    TaskScheduler.Default);
             }
         }
 
@@ -281,6 +282,7 @@ private async Task TransactionWillRecoverAfterTargetedPhase(
 
             TransactionRecoveryEventObserver.PhaseGate? cleanupGate = null;
             Task<TransactionRecoveryEventObserver.RecoveryTransition>? cleanupObservation = null;
+            using var cleanupCancellation = new CancellationTokenSource();
             if (requireParticipantConfirmation)
             {
                 cleanupGate = recoveryEvents.GateNextTransition(candidate =>
@@ -289,9 +291,15 @@ private async Task TransactionWillRecoverAfterTargetedPhase(
                     && candidate.TransactionId is { } transactionId
                     && transition.TransactionIds.Contains(transactionId)
                     && candidate.GrainId != transition.GrainId);
-                cleanupObservation = ObserveAndReleaseGateAsync(cleanupGate, GetDeadline(this.recoveryTimeout));
+#pragma warning disable CA2025 // The observation is cancelled and awaited in finally before cleanupGate is disposed.
+                cleanupObservation = ObserveAndReleaseGateAsync(
+                    cleanupGate,
+                    GetDeadline(this.recoveryTimeout),
+                    cleanupCancellation.Token);
+#pragma warning restore CA2025
             }
 
+            var operationCompleted = false;
             try
             {
                 var siloToTerminate = this.testCluster.Silos.Single(
@@ -302,7 +310,9 @@ private async Task TransactionWillRecoverAfterTargetedPhase(
                 var stopping = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
                 using var stoppingRegistration = applicationLifetime.ApplicationStopping.Register(
                     () => stopping.TrySetResult());
+#pragma warning disable CA2025 // The original shutdown task is awaited normally and on timeout before the fixture can dispose testCluster or siloToTerminate.
                 var shutdown = this.testCluster.KillSiloAsync(siloToTerminate);
+#pragma warning restore CA2025
                 try
                 {
                     await stopping.Task.WaitAsync(this.failureDetectionTimeout);
@@ -312,7 +322,7 @@ private async Task TransactionWillRecoverAfterTargetedPhase(
                 catch (TimeoutException)
                 {
                     phaseGate.Release();
-                    shutdown.Ignore();
+                    await shutdown;
                     throw;
                 }
 
@@ -356,22 +366,40 @@ private async Task TransactionWillRecoverAfterTargetedPhase(
                 }
 
                 await ValidateResults(txGrains, transactionGroups);
+                operationCompleted = true;
             }
             finally
             {
                 phaseGate.Release();
-                cleanupObservation?.Ignore();
+                await cleanupCancellation.CancelAsync();
+                if (cleanupObservation is not null)
+                {
+                    try
+                    {
+                        await cleanupObservation;
+                    }
+                    catch (OperationCanceledException) when (cleanupCancellation.IsCancellationRequested)
+                    {
+                    }
+                    catch (TimeoutException) when (!operationCompleted)
+                    {
+                        this.Log(
+                            "Recovery phase=participant-cleanup-watchdog also expired while the recovery operation failed.");
+                    }
+                }
+
                 cleanupGate?.Dispose();
             }
         }
 
-        private static async Task<TransactionRecoveryEventObserver.RecoveryTransition> ObserveAndReleaseGateAsync(
+        internal static async Task<TransactionRecoveryEventObserver.RecoveryTransition> ObserveAndReleaseGateAsync(
             TransactionRecoveryEventObserver.PhaseGate gate,
-            long deadline)
+            long deadline,
+            CancellationToken cancellationToken)
         {
             try
             {
-                return await gate.WaitAsync(deadline);
+                return await gate.WaitAsync(deadline, cancellationToken);
             }
             finally
             {
@@ -438,16 +466,15 @@ protected virtual async Task TransactionWillRecoverAfterRandomSiloFailure(
             Task<List<ExpectedGrainActivity>[]?> producer = RunWhileSucceeding(
                 transactionGroups,
                 getIndex,
-                stopProducing,
+                stopProducing.Token,
                 firstFailure,
                 firstInFlightBatch);
-            using var producerCancellationScope = new ProducerCancellationScope(stopProducing, producer);
+            await using var producerCancellationScope = new ProducerCancellationScope(stopProducing, producer);
             var inFlightBatch = await firstInFlightBatch.Task.WaitAsync(this.failureDetectionTimeout, cancellationToken);
 
             if (firstFailure.Task.IsCompleted)
             {
-                stopProducing.Cancel();
-                producer.Ignore();
+                await stopProducing.CancelAsync();
                 var prematureFailure = await firstFailure.Task.WaitAsync(cancellationToken);
                 throw new InvalidOperationException(
                     $"A transaction failed before the silo was terminated. Index: {prematureFailure.Index}. "
@@ -629,7 +656,7 @@ private static Task WakeupGrains(List<ITransactionalBitArrayGrain> grains)
         private async Task<List<ExpectedGrainActivity>[]?> RunWhileSucceeding(
             List<ExpectedGrainActivity>[] transactionGroups,
             Func<int> getIndex,
-            CancellationTokenSource stopProducing,
+            CancellationToken stopProducing,
             TaskCompletionSource<TransactionFailure> firstFailure,
             TaskCompletionSource<InFlightBatch> firstInFlightBatch)
         {
@@ -641,10 +668,7 @@ private static Task WakeupGrains(List<ITransactionalBitArrayGrain> grains)
                     transactionIndex,
                     failure =>
                     {
-                        if (firstFailure.TrySetResult(failure))
-                        {
-                            stopProducing.Cancel();
-                        }
+                        firstFailure.TrySetResult(failure);
                     },
                     tasks =>
                     {

test/Transactions/Orleans.Transactions.Tests/TransactionRecoveryTestsRunnerLifetimeTests.cs

@@ -0,0 +1,150 @@
+using System.Diagnostics;
+using Orleans.Transactions.Diagnostics;
+using Orleans.Transactions.TestKit;
+using TestExtensions;
+using Xunit;
+
+namespace Orleans.Transactions.Tests;
+
+[TestSuite("BVT")]
+[TestProvider("None")]
+[TestArea("Transactions")]
+[TestCategory("BVT"), TestCategory("Transactions")]
+public class TransactionRecoveryTestsRunnerLifetimeTests
+{
+    private static readonly TimeSpan SynchronizationTimeout = TimeSpan.FromSeconds(10);
+
+    [Fact]
+    public async Task ProducerCancellationScope_CompletedProducerDisposesCancellationSource()
+    {
+        var stopProducing = new CancellationTokenSource();
+        var scope = new TransactionRecoveryTestsRunner.ProducerCancellationScope(
+            stopProducing,
+            Task.CompletedTask);
+
+        await scope.DisposeAsync();
+
+        Assert.Throws<ObjectDisposedException>(() => stopProducing.Token);
+    }
+
+    [Fact]
+    public async Task ProducerCancellationScope_FaultedProducerDoesNotThrowDuringDisposal()
+    {
+        var stopProducing = new CancellationTokenSource();
+        var scope = new TransactionRecoveryTestsRunner.ProducerCancellationScope(
+            stopProducing,
+            Task.FromException(new InvalidOperationException("boom")));
+
+        await scope.DisposeAsync();
+
+        Assert.Throws<ObjectDisposedException>(() => stopProducing.Token);
+    }
+
+    [Fact]
+    public async Task ProducerCancellationScope_ActiveProducerIsAwaitedBeforeCancellationSourceDisposal()
+    {
+        var stopProducing = new CancellationTokenSource();
+        var cancellationObserved = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        using var registration = stopProducing.Token.Register(() => cancellationObserved.TrySetResult());
+        var producer = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var scope = new TransactionRecoveryTestsRunner.ProducerCancellationScope(
+            stopProducing,
+            producer.Task);
+
+        var disposal = scope.DisposeAsync().AsTask();
+        await cancellationObserved.Task.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+
+        Assert.False(disposal.IsCompleted);
+        Assert.True(stopProducing.Token.IsCancellationRequested);
+
+        producer.TrySetResult();
+        await disposal.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+
+        Assert.Throws<ObjectDisposedException>(() => stopProducing.Token);
+    }
+
+    [Fact]
+    public async Task ObserveAndReleaseGateAsync_ReachedGateCompletesBeforeDisposal()
+    {
+        var observer = new TransactionRecoveryEventObserver(_ => true);
+        var gate = observer.GateNextTransition(_ => true);
+        var transition = CreateTransition();
+        var observation = TransactionRecoveryTestsRunner.ObserveAndReleaseGateAsync(
+            gate,
+            GetDeadline(),
+            TestContext.Current.CancellationToken);
+        var blockingTask = RunLongRunning(() =>
+        {
+            Assert.True(gate.TryReach(transition));
+            gate.Block();
+        });
+
+        try
+        {
+            Assert.Same(transition, await observation.WaitAsync(
+                SynchronizationTimeout,
+                TestContext.Current.CancellationToken));
+            await blockingTask.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+        }
+        finally
+        {
+            gate.Dispose();
+            observer.Dispose();
+            await blockingTask.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+        }
+    }
+
+    [Fact]
+    public async Task ObserveAndReleaseGateAsync_CancellationCompletesBeforeDisposal()
+    {
+        var observer = new TransactionRecoveryEventObserver(_ => true);
+        var gate = observer.GateNextTransition(_ => true);
+        using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(
+            TestContext.Current.CancellationToken);
+        var observation = TransactionRecoveryTestsRunner.ObserveAndReleaseGateAsync(
+            gate,
+            GetDeadline(),
+            cancellation.Token);
+        var blockingTask = RunLongRunning(gate.Block);
+
+        try
+        {
+            await cancellation.CancelAsync();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => observation);
+            await blockingTask.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+        }
+        finally
+        {
+            gate.Dispose();
+            observer.Dispose();
+            await blockingTask.WaitAsync(SynchronizationTimeout, TestContext.Current.CancellationToken);
+        }
+    }
+
+    private static TransactionRecoveryEventObserver.RecoveryTransition CreateTransition() =>
+        new(
+            Sequence: 1,
+            ObservedAtUtc: DateTime.UtcNow,
+            Elapsed: TimeSpan.Zero,
+            TransactionRecoveryEventObserver.RecoveryTransitionKind.TransactionConfirmCompleted,
+            TransactionIds: [],
+            TransactionDiagnosticEvents.TransactionProtocolRole.RemoteParticipant,
+            TransactionDiagnosticEvents.TransactionPhase.Confirm,
+            ResourceName: "test",
+            GrainId: null,
+            SiloAddress: null,
+            ActivationId: default,
+            Status: null,
+            CommitCount: null,
+            Succeeded: true);
+
+    private static long GetDeadline()
+        => Stopwatch.GetTimestamp() + (long)(SynchronizationTimeout.TotalSeconds * Stopwatch.Frequency);
+
+    private static Task RunLongRunning(Action action)
+        => Task.Factory.StartNew(
+            action,
+            CancellationToken.None,
+            TaskCreationOptions.LongRunning,
+            TaskScheduler.Default);
+}