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/ZooKeeperBasedMembershipTable.cs

@@ -12,6 +12,7 @@
 using Microsoft.Extensions.Options;
 using Orleans.Configuration;
 using Orleans.Runtime.Host;
+using Polly;
 
 namespace Orleans.Runtime.Membership
 {
@@ -43,6 +44,9 @@ public partial class ZooKeeperBasedMembershipTable : IMembershipTable
         private const int ZOOKEEPER_SESSION_TIMEOUT = 10_000;
 
         private readonly ZooKeeperWatcher watcher;
+        private readonly Func<bool, ZooKeeperSession> _createSession;
+        private readonly ResiliencePipeline _readRetryPipeline;
+        private readonly Action<Task>? _observeOperation;
 
         /// <summary>
         /// The deployment connection string. for eg. "192.168.1.1,192.168.1.2/ClusterId"
@@ -73,6 +77,17 @@ public ZooKeeperBasedMembershipTable(
             ILogger<ZooKeeperBasedMembershipTable> logger,
             IOptions<ZooKeeperClusteringSiloOptions> membershipTableOptions,
             IOptions<ClusterOptions> clusterOptions)
+            : this(logger, membershipTableOptions, clusterOptions, null, null)
+        {
+        }
+
+        internal ZooKeeperBasedMembershipTable(
+            ILogger<ZooKeeperBasedMembershipTable> logger,
+            IOptions<ZooKeeperClusteringSiloOptions> membershipTableOptions,
+            IOptions<ClusterOptions> clusterOptions,
+            Func<bool, ZooKeeperSession>? createSession,
+            ResiliencePipeline? readRetryPipeline,
+            Action<Task>? observeOperation = null)
         {
             ArgumentNullException.ThrowIfNull(logger);
             ArgumentNullException.ThrowIfNull(membershipTableOptions);
@@ -84,6 +99,9 @@ public ZooKeeperBasedMembershipTable(
             this.clusterPath = "/" + clusterOptions.Value.ClusterId;
             rootConnectionString = options.ConnectionString;
             deploymentConnectionString = options.ConnectionString + this.clusterPath;
+            _createSession = createSession ?? (readOnly => CreateSession(deploymentConnectionString, watcher, readOnly));
+            _readRetryPipeline = readRetryPipeline ?? ZooKeeperReadRetryPolicy.CreatePipeline(logger, TimeProvider.System);
+            _observeOperation = observeOperation;
         }
 
         /// <summary>
@@ -128,10 +146,13 @@ await UsingZookeeper(rootConnectionString, async zk =>
         public Task<MembershipTableData> ReadRow(SiloAddress siloAddress) => ReadRowAsync(siloAddress, CancellationToken.None);
 
         /// <inheritdoc />
+        /// <remarks>
+        /// Connection-loss failures in native reads are retried up to four times on the operation's session.
+        /// The table and child versions fence each complete snapshot pass.
+        /// </remarks>
         public Task<MembershipTableData> ReadRowAsync(SiloAddress siloAddress, CancellationToken cancellationToken = default)
         {
-            return UsingZookeeper(zk => ReadCoreAsync(zk, siloAddress, cancellationToken),
-                this.deploymentConnectionString, this.watcher, cancellationToken, canBeReadOnly: true);
+            return ReadAsync(() => _createSession(true), _readRetryPipeline, siloAddress, cancellationToken);
         }
 
         /// <summary>
@@ -145,15 +166,36 @@ public Task<MembershipTableData> ReadRowAsync(SiloAddress siloAddress, Cancellat
         public Task<MembershipTableData> ReadAll() => ReadAllAsync(CancellationToken.None);
 
         /// <inheritdoc />
+        /// <remarks>
+        /// Rows are read concurrently on an operation-owned connection,
+        /// with each membership record read before its heartbeat.
+        /// Table and child-version checks fence the complete snapshot.
+        /// Connection-loss failures in native reads are retried up to four times on the same session.
+        /// Caller cancellation stops further requests while admitted requests and client close complete.
+        /// </remarks>
         public Task<MembershipTableData> ReadAllAsync(CancellationToken cancellationToken = default)
         {
-            return ReadAllAsync(this.deploymentConnectionString, this.watcher, cancellationToken);
+            return ReadAsync(() => _createSession(true), _readRetryPipeline, null, cancellationToken);
         }
 
-        internal static Task<MembershipTableData> ReadAllAsync(string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken)
+        internal static Task<MembershipTableData> ReadAsync(
+            Func<ZooKeeperSession> createSession,
+            ResiliencePipeline pipeline,
+            SiloAddress? siloAddress,
+            CancellationToken cancellationToken)
         {
-            return UsingZookeeper(zk => ReadCoreAsync(zk, null, cancellationToken),
-                deploymentConnectionString, watcher, cancellationToken, canBeReadOnly: true);
+            return ZooKeeperSession.ExecuteAsync(createSession,
+                native => ReadCoreAsync(ZooKeeperReadRetryPolicy.Wrap(native, pipeline, cancellationToken), siloAddress, cancellationToken),
+                cancellationToken);
+        }
+
+        internal static ZooKeeperSession CreateSession(string connectionString, ZooKeeperWatcher watcher, bool readOnly)
+        {
+            var client = new ZooKeeper(connectionString, ZOOKEEPER_SESSION_TIMEOUT, watcher, readOnly);
+            return new ZooKeeperSession(
+                new NativeOperations(path => client.getDataAsync(path), path => client.getChildrenAsync(path),
+                    client.sync, operations => client.multiAsync(operations), client.setDataAsync),
+                client.closeAsync);
         }
 
         internal static async Task<MembershipTableData> ReadCoreAsync(
@@ -179,36 +221,36 @@ internal static async Task<MembershipTableData> ReadCoreAsync(
                     addresses = [siloAddress];
                 }
 
+                var rows = new List<Tuple<MembershipEntry, string>>();
+                KeeperException.NoNodeException? missingRow = null;
+                cancellationToken.ThrowIfCancellationRequested();
                 var pendingRows = Task.WhenAll(addresses.Select(address => GetRow(zk, address, siloAddress is not null, cancellationToken)));
-                Tuple<MembershipEntry, string>?[] rows;
                 try
                 {
-                    rows = await pendingRows;
+                    rows.AddRange((await pendingRows).OfType<Tuple<MembershipEntry, string>>());
                 }
-                catch (KeeperException.NoNodeException)
+                catch (KeeperException.NoNodeException exception)
                 {
-                    // Observe every parallel read: a removed row must not hide another request's failure.
-                    var failure = pendingRows.Exception!.InnerExceptions.FirstOrDefault(exception => exception is not KeeperException.NoNodeException);
+                    // Join every admitted read so a missing row cannot hide another native failure.
+                    var failure = pendingRows.Exception!.InnerExceptions.FirstOrDefault(error => error is not KeeperException.NoNodeException);
                     if (failure is not null)
                     {
                         System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(failure).Throw();
                     }
 
-                    cancellationToken.ThrowIfCancellationRequested();
-                    var current = await zk.GetData("/");
-                    if (SameVersion(before, current.Stat))
-                    {
-                        throw;
-                    }
-
-                    continue;
+                    missingRow = exception;
                 }
 
                 cancellationToken.ThrowIfCancellationRequested();
                 var after = await zk.GetData("/");
                 if (SameVersion(before, after.Stat))
                 {
-                    return new MembershipTableData(rows.OfType<Tuple<MembershipEntry, string>>().ToList(), ConvertToTableVersion(after.Stat));
+                    if (missingRow is not null)
+                    {
+                        System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(missingRow).Throw();
+                    }
+
+                    return new MembershipTableData(rows, ConvertToTableVersion(after.Stat));
                 }
             }
         }
@@ -238,14 +280,18 @@ private static bool SameVersion(Stat before, Stat after) =>
         public Task<bool> InsertRow(MembershipEntry entry, TableVersion tableVersion) => InsertRowAsync(entry, tableVersion, CancellationToken.None);
 
         /// <inheritdoc />
+        /// <remarks>
+        /// The conditional transaction is submitted once. Connection loss propagates when its
+        /// commit outcome is unknown, preserving the caller's ability to resolve that outcome.
+        /// </remarks>
         public Task<bool> InsertRowAsync(MembershipEntry entry, TableVersion tableVersion, CancellationToken cancellationToken = default)
         {
             ArgumentNullException.ThrowIfNull(entry);
             ArgumentNullException.ThrowIfNull(tableVersion);
             cancellationToken.ThrowIfCancellationRequested();
 
-            return UsingZookeeper(zk => InsertRowCoreAsync(zk, entry, tableVersion, cancellationToken),
-                this.deploymentConnectionString, this.watcher, cancellationToken);
+            return ZooKeeperSession.ExecuteAsync(() => _createSession(false),
+                zk => InsertRowCoreAsync(zk, entry, tableVersion, cancellationToken), cancellationToken);
         }
 
         internal static async Task<bool> InsertRowCoreAsync(
@@ -300,15 +346,19 @@ await zk.Multi(
         public Task<bool> UpdateRow(MembershipEntry entry, string etag, TableVersion tableVersion) => UpdateRowAsync(entry, etag, tableVersion, CancellationToken.None);
 
         /// <inheritdoc />
+        /// <remarks>
+        /// The conditional transaction is submitted once. Connection loss propagates when its
+        /// commit outcome is unknown, preserving the caller's ability to resolve that outcome.
+        /// </remarks>
         public Task<bool> UpdateRowAsync(MembershipEntry entry, string etag, TableVersion tableVersion, CancellationToken cancellationToken = default)
         {
             ArgumentNullException.ThrowIfNull(entry);
             ArgumentNullException.ThrowIfNull(etag);
             ArgumentNullException.ThrowIfNull(tableVersion);
             cancellationToken.ThrowIfCancellationRequested();
 
-            return UsingZookeeper(zk => UpdateRowCoreAsync(zk, entry, etag, tableVersion, cancellationToken),
-                this.deploymentConnectionString, this.watcher, cancellationToken);
+            return ZooKeeperSession.ExecuteAsync(() => _createSession(false),
+                zk => UpdateRowCoreAsync(zk, entry, etag, tableVersion, cancellationToken), cancellationToken);
         }
 
         internal static async Task<bool> UpdateRowCoreAsync(
@@ -431,7 +481,7 @@ internal sealed class NativeOperations(
             internal Func<string, byte[], int, Task<Stat>> SetData { get; } = setData;
         }
 
-        private static async Task<T> UsingZookeeper<T>(Func<NativeOperations, Task<T>> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken, bool canBeReadOnly = false)
+        private async Task<T> UsingZookeeper<T>(Func<NativeOperations, Task<T>> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken, bool canBeReadOnly = false)
         {
             cancellationToken.ThrowIfCancellationRequested();
             var operation = ZooKeeper.Using(deploymentConnectionString, ZOOKEEPER_SESSION_TIMEOUT, watcher, zk =>
@@ -442,13 +492,15 @@ private static async Task<T> UsingZookeeper<T>(Func<NativeOperations, Task<T>> z
                     operations => zk.multiAsync(operations), zk.setDataAsync));
             }, canBeReadOnly);
 
-            return await AwaitOperationAsync(operation, cancellationToken);
+            return await AwaitOperationAsync(operation, cancellationToken, _observeOperation);
         }
 
-        internal static async Task<T> AwaitOperationAsync<T>(Task<T> operation, CancellationToken cancellationToken)
+        internal static async Task<T> AwaitOperationAsync<T>(
+            Task<T> operation, CancellationToken cancellationToken, Action<Task>? observeOperation = null)
         {
             // ZooKeeperNetEx is tokenless. Keep the client alive until pending requests and
             // asynchronous disposal finish, observing failures even if the caller stops waiting.
+            observeOperation?.Invoke(operation);
             operation.Ignore();
             return await operation.WaitAsync(cancellationToken);
         }
@@ -462,6 +514,7 @@ private async Task UsingZookeeper(string connectString, Func<ZooKeeper, Task> zk
                 return zkMethod(zk);
             });
 
+            _observeOperation?.Invoke(operation);
             operation.Ignore();
             await operation.WaitAsync(cancellationToken);
         }

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,75 @@
+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) =>
+        new(
+            path => ExecuteAsync("GetData", () => native.GetData(path), pipeline, cancellationToken),
+            path => ExecuteAsync("GetChildren", () => native.GetChildren(path), pipeline, cancellationToken),
+            path => ExecuteAsync("Sync", async () =>
+            {
+                await native.Sync(path);
+                return true;
+            }, pipeline, cancellationToken),
+            native.Multi,
+            native.SetData);
+
+    private static async Task<T> ExecuteAsync<T>(
+        string operationName,
+        Func<Task<T>> operation,
+        ResiliencePipeline pipeline,
+        CancellationToken cancellationToken)
+    {
+        var context = ResilienceContextPool.Shared.Get(operationName, cancellationToken);
+        try
+        {
+            return await pipeline.ExecuteAsync(async context =>
+            {
+                context.CancellationToken.ThrowIfCancellationRequested();
+                // Await actual native completion: cancellation ends admission, not an in-flight request.
+                return await operation();
+            }, 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);
+}