fix(directory): wait for active membership during startup

src/Orleans.Runtime/GrainDirectory/DistributedGrainDirectory.cs

@@ -241,22 +241,16 @@ internal Task UnregisterAsync(GrainAddress address, CancellationToken cancellati
         {
             cancellationToken.ThrowIfCancellationRequested();
             var initialRecoveryMembershipVersion = _recoveryMembershipVersion;
-            if (view.Version.Value < initialRecoveryMembershipVersion || !view.TryGetOwner(grainId, out var owner, out var partitionReference))
+            var resolved = await GetViewWithOwnerAsync(grainId, view, initialRecoveryMembershipVersion, cancellationToken);
+            if (resolved is not { } ownerView)
             {
-                // If there are no members and this view is current, bail out with the default return value.
-                // Otherwise, wait for the directory to observe the newer cluster membership view.
-                if (view.Members.Length == 0
-                    && view.Version.Value > 0
-                    && LatestClusterMembershipSnapshot.Version <= view.Version)
-                {
-                    return default;
-                }
-
-                var targetVersion = Math.Max(view.Version.Value + 1, initialRecoveryMembershipVersion);
-                view = await _membershipService.RefreshViewAsync(new(targetVersion), cancellationToken);
-                continue;
+                return default;
             }
 
+            view = ownerView.View;
+            var owner = ownerView.Owner;
+            var partitionReference = ownerView.PartitionReference;
+
 #if false
             if (logger.IsEnabled(LogLevel.Trace))
             {
@@ -322,6 +316,47 @@ exception is OrleansMessageRejectionException or OperationCanceledException
         }
     }
 
+    private async ValueTask<(DirectoryMembershipSnapshot View, SiloAddress Owner, IGrainDirectoryPartition PartitionReference)?> GetViewWithOwnerAsync(
+        GrainId grainId,
+        DirectoryMembershipSnapshot view,
+        long minimumVersion,
+        CancellationToken cancellationToken)
+    {
+        while (true)
+        {
+            if (view.Version.Value >= minimumVersion
+                && view.TryGetOwner(grainId, out var owner, out var partitionReference))
+            {
+                return (view, owner, partitionReference);
+            }
+
+            // Cluster membership advances before the directory's asynchronous view. If the latest cluster view
+            // still has an active member, this empty directory view is stale; otherwise it is terminal, as on shutdown.
+            if (view.Members.Length == 0
+                && view.Version.Value > 0
+                && !HasActiveMembers(LatestClusterMembershipSnapshot))
+            {
+                return null;
+            }
+
+            var targetVersion = Math.Max(view.Version.Value + 1, minimumVersion);
+            view = await _membershipService.RefreshViewAsync(new(targetVersion), cancellationToken);
+        }
+    }
+
+    private static bool HasActiveMembers(ClusterMembershipSnapshot snapshot)
+    {
+        foreach (var member in snapshot.Members.Values)
+        {
+            if (member.Status == SiloStatus.Active)
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     internal static bool CanInvokeClusterMember(ClusterMembershipSnapshot snapshot, SiloAddress siloAddress)
         => snapshot.GetSiloStatus(siloAddress) is SiloStatus.Active or SiloStatus.Joining or SiloStatus.ShuttingDown;
 
@@ -599,6 +634,16 @@ private async Task RequestExpiredLeaseCleanups()
         return owner;
     }
 
+    async Task<SiloAddress?> ITestHooks.WaitForPrimaryForGrain(GrainId grainId, CancellationToken cancellationToken)
+    {
+        var result = await GetViewWithOwnerAsync(
+            grainId,
+            _membershipService.CurrentView,
+            _recoveryMembershipVersion,
+            cancellationToken);
+        return result?.Owner;
+    }
+
     async Task<GrainAddress?> ITestHooks.GetLocalRecord(GrainId grainId)
     {
         var view = _membershipService.CurrentView;
@@ -617,6 +662,7 @@ private async Task RequestExpiredLeaseCleanups()
     internal interface ITestHooks
     {
         SiloAddress? GetPrimaryForGrain(GrainId grainId);
+        Task<SiloAddress?> WaitForPrimaryForGrain(GrainId grainId, CancellationToken cancellationToken);
         Task<GrainAddress?> GetLocalRecord(GrainId grainId);
     }
 

test/Orleans.GrainDirectory.Tests/GrainDirectory/DistributedGrainDirectoryTests.cs

@@ -1,5 +1,11 @@
 #nullable enable
+using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using Orleans.Configuration;
 using Orleans.GrainDirectory;
 using Orleans.Runtime.GrainDirectory;
 using Orleans.TestingHost;
@@ -23,3 +29,103 @@ public sealed class DefaultGrainDirectoryTests(DefaultClusterFixture fixture, IT
     protected override IGrainDirectory CreateGrainDirectory() =>
         Primary.SiloHost.Services.GetRequiredService<GrainDirectoryResolver>().DefaultGrainDirectory!;
 }
+
+[TestSuite("BVT")]
+[TestProvider("None")]
+[TestArea("GrainDirectory")]
+[TestCategory("BVT"), TestCategory("Directory")]
+public sealed class DistributedGrainDirectoryMembershipTests
+{
+    [Fact]
+    public async Task OwnerResolutionWaitsForActiveDirectoryMembership()
+    {
+        var builder = new InProcessTestClusterBuilder(1);
+#pragma warning disable ORLEANSEXP003
+        builder.Options.UseDistributedGrainDirectory = true;
+#pragma warning restore ORLEANSEXP003
+        builder.ConfigureSilo((_, siloBuilder) => siloBuilder.ConfigureServices(services =>
+        {
+            services.AddSingleton<ControlledMembershipService>();
+            services.Replace(ServiceDescriptor.Singleton(static sp => new DirectoryMembershipService(
+                sp.GetRequiredService<ControlledMembershipService>(),
+                sp.GetRequiredService<IInternalGrainFactory>(),
+                sp.GetRequiredService<ILogger<DirectoryMembershipService>>(),
+                sp.GetRequiredService<IOptions<GrainDirectoryOptions>>().Value.PartitionsPerSilo,
+                DirectoryMembershipSnapshot.DefaultGetRingBoundaries)));
+        }));
+        var cluster = builder.Build();
+
+        try
+        {
+            await cluster.DeployAsync();
+            var silo = cluster.Silos[0];
+            var membership = silo.ServiceProvider.GetRequiredService<DirectoryMembershipService>();
+            var directory = silo.ServiceProvider.GetRequiredService<DistributedGrainDirectory>();
+            var testHooks = (DistributedGrainDirectory.ITestHooks)directory;
+            var controlledMembership = silo.ServiceProvider.GetRequiredService<ControlledMembershipService>();
+            using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+
+            await controlledMembership.ActiveUpdateBlocked.WaitAsync(timeout.Token);
+            var siloAddress = silo.SiloAddress;
+            Assert.Empty(membership.CurrentView.Members);
+
+            var grainId = GrainId.Create("directory-readiness", "registration");
+            var primaryTask = testHooks.WaitForPrimaryForGrain(grainId, timeout.Token);
+            Assert.False(primaryTask.IsCompleted);
+
+            controlledMembership.ReleaseActiveUpdate();
+
+            Assert.Equal(siloAddress, await primaryTask);
+        }
+        finally
+        {
+            await cluster.DisposeAsync();
+        }
+    }
+
+    private sealed class ControlledMembershipService(IClusterMembershipService inner) : IClusterMembershipService
+    {
+        private readonly TaskCompletionSource _activeUpdateBlocked = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly TaskCompletionSource _releaseActiveUpdate = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private int _hasBlockedActiveUpdate;
+
+        public ClusterMembershipSnapshot CurrentSnapshot => inner.CurrentSnapshot;
+
+        public IAsyncEnumerable<ClusterMembershipSnapshot> MembershipUpdates => GetMembershipUpdates();
+
+        public Task ActiveUpdateBlocked => _activeUpdateBlocked.Task;
+
+        public ValueTask Refresh(MembershipVersion minimumVersion = default, CancellationToken cancellationToken = default) =>
+            inner.Refresh(minimumVersion, cancellationToken);
+
+        public Task<bool> TryKill(SiloAddress siloAddress) => inner.TryKill(siloAddress);
+
+        public void ReleaseActiveUpdate() => _releaseActiveUpdate.TrySetResult();
+
+        private async IAsyncEnumerable<ClusterMembershipSnapshot> GetMembershipUpdates(
+            [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            await foreach (var update in inner.MembershipUpdates.WithCancellation(cancellationToken))
+            {
+                if (update.Members.Values.Any(static member => member.Status == SiloStatus.Active)
+                    && Interlocked.CompareExchange(ref _hasBlockedActiveUpdate, 1, 0) == 0)
+                {
+                    var joiningVersion = new MembershipVersion(update.Version.Value - 1);
+                    if (joiningVersion <= MembershipVersion.MinValue)
+                    {
+                        throw new InvalidOperationException($"Expected an Active membership version greater than one, encountered {update.Version}.");
+                    }
+
+                    var joiningMembers = update.Members.ToImmutableDictionary(
+                        static member => member.Key,
+                        static member => new ClusterMember(member.Key, SiloStatus.Joining, member.Value.Name));
+                    yield return new ClusterMembershipSnapshot(joiningMembers, joiningVersion);
+                    _activeUpdateBlocked.TrySetResult();
+                    await _releaseActiveUpdate.Task.WaitAsync(cancellationToken);
+                }
+
+                yield return update;
+            }
+        }
+    }
+}