fix(testing): await reminder ring listener updates

Author: ReubenBond

src/Orleans.Reminders/ReminderService/LocalReminderService.cs

@@ -10,6 +10,7 @@
 using Orleans.Reminders.Diagnostics;
 using Orleans.Runtime.ConsistentRing;
 using Orleans.Runtime.Internal;
+using Orleans.Runtime.MembershipService;
 using Orleans.Runtime.Scheduler;
 
 namespace Orleans.Runtime.ReminderService
@@ -28,6 +29,7 @@ internal sealed partial class LocalReminderService : GrainService, IReminderServ
         private readonly IAsyncTimer listRefreshTimer; // timer that refreshes our list of reminders to reflect global reminder table
         private readonly GrainReferenceActivator _referenceActivator;
         private readonly GrainInterfaceType _grainInterfaceType;
+        private readonly SiloStatusListenerManager _siloStatusListenerManager;
         private readonly TimeProvider _timeProvider;
         private readonly ReminderInstruments _reminderInstruments;
         private long localTableSequence;
@@ -50,6 +52,7 @@ public LocalReminderService(
             IAsyncTimerFactory asyncTimerFactory,
             IOptions<ReminderOptions> reminderOptions,
             IConsistentRingProvider ringProvider,
+            SiloStatusListenerManager siloStatusListenerManager,
             [FromKeyedServices(ReminderTimeProviderNames.Reminders)] TimeProvider timeProvider,
             ReminderInstruments reminderInstruments,
             SystemTargetShared shared)
@@ -60,6 +63,7 @@ public LocalReminderService(
         {
             _referenceActivator = referenceActivator;
             _grainInterfaceType = interfaceTypeResolver.GetGrainInterfaceType(typeof(IRemindable));
+            _siloStatusListenerManager = siloStatusListenerManager;
             this.reminderOptions = reminderOptions.Value;
             this.reminderTable = reminderTable;
             _timeProvider = timeProvider;
@@ -370,6 +374,9 @@ internal Task TestOnlyRefresh()
             return refreshTask.Unwrap();
         }
 
+        internal Task TestOnlyWaitForSiloStatusListeners(CancellationToken cancellationToken)
+            => _siloStatusListenerManager.TestOnlyWaitForCurrentMembershipVersion(cancellationToken);
+
         private void RemoveOutOfRangeReminders(List<Task> removedReminderTasks)
         {
             CheckRuntimeContext();

src/Orleans.Runtime/MembershipService/SiloStatusListenerManager.cs

@@ -22,7 +22,10 @@ internal sealed partial class SiloStatusListenerManager : ILifecycleParticipant<
     private readonly IMembershipManager _membershipService;
     private readonly ILogger<SiloStatusListenerManager> _logger;
     private readonly IFatalErrorHandler _fatalErrorHandler;
+    private readonly object _processedMembershipLock = new();
     private ImmutableList<WeakReference<ISiloStatusListener>> _listeners = [];
+    private MembershipVersion _processedMembershipVersion = MembershipVersion.MinValue;
+    private TaskCompletionSource _processedMembershipVersionChanged = CreateCompletion();
 
     public SiloStatusListenerManager(
         IMembershipManager membershipManager,
@@ -88,6 +91,7 @@ private async Task ProcessMembershipUpdates()
                 var update = (previous is null || snapshot.Version == MembershipVersion.MinValue) ? snapshot.AsUpdate() : snapshot.CreateUpdate(previous);
                 NotifyObservers(update);
                 previous = snapshot;
+                MarkMembershipVersionProcessed(snapshot.Version);
             }
         }
         catch (OperationCanceledException) when (_cancellation.IsCancellationRequested)
@@ -144,6 +148,47 @@ private void NotifyObservers(ClusterMembershipUpdate update)
         }
     }
 
+    internal async Task TestOnlyWaitForCurrentMembershipVersion(CancellationToken cancellationToken)
+    {
+        var targetVersion = _membershipService.CurrentSnapshot.Version;
+        while (true)
+        {
+            Task versionChanged;
+            lock (_processedMembershipLock)
+            {
+                if (_processedMembershipVersion >= targetVersion)
+                {
+                    return;
+                }
+
+                versionChanged = _processedMembershipVersionChanged.Task;
+            }
+
+            await versionChanged.WaitAsync(cancellationToken);
+        }
+    }
+
+    private void MarkMembershipVersionProcessed(MembershipVersion version)
+    {
+        TaskCompletionSource versionChanged;
+        lock (_processedMembershipLock)
+        {
+            if (version <= _processedMembershipVersion)
+            {
+                return;
+            }
+
+            _processedMembershipVersion = version;
+            versionChanged = _processedMembershipVersionChanged;
+            _processedMembershipVersionChanged = CreateCompletion();
+        }
+
+        versionChanged.TrySetResult();
+    }
+
+    private static TaskCompletionSource CreateCompletion()
+        => new(TaskCreationOptions.RunContinuationsAsynchronously);
+
     void ILifecycleParticipant<ISiloLifecycle>.Participate(ISiloLifecycle lifecycle)
     {
         Task? task = null;

test/Orleans.Reminders.Tests/TimerTests/LocalReminderServiceTests.cs

@@ -254,6 +254,51 @@ public async Task InitialRead_TreatsNullTableResultAsNoWork()
         Assert.True(reminderTable.RangeReadCount > 0);
     }
 
+    [TestSuite("BVT")]
+    [TestProvider("None")]
+    [Fact, TestCategory("BVT")]
+    public async Task SiloStatusListenerBarrier_WaitsForCurrentMembershipVersion()
+    {
+        var initialSilo = Assert.Single(fixture.HostedCluster.Silos);
+        using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        cancellation.CancelAfter(TestConstants.InitTimeout);
+        var listener = new BlockingSiloStatusListener(initialSilo.SiloAddress);
+        var statusOracle = initialSilo.ServiceProvider.GetRequiredService<ISiloStatusOracle>();
+        Assert.True(statusOracle.SubscribeToSiloStatusEvents(listener));
+        Task<List<InProcessSiloHandle>>? startTask = null;
+        InProcessSiloHandle? joinedSilo = null;
+        try
+        {
+            startTask = fixture.HostedCluster.StartSilosAsync(1, cancellation.Token);
+            await listener.WaitUntilBlockedAsync(cancellation.Token);
+
+            var reminderService = initialSilo.ServiceProvider.GetRequiredService<LocalReminderService>();
+            var barrier = reminderService.TestOnlyWaitForSiloStatusListeners(cancellation.Token);
+            Assert.False(barrier.IsCompleted);
+
+            listener.Release();
+            joinedSilo = Assert.Single(await startTask.WaitAsync(cancellation.Token));
+            await barrier;
+        }
+        finally
+        {
+            listener.Release();
+            statusOracle.UnSubscribeFromSiloStatusEvents(listener);
+            if (joinedSilo is null && startTask is not null)
+            {
+                using var startupCleanup = new CancellationTokenSource(TestConstants.InitTimeout);
+                joinedSilo = Assert.Single(await startTask.WaitAsync(startupCleanup.Token));
+            }
+
+            if (joinedSilo is not null)
+            {
+                using var siloCleanup = new CancellationTokenSource(TestConstants.InitTimeout);
+                await fixture.HostedCluster.StopSiloAsync(joinedSilo, siloCleanup.Token);
+                await fixture.HostedCluster.WaitForLivenessToStabilizeAsync();
+            }
+        }
+    }
+
     [TestSuite("BVT")]
     [TestProvider("None")]
     [Fact, TestCategory("BVT")]
@@ -546,6 +591,31 @@ public override async ValueTask DisposeAsync()
         }
     }
 
+    private sealed class BlockingSiloStatusListener(SiloAddress localSilo) : ISiloStatusListener
+    {
+        private readonly TaskCompletionSource _blocked = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly ManualResetEventSlim _release = new();
+        private int _hasBlocked;
+
+        public void SiloStatusChangeNotification(SiloAddress updatedSilo, SiloStatus status)
+        {
+            if (updatedSilo.Equals(localSilo)
+                || status != SiloStatus.Active
+                || Interlocked.Exchange(ref _hasBlocked, 1) != 0)
+            {
+                return;
+            }
+
+            _blocked.TrySetResult();
+            _release.Wait();
+        }
+
+        public Task WaitUntilBlockedAsync(CancellationToken cancellationToken)
+            => _blocked.Task.WaitAsync(cancellationToken);
+
+        public void Release() => _release.Set();
+    }
+
     private sealed class NullReturningReminderTable : IReminderTable
     {
         private int rangeReadCount;

test/Orleans.Testing.Reminders/ReminderServiceLifecycleHarness.cs

@@ -172,8 +172,12 @@ public async Task WaitForTopologyReconciliationAsync(CancellationToken cancellat
         await Task.WhenAll(
             _cluster.WaitForLivenessToStabilizeAsync().WaitAsync(cancellationToken),
             _cluster.WaitForClusterManifestToStabilizeAsync().WaitAsync(cancellationToken));
+        var activeSilos = _cluster.GetActiveSilos();
+        await Task.WhenAll(activeSilos.Select(silo =>
+            silo.ServiceProvider.GetRequiredService<LocalReminderService>()
+                .TestOnlyWaitForSiloStatusListeners(cancellationToken)));
         await RefreshAsync(cancellationToken);
-        var barriers = _cluster.GetActiveSilos().Select(silo =>
+        var barriers = activeSilos.Select(silo =>
             silo.ServiceProvider.GetRequiredService<LocalReminderService>()
                 .TestOnlyWaitForRangeChangeReconciliation(cancellationToken));
         await Task.WhenAll(barriers);