fix(zookeeper): make membership reads sequential and resilient

src/Orleans.Clustering.ZooKeeper/Orleans.Clustering.ZooKeeper.csproj

@@ -11,6 +11,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\Orleans.Runtime\Orleans.Runtime.csproj" />
+    <PackageReference Include="Polly" />
     <PackageReference Include="ZooKeeperNetEx" />
     <InternalsVisibleTo Include="Orleans.Clustering.ZooKeeper.Tests" />
   </ItemGroup>

src/Orleans.Clustering.ZooKeeper/ZooKeeperConnectionMonitor.cs

@@ -0,0 +1,14 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Orleans.Runtime.Membership;
+
+internal interface IZooKeeperConnectionMonitor
+{
+    long CaptureAttemptGeneration();
+
+    ValueTask<bool> WaitForConnectionAfterAsync(long connectedGeneration, CancellationToken cancellationToken);
+
+    void ReportConnectionLoss(long connectedGeneration);
+}

src/Orleans.Clustering.ZooKeeper/ZooKeeperGatewayListProvider.cs

@@ -7,6 +7,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Orleans.Configuration;
+using Polly;
 
 namespace Orleans.Runtime.Membership
 {
@@ -27,6 +28,8 @@ public class ZooKeeperGatewayListProvider : IGatewayListProvider
         /// </summary>
         private readonly string _deploymentConnectionString;
         private readonly TimeSpan _maxStaleness;
+        private readonly Func<ZooKeeperSession> _createSession;
+        private readonly ResiliencePipeline _readRetryPipeline;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="ZooKeeperGatewayListProvider"/> class.
@@ -44,6 +47,17 @@ public ZooKeeperGatewayListProvider(
             IOptions<ZooKeeperGatewayListProviderOptions> options,
             IOptions<GatewayOptions> gatewayOptions,
             IOptions<ClusterOptions> clusterOptions)
+            : this(logger, options, gatewayOptions, clusterOptions, null, null)
+        {
+        }
+
+        internal ZooKeeperGatewayListProvider(
+            ILogger<ZooKeeperGatewayListProvider> logger,
+            IOptions<ZooKeeperGatewayListProviderOptions> options,
+            IOptions<GatewayOptions> gatewayOptions,
+            IOptions<ClusterOptions> clusterOptions,
+            Func<ZooKeeperSession>? createSession,
+            ResiliencePipeline? readRetryPipeline)
         {
             ArgumentNullException.ThrowIfNull(logger);
             ArgumentNullException.ThrowIfNull(options);
@@ -54,6 +68,8 @@ public ZooKeeperGatewayListProvider(
             _deploymentPath = "/" + clusterOptions.Value.ClusterId;
             _deploymentConnectionString = options.Value.ConnectionString + _deploymentPath;
             _maxStaleness = gatewayOptions.Value.GatewayListRefreshPeriod;
+            _createSession = createSession ?? (() => ZooKeeperBasedMembershipTable.CreateSession(_deploymentConnectionString, _watcher, true));
+            _readRetryPipeline = readRetryPipeline ?? ZooKeeperReadRetryPolicy.CreatePipeline(logger, TimeProvider.System);
         }
 
         /// <summary>
@@ -65,9 +81,13 @@ public ZooKeeperGatewayListProvider(
         /// Returns the list of gateways (silos) that can be used by a client to connect to Orleans cluster.
         /// The Uri is in the form of: "gwy.tcp://IP:port/Generation". See Utils.ToGatewayUri and Utils.ToSiloAddress for more details about Uri format.
         /// </summary>
+        /// <remarks>
+        /// Gateway discovery uses a version-fenced membership snapshot. Native reads retry connection-loss
+        /// failures up to four times on the operation's session before propagating the final failure.
+        /// </remarks>
         public async Task<IList<Uri>> GetGateways()
         {
-            var membershipTableData = await ZooKeeperBasedMembershipTable.ReadAllAsync(this._deploymentConnectionString, this._watcher, CancellationToken.None);
+            var membershipTableData = await ZooKeeperBasedMembershipTable.ReadAsync(_createSession, _readRetryPipeline, null, CancellationToken.None);
             return membershipTableData.Members.Select(e => e.Item1).
                 Where(m => m.Status == SiloStatus.Active && m.ProxyPort != 0).
                 Select(m =>

src/Orleans.Clustering.ZooKeeper/ZooKeeperReadRetryPolicy.cs

@@ -0,0 +1,96 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using org.apache.zookeeper;
+using Polly;
+using Polly.Retry;
+
+namespace Orleans.Runtime.Membership;
+
+internal static partial class ZooKeeperReadRetryPolicy
+{
+    internal const int MaxRetryAttempts = 4;
+    internal static readonly TimeSpan RetryDelay = TimeSpan.FromMilliseconds(250);
+
+    internal static ResiliencePipeline CreatePipeline(ILogger logger, TimeProvider timeProvider) =>
+        new ResiliencePipelineBuilder { TimeProvider = timeProvider }
+            .AddRetry(new RetryStrategyOptions
+            {
+                MaxRetryAttempts = MaxRetryAttempts,
+                BackoffType = DelayBackoffType.Exponential,
+                Delay = RetryDelay,
+                ShouldHandle = new PredicateBuilder().Handle<KeeperException.ConnectionLossException>(),
+                OnRetry = args =>
+                {
+                    LogWarningRetryRead(logger, args.Outcome.Exception!, args.Context.OperationKey!,
+                        args.AttemptNumber + 1, MaxRetryAttempts, args.RetryDelay.TotalMilliseconds);
+                    return default;
+                }
+            })
+            .Build();
+
+    internal static ZooKeeperBasedMembershipTable.NativeOperations Wrap(
+        ZooKeeperBasedMembershipTable.NativeOperations native,
+        ResiliencePipeline pipeline,
+        CancellationToken cancellationToken,
+        IZooKeeperConnectionMonitor? connectionMonitor = null) =>
+        new(
+            path => ExecuteAsync("GetData", () => native.GetData(path), pipeline, cancellationToken, connectionMonitor),
+            path => ExecuteAsync("GetChildren", () => native.GetChildren(path), pipeline, cancellationToken, connectionMonitor),
+            path => ExecuteAsync("Sync", async () =>
+            {
+                await native.Sync(path);
+                return true;
+            }, pipeline, cancellationToken, connectionMonitor),
+            native.Multi,
+            native.SetData);
+
+    private static async Task<T> ExecuteAsync<T>(
+        string operationName,
+        Func<Task<T>> operation,
+        ResiliencePipeline pipeline,
+        CancellationToken cancellationToken,
+        IZooKeeperConnectionMonitor? connectionMonitor)
+    {
+        var context = ResilienceContextPool.Shared.Get(operationName, cancellationToken);
+        long? reconnectAfter = null;
+        try
+        {
+            return await pipeline.ExecuteAsync(async context =>
+            {
+                context.CancellationToken.ThrowIfCancellationRequested();
+                if (reconnectAfter is { } connectedGeneration && connectionMonitor is not null)
+                {
+                    await connectionMonitor.WaitForConnectionAfterAsync(
+                        connectedGeneration,
+                        context.CancellationToken);
+                }
+
+                context.CancellationToken.ThrowIfCancellationRequested();
+                var attemptGeneration = connectionMonitor?.CaptureAttemptGeneration() ?? 0;
+                // Await actual native completion: cancellation ends admission, not an in-flight request.
+                try
+                {
+                    return await operation();
+                }
+                catch (KeeperException.ConnectionLossException)
+                {
+                    connectionMonitor?.ReportConnectionLoss(attemptGeneration);
+                    reconnectAfter = attemptGeneration;
+                    throw;
+                }
+            }, context);
+        }
+        finally
+        {
+            ResilienceContextPool.Shared.Return(context);
+        }
+    }
+
+    [LoggerMessage(
+        Level = LogLevel.Warning,
+        Message = "ZooKeeper {Operation} lost its connection. Retrying in {DelayMilliseconds}ms ({Retry}/{MaxRetries}).")]
+    private static partial void LogWarningRetryRead(
+        ILogger logger, Exception exception, string operation, int retry, int maxRetries, double delayMilliseconds);
+}

src/Orleans.Clustering.ZooKeeper/ZooKeeperSession.cs

@@ -0,0 +1,73 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Orleans.Runtime.Membership;
+
+internal sealed class ZooKeeperSession
+{
+    private readonly ZooKeeperBasedMembershipTable.NativeOperations _operations;
+    private readonly Func<Task> _close;
+    // Bind synchronously; the owned task supplies the completion semantics.
+    private readonly TaskCompletionSource<Task> _completion = new();
+
+    internal ZooKeeperSession(
+        ZooKeeperBasedMembershipTable.NativeOperations operations,
+        Func<Task> close,
+        IZooKeeperConnectionMonitor? connectionMonitor = null)
+    {
+        _operations = operations;
+        _close = close;
+        ConnectionMonitor = connectionMonitor;
+        Completion = _completion.Task.Unwrap();
+        Completion.Ignore();
+    }
+
+    internal Task Completion { get; }
+    internal IZooKeeperConnectionMonitor? ConnectionMonitor { get; }
+    internal ZooKeeperBasedMembershipTable.NativeOperations Operations => _operations;
+
+    internal static Task<T> ExecuteAsync<T>(
+        Func<ZooKeeperSession> createSession,
+        Func<ZooKeeperBasedMembershipTable.NativeOperations, Task<T>> operation,
+        CancellationToken cancellationToken)
+        => ExecuteSessionAsync(createSession, session => operation(session.Operations), cancellationToken);
+
+    internal static Task<T> ExecuteSessionAsync<T>(
+        Func<ZooKeeperSession> createSession,
+        Func<ZooKeeperSession, Task<T>> operation,
+        CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+        var session = createSession();
+        var completion = session.RunAsync(operation);
+        session._completion.SetResult(completion);
+        return ZooKeeperBasedMembershipTable.AwaitOperationAsync(completion, cancellationToken);
+    }
+
+    private async Task<T> RunAsync<T>(Func<ZooKeeperSession, Task<T>> operation)
+    {
+        T result;
+        try
+        {
+            result = await operation(this);
+        }
+        catch (Exception primary)
+        {
+            try
+            {
+                await _close();
+            }
+            catch (Exception secondary)
+            {
+                throw new AggregateException("The ZooKeeper operation and its close both failed.", primary, secondary);
+            }
+
+            throw;
+        }
+
+        // The callback joins native requests before this tokenless close.
+        await _close();
+        return result;
+    }
+}

test/Extensions/Orleans.Clustering.ZooKeeper.Tests/Orleans.Clustering.ZooKeeper.Tests.csproj

@@ -6,7 +6,9 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
     <ProjectReference Include="$(SourceRoot)src\Orleans.Clustering.ZooKeeper\Orleans.Clustering.ZooKeeper.csproj" />
+    <ProjectReference Include="$(SourceRoot)src\Orleans.Clustering.TestKit\Orleans.Clustering.TestKit.csproj" />
     <ProjectReference Include="$(SourceRoot)test\Orleans.Runtime.Internal.Tests\Orleans.Runtime.Internal.Tests.csproj" />
   </ItemGroup>
 </Project>