From 706f5d228e0ca8c2f0ce260248396223af65b03e Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 17 Sep 2026 15:19:43 -0700 Subject: [PATCH 1/8] test(membership): correct legacy cleanup baseline --- .../MembershipTableTestsBase.cs | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/test/Orleans.Runtime.Internal.Tests/MembershipTests/MembershipTableTestsBase.cs b/test/Orleans.Runtime.Internal.Tests/MembershipTests/MembershipTableTestsBase.cs index e99d16a1c5e..0ad65748915 100644 --- a/test/Orleans.Runtime.Internal.Tests/MembershipTests/MembershipTableTestsBase.cs +++ b/test/Orleans.Runtime.Internal.Tests/MembershipTests/MembershipTableTestsBase.cs @@ -460,17 +460,18 @@ protected async Task MembershipTable_CleanupDefunctSiloEntries(bool extendedProt Assert.True(ok, "InsertRow Dead failed"); newTableVersion = table.Version.Next(); - var oldEntryJoining = CreateMembershipEntryForTest(); - oldEntryJoining.IAmAliveTime = oldEntryJoining.IAmAliveTime.AddDays(-10); - oldEntryJoining.StartTime = oldEntryJoining.StartTime.AddDays(-10); - oldEntryJoining.Status = SiloStatus.Joining; - ok = await membershipTable.InsertRowAsync(oldEntryJoining, newTableVersion, cancellationToken); + var oldEntryActive = CreateMembershipEntryForTest(); + oldEntryActive.IAmAliveTime = oldEntryActive.IAmAliveTime.AddDays(-10); + oldEntryActive.StartTime = oldEntryActive.StartTime.AddDays(-10); + oldEntryActive.Status = SiloStatus.Active; + ok = await membershipTable.InsertRowAsync(oldEntryActive, newTableVersion, cancellationToken); table = await membershipTable.ReadAllAsync(cancellationToken); - Assert.True(ok, "InsertRow Joining failed"); + Assert.True(ok, "InsertRow Active failed"); newTableVersion = table.Version.Next(); var newEntry = CreateMembershipEntryForTest(); + newEntry.Status = SiloStatus.Active; ok = await membershipTable.InsertRowAsync(newEntry, newTableVersion, cancellationToken); Assert.True(ok, "InsertRow failed"); @@ -481,8 +482,9 @@ protected async Task MembershipTable_CleanupDefunctSiloEntries(bool extendedProt Assert.Equal(3, data.Members.Count); - // Every status other than Active should get cleared out if old - foreach (var siloStatus in Enum.GetValues()) + var activeAddresses = new List { oldEntryActive.SiloAddress, newEntry.SiloAddress }; + var deadAddresses = new List { oldEntryDead.SiloAddress }; + foreach (var siloStatus in new[] { SiloStatus.Active, SiloStatus.Dead }) { var oldEntry = CreateMembershipEntryForTest(); oldEntry.IAmAliveTime = oldEntry.IAmAliveTime.AddDays(-10); @@ -493,6 +495,7 @@ protected async Task MembershipTable_CleanupDefunctSiloEntries(bool extendedProt Assert.True(ok, "InsertRow failed"); + (siloStatus == SiloStatus.Active ? activeAddresses : deadAddresses).Add(oldEntry.SiloAddress); newTableVersion = table.Version.Next(); } @@ -501,7 +504,10 @@ protected async Task MembershipTable_CleanupDefunctSiloEntries(bool extendedProt data = await membershipTable.ReadAllAsync(cancellationToken); logger.LogInformation("Membership.ReadAll returned TableVersion={TableVersion} Data={Data}", data.Version, data); - Assert.Equal(2, data.Members.Count); + Assert.Equal(3, data.Members.Count); + Assert.All(data.Members, row => Assert.Equal(SiloStatus.Active, row.Item1.Status)); + Assert.All(activeAddresses, address => Assert.NotNull(data.TryGet(address))); + Assert.All(deadAddresses, address => Assert.Null(data.TryGet(address))); } // Utility methods From ba6ef35350a0dd18f95982887333db2ea137f9d7 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 17 Sep 2026 16:14:50 -0700 Subject: [PATCH 2/8] fix(zookeeper): preserve membership consistency --- .../ZooKeeperBasedMembershipTable.cs | 336 ++++++-- .../ZooKeeperBasedMembershipTableUnitTests.cs | 810 ++++++++++++++++++ .../ZooKeeperNativeFake.cs | 186 ++++ 3 files changed, 1249 insertions(+), 83 deletions(-) create mode 100644 test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index 2b260581146..577d6302d81 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -29,7 +29,7 @@ namespace Orleans.Runtime.Membership /// Every Orleans deployment has a node /UniqueDeploymentId /// Every Silo's state is saved in /UniqueDeploymentId/IP:Port@Gen /// Every Silo's IAmAlive is saved in /UniqueDeploymentId/IP:Port@Gen/IAmAlive - /// IAmAlive is saved in a separate node because its updates are unconditional. + /// IAmAlive is saved in a separate node so compare-and-set heartbeat updates preserve the membership row's version. /// /// a node's ZK version is its ETag: /// the table version is the version of /UniqueDeploymentId @@ -131,25 +131,8 @@ await UsingZookeeper(rootConnectionString, async zk => /// public Task ReadRowAsync(SiloAddress siloAddress, CancellationToken cancellationToken = default) { - return UsingZookeeper(async zk => - { - var getRowTask = GetRow(zk, siloAddress, cancellationToken); - var getTableNodeTask = zk.getDataAsync("/");//get the current table version - - List> rows = new List>(1); - try - { - await Task.WhenAll(getRowTask, getTableNodeTask); - rows.Add(await getRowTask); - } - catch (KeeperException.NoNodeException) - { - //that's ok because orleans expects an empty list in case of a missing row - } - - var tableVersion = ConvertToTableVersion((await getTableNodeTask).Stat); - return new MembershipTableData(rows, tableVersion); - }, this.deploymentConnectionString, this.watcher, cancellationToken, true); + return UsingZookeeper(zk => ReadCoreAsync(zk, siloAddress, cancellationToken), + this.deploymentConnectionString, this.watcher, cancellationToken); } /// @@ -170,21 +153,69 @@ public Task ReadAllAsync(CancellationToken cancellationToke internal static Task ReadAllAsync(string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken) { - return UsingZookeeper(async zk => + return UsingZookeeper(zk => ReadCoreAsync(zk, null, cancellationToken), + deploymentConnectionString, watcher, cancellationToken); + } + + internal static async Task ReadCoreAsync( + NativeOperations zk, SiloAddress? siloAddress, CancellationToken cancellationToken) + { + while (true) { - var childrenResult = await zk.getChildrenAsync("/");//get all the child nodes (without the data) + cancellationToken.ThrowIfCancellationRequested(); + await zk.Sync("/"); + cancellationToken.ThrowIfCancellationRequested(); + Stat before; + IEnumerable addresses; + if (siloAddress is null) + { + var children = await zk.GetChildren("/"); + before = children.Stat; + addresses = children.Children.Select(SiloAddress.FromParsableString); + } + else + { + before = (await zk.GetData("/")).Stat; + addresses = [siloAddress]; + } - var childrenTasks = //get the data from each child node - childrenResult.Children.Select(child => GetRow(zk, SiloAddress.FromParsableString(child), cancellationToken)).ToList(); + var pendingRows = Task.WhenAll(addresses.Select(address => GetRow(zk, address, siloAddress is not null, cancellationToken))); + Tuple?[] rows; + try + { + rows = await pendingRows; + } + catch (KeeperException.NoNodeException) + { + // 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); + if (failure is not null) + { + System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(failure).Throw(); + } - var childrenTaskResults = await Task.WhenAll(childrenTasks); + cancellationToken.ThrowIfCancellationRequested(); + var current = await zk.GetData("/"); + if (SameVersion(before, current.Stat)) + { + throw; + } - var tableVersion = ConvertToTableVersion(childrenResult.Stat);//this is the current table version + continue; + } - return new MembershipTableData(childrenTaskResults.ToList(), tableVersion); - }, deploymentConnectionString, watcher, cancellationToken, true); + cancellationToken.ThrowIfCancellationRequested(); + var after = await zk.GetData("/"); + if (SameVersion(before, after.Stat)) + { + return new MembershipTableData(rows.OfType>().ToList(), ConvertToTableVersion(after.Stat)); + } + } } + private static bool SameVersion(Stat before, Stat after) => + before.getVersion() == after.getVersion() && before.getCversion() == after.getCversion(); + /// /// Atomically tries to insert (add) a new MembershipEntry for one silo and also update the TableVersion. /// If operation succeeds, the following changes would be made to the table: @@ -213,6 +244,14 @@ public Task InsertRowAsync(MembershipEntry entry, TableVersion tableVersio ArgumentNullException.ThrowIfNull(tableVersion); cancellationToken.ThrowIfCancellationRequested(); + return UsingZookeeper(zk => InsertRowCoreAsync(zk, entry, tableVersion, cancellationToken), + this.deploymentConnectionString, this.watcher, cancellationToken); + } + + internal static async Task InsertRowCoreAsync( + NativeOperations zk, MembershipEntry entry, TableVersion tableVersion, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); string rowPath = ConvertToRowPath(entry.SiloAddress); string rowIAmAlivePath = ConvertToRowIAmAlivePath(entry.SiloAddress); byte[] newRowData = Serialize(entry); @@ -220,10 +259,26 @@ public Task InsertRowAsync(MembershipEntry entry, TableVersion tableVersio int expectedTableVersion = int.Parse(tableVersion.VersionEtag, CultureInfo.InvariantCulture); - return TryTransaction(t => t - .setData("/", null, expectedTableVersion)//increments the version of node "/" - .create(rowPath, newRowData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT) - .create(rowIAmAlivePath, newRowIAmAliveData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), cancellationToken); + try + { + await zk.Multi( + [ + Op.setData("/", null, expectedTableVersion), + Op.create(rowPath, newRowData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), + Op.create(rowIAmAlivePath, newRowIAmAliveData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT) + ]); + return true; + } + catch (KeeperException e) when (e is KeeperException.NodeExistsException or KeeperException.BadVersionException) + { + return false; + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); + return false; + } } /// @@ -258,18 +313,58 @@ public Task UpdateRowAsync(MembershipEntry entry, string etag, TableVersio ArgumentNullException.ThrowIfNull(tableVersion); cancellationToken.ThrowIfCancellationRequested(); + return UsingZookeeper(zk => UpdateRowCoreAsync(zk, entry, etag, tableVersion, cancellationToken), + this.deploymentConnectionString, this.watcher, cancellationToken); + } + + internal static async Task UpdateRowCoreAsync( + NativeOperations zk, MembershipEntry entry, string etag, TableVersion tableVersion, CancellationToken cancellationToken) + { string rowPath = ConvertToRowPath(entry.SiloAddress); string rowIAmAlivePath = ConvertToRowIAmAlivePath(entry.SiloAddress); var newRowData = Serialize(entry); - var newRowIAmAliveData = Serialize(entry.IAmAliveTime); - int expectedTableVersion = int.Parse(tableVersion.VersionEtag, CultureInfo.InvariantCulture); int expectedRowVersion = int.Parse(etag, CultureInfo.InvariantCulture); - return TryTransaction(t => t - .setData("/", null, expectedTableVersion)//increments the version of node "/" - .setData(rowPath, newRowData, expectedRowVersion)//increments the version of node "/IP:Port@Gen" - .setData(rowIAmAlivePath, newRowIAmAliveData), cancellationToken); + while (true) + { + cancellationToken.ThrowIfCancellationRequested(); + if ((await zk.GetData("/")).Stat.getVersion() != expectedTableVersion) + { + return false; + } + + try + { + cancellationToken.ThrowIfCancellationRequested(); + if ((await zk.GetData(rowPath)).Stat.getVersion() != expectedRowVersion) + { + return false; + } + + cancellationToken.ThrowIfCancellationRequested(); + var heartbeat = await zk.GetData(rowIAmAlivePath); + var time = new DateTime(Math.Max(entry.IAmAliveTime.Ticks, Deserialize(heartbeat.Data).Ticks), DateTimeKind.Utc); + cancellationToken.ThrowIfCancellationRequested(); + await zk.Multi( + [ + Op.setData("/", null, expectedTableVersion), + Op.setData(rowPath, newRowData, expectedRowVersion), + Op.setData(rowIAmAlivePath, Serialize(time), heartbeat.Stat.getVersion()) + ]); + return true; + } + catch (KeeperException.BadVersionException) + { + // Retry heartbeat races while retaining the caller's row and table preconditions. + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); + return false; + } + } } /// @@ -294,10 +389,43 @@ public Task UpdateIAmAliveAsync(MembershipEntry entry, CancellationToken cancell ArgumentNullException.ThrowIfNull(entry); cancellationToken.ThrowIfCancellationRequested(); - string rowIAmAlivePath = ConvertToRowIAmAlivePath(entry.SiloAddress); - byte[] newRowIAmAliveData = Serialize(entry.IAmAliveTime); - //update the data for IAmAlive unconditionally - return UsingZookeeper(zk => zk.setDataAsync(rowIAmAlivePath, newRowIAmAliveData), this.deploymentConnectionString, this.watcher, cancellationToken); + return UsingZookeeper(zk => UpdateIAmAliveCoreAsync(entry, zk.GetData, zk.SetData, cancellationToken), + this.deploymentConnectionString, this.watcher, cancellationToken); + } + + internal static async Task UpdateIAmAliveCoreAsync( + MembershipEntry entry, + Func> getData, + Func> setData, + CancellationToken cancellationToken) + { + var path = ConvertToRowIAmAlivePath(entry.SiloAddress); + while (true) + { + cancellationToken.ThrowIfCancellationRequested(); + try + { + var heartbeat = await getData(path); + if (Deserialize(heartbeat.Data) >= entry.IAmAliveTime) + { + return true; + } + + cancellationToken.ThrowIfCancellationRequested(); + await setData(path, Serialize(entry.IAmAliveTime), heartbeat.Stat.getVersion()); + return true; + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await getData("/"); + return true; + } + catch (KeeperException.BadVersionException) + { + // Preserve the maximum when another heartbeat or status update wins the race. + } + } } /// @@ -318,59 +446,70 @@ public Task DeleteMembershipTableEntriesAsync(string clusterId, CancellationToke }, cancellationToken); } - private async Task TryTransaction(Func transactionFunc, CancellationToken cancellationToken) - { - try - { - await UsingZookeeper(zk => transactionFunc(zk.transaction()).commitAsync(), this.deploymentConnectionString, this.watcher, cancellationToken); - return true; - } - catch (KeeperException e) - { - //these exceptions are thrown when the transaction fails to commit due to semantical reasons - if (e is KeeperException.NodeExistsException || e is KeeperException.NoNodeException || - e is KeeperException.BadVersionException) - { - return false; - } - throw; - } - } - /// /// Reads the nodes /IP:Port@Gen and /IP:Port@Gen/IAmAlive (which together is one row) /// /// The zookeeper instance used for the read /// The silo address. + /// Whether a missing row is represented by an empty point read. /// A token which cancels the operation. - private static async Task> GetRow(ZooKeeper zk, SiloAddress siloAddress, CancellationToken cancellationToken) + private static async Task?> GetRow(NativeOperations zk, SiloAddress siloAddress, bool allowMissing, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); string rowPath = ConvertToRowPath(siloAddress); string rowIAmAlivePath = ConvertToRowIAmAlivePath(siloAddress); - var rowDataTask = zk.getDataAsync(rowPath); - var rowIAmAliveDataTask = zk.getDataAsync(rowIAmAlivePath); - - await Task.WhenAll(rowDataTask, rowIAmAliveDataTask); + DataResult row; + try + { + row = await zk.GetData(rowPath); + } + catch (KeeperException.NoNodeException) when (allowMissing) + { + return null; + } - MembershipEntry me = Deserialize((await rowDataTask).Data); - me.IAmAliveTime = Deserialize((await rowIAmAliveDataTask).Data); + cancellationToken.ThrowIfCancellationRequested(); + var heartbeat = await zk.GetData(rowIAmAlivePath); + MembershipEntry me = Deserialize(row.Data); + me.IAmAliveTime = Deserialize(heartbeat.Data); - int rowVersion = (await rowDataTask).Stat.getVersion(); + int rowVersion = row.Stat.getVersion(); return new Tuple(me, rowVersion.ToString(CultureInfo.InvariantCulture)); } - private static async Task UsingZookeeper(Func> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken, bool canBeReadOnly = false) + // These delegates expose the native request boundary for deterministic tests of the real operation loops. + internal sealed class NativeOperations( + Func> getData, + Func> getChildren, + Func sync, + Func, Task> multi, + Func> setData) + { + internal Func> GetData { get; } = getData; + internal Func> GetChildren { get; } = getChildren; + internal Func Sync { get; } = sync; + internal Func, Task> Multi { get; } = multi; + internal Func> SetData { get; } = setData; + } + + private static async Task UsingZookeeper(Func> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); var operation = ZooKeeper.Using(deploymentConnectionString, ZOOKEEPER_SESSION_TIMEOUT, watcher, zk => { cancellationToken.ThrowIfCancellationRequested(); - return zkMethod(zk); - }, canBeReadOnly); + return zkMethod(new NativeOperations( + path => zk.getDataAsync(path), path => zk.getChildrenAsync(path), zk.sync, + operations => zk.multiAsync(operations), zk.setDataAsync)); + }); + return await AwaitOperationAsync(operation, cancellationToken); + } + + internal static async Task AwaitOperationAsync(Task operation, CancellationToken cancellationToken) + { // ZooKeeperNetEx is tokenless. Keep the client alive until pending requests and // asynchronous disposal finish, observing failures even if the caller stops waiting. operation.Ignore(); @@ -446,23 +585,54 @@ internal static T Deserialize(byte[] data) /// public Task CleanupDefunctSiloEntriesAsync(DateTimeOffset beforeDate, CancellationToken cancellationToken = default) { - return UsingZookeeper(async zk => + return UsingZookeeper(zk => CleanupCoreAsync(zk, beforeDate, cancellationToken), + this.deploymentConnectionString, this.watcher, cancellationToken); + } + + internal static async Task CleanupCoreAsync(NativeOperations zk, DateTimeOffset beforeDate, CancellationToken cancellationToken) + { + var cutoff = beforeDate.UtcDateTime; + while (true) { - var childrenResult = await zk.getChildrenAsync("/"); - var rows = await Task.WhenAll( - childrenResult.Children.Select(child => GetRow(zk, SiloAddress.FromParsableString(child), cancellationToken))); + var table = await ReadCoreAsync(zk, null, cancellationToken); + var candidates = table.Members.Where(row => row.Item1.Status == SiloStatus.Dead + && row.Item1.StartTime < cutoff && row.Item1.IAmAliveTime < cutoff + && row.Item1.SuspectTimes?.Any(vote => vote.Item2 >= cutoff) != true).ToList(); + if (candidates.Count == 0) + { + return true; + } - foreach (var (entry, _) in rows) + foreach (var (entry, etag) in candidates) { - if (entry.Status != SiloStatus.Active - && Math.Max(entry.IAmAliveTime.Ticks, entry.StartTime.Ticks) < beforeDate.Ticks) + cancellationToken.ThrowIfCancellationRequested(); + try + { + var heartbeatPath = ConvertToRowIAmAlivePath(entry.SiloAddress); + var heartbeat = await zk.GetData(heartbeatPath); + if (Deserialize(heartbeat.Data) >= cutoff) + { + continue; + } + + cancellationToken.ThrowIfCancellationRequested(); + await zk.Multi( + [ + Op.delete(heartbeatPath, heartbeat.Stat.getVersion()), + Op.delete(ConvertToRowPath(entry.SiloAddress), int.Parse(etag, CultureInfo.InvariantCulture)) + ]); + } + catch (KeeperException.BadVersionException) + { + // Re-evaluate eligibility after a concurrent row or heartbeat update. + } + catch (KeeperException.NoNodeException) { - await DeleteRecursive(zk, ConvertToRowPath(entry.SiloAddress), cancellationToken); + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); } } - - return true; - }, this.deploymentConnectionString, this.watcher, cancellationToken); + } } [LoggerMessage( diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index 4b826de25ac..01c22398f05 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -1,10 +1,13 @@ using System; +using System.Collections.Generic; +using System.Linq; using System.Net; using System.Reflection; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.Logging.Abstractions; using Microsoft.Extensions.Options; +using org.apache.zookeeper; using Orleans.Configuration; using Orleans.Runtime; using Orleans.Runtime.Membership; @@ -247,6 +250,813 @@ public void ConvertToRowIAmAlivePath_ValidAddress_AppendsIAmAliveSegment() Assert.Equal(InvokePrivatePathMethod("ConvertToRowPath", address) + "/IAmAlive", result); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Read_ConcurrentStatusUpdate_RetriesWithCoherentVersion(bool pointRead) + { + var (fake, entry) = await CreateNativeTable(); + fake.AfterRead = async path => + { + if (path == ZooKeeperNativeFake.RowPath(entry.SiloAddress)) + { + fake.AfterRead = null; + entry.Status = SiloStatus.Dead; + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + } + }; + + var result = await Read(fake, pointRead ? entry.SiloAddress : null); + + var row = Assert.Single(result.Members); + Assert.Equal(SiloStatus.Dead, row.Item1.Status); + Assert.Equal("1", row.Item2); + Assert.Equal(2, result.Version.Version); + Assert.Equal("2", result.Version.VersionEtag); + Assert.Equal(2, fake.Calls.Count(call => call == "sync /")); + Assert.Equal("sync /", fake.Calls[0]); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Read_ConcurrentUnversionedCleanup_RetriesChildVersion(bool pointRead) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + fake.AfterRead = async path => + { + if (path == ZooKeeperNativeFake.RowPath(entry.SiloAddress)) + { + fake.AfterRead = null; + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + } + }; + + var result = await Read(fake, pointRead ? entry.SiloAddress : null); + + Assert.Empty(result.Members); + Assert.Equal(1, result.Version.Version); + Assert.Equal("1", result.Version.VersionEtag); + Assert.Single(fake.Nodes); + Assert.Equal(2, fake.Nodes["/"].ChildrenVersion); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Read_CleanupAfterHeartbeatRead_FencesCompletedRows(bool pointRead) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + fake.AfterRead = async path => + { + if (path == ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)) + { + fake.AfterRead = null; + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + } + }; + + var result = await Read(fake, pointRead ? entry.SiloAddress : null); + + Assert.Empty(result.Members); + Assert.Equal(1, result.Version.Version); + Assert.Single(fake.Nodes); + } + + [Fact] + public async Task Read_MissingPointRow_ReturnsEmptyViewWithCurrentVersion() + { + var fake = new ZooKeeperNativeFake(); + fake.Nodes["/"] = new([], 17); + + var result = await Read(fake, CreateSiloAddress()); + + Assert.Empty(result.Members); + Assert.Equal(17, result.Version.Version); + Assert.Equal("17", result.Version.VersionEtag); + Assert.Equal(["sync /", "read /", "read /127.0.0.1:11111@12345", "read /"], fake.Calls); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Read_ConcurrentHeartbeat_UsesLatestSeparateTimestamp(bool pointRead) + { + var (fake, entry) = await CreateNativeTable(); + fake.AfterRead = async path => + { + if (path == ZooKeeperNativeFake.RowPath(entry.SiloAddress)) + { + fake.AfterRead = null; + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + await Heartbeat(fake, entry); + } + }; + + var result = await Read(fake, pointRead ? entry.SiloAddress : null); + + Assert.Equal(entry.IAmAliveTime, Assert.Single(result.Members).Item1.IAmAliveTime); + Assert.Equal(1, result.Version.Version); + Assert.Equal("0", result.Members[0].Item2); + } + + [Fact] + public async Task Read_ParallelMissingRowAndAuthorizationFailure_PropagatesAuthorizationFailure() + { + var (fake, entry) = await CreateNativeTable(); + var second = CreateTimedEntry(12346); + Assert.True(await Insert(fake, second, 1)); + var failure = new KeeperException.NoAuthException(); + fake.BeforeRead = path => path == ZooKeeperNativeFake.RowPath(entry.SiloAddress) + ? Task.FromException(new KeeperException.NoNodeException(path)) + : path == ZooKeeperNativeFake.RowPath(second.SiloAddress) + ? Task.FromException(failure) : Task.CompletedTask; + + var actual = await Record.ExceptionAsync(() => Read(fake)); + + Assert.Same(failure, actual); + Assert.Equal(1, fake.Calls.Count(call => call == "sync /")); + } + + [Fact] + public async Task Read_StableMissingHeartbeat_PropagatesFailure() + { + var (fake, entry) = await CreateNativeTable(); + var path = ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress); + fake.Nodes.Remove(path); + + var failure = await Assert.ThrowsAsync(() => Read(fake)); + + Assert.Equal(path, failure.getPath()); + Assert.Equal(1, fake.Calls.Count(call => call == "sync /")); + } + + [Fact] + public async Task Insert_PersistentNodesAndVersion_CommitAtomically() + { + var fake = new ZooKeeperNativeFake(); + var entry = CreateTimedEntry(); + + Assert.True(await Insert(fake, entry, 0)); + + Assert.Collection(Assert.Single(fake.Transactions), + operation => AssertSet(operation, "/", 0), + operation => AssertCreate(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress)), + operation => AssertCreate(operation, ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress))); + var result = await Read(fake); + Assert.Equal(1, result.Version.Version); + Assert.Equal("0", Assert.Single(result.Members).Item2); + Assert.Equal(entry.SiloName, result.Members[0].Item1.SiloName); + Assert.Equal(entry.IAmAliveTime, result.Members[0].Item1.IAmAliveTime); + Assert.All(fake.Nodes.Values, node => Assert.Equal(0, node.Flags)); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task Insert_DuplicateOrStaleVersion_RollsBackTransaction(bool duplicate) + { + var (fake, entry) = await CreateNativeTable(); + var original = fake.Nodes.ToArray(); + + Assert.False(await Insert(fake, duplicate ? entry : CreateTimedEntry(12346), duplicate ? 1 : 0)); + + Assert.Equal(original, fake.Nodes.ToArray()); + Assert.Single(fake.Transactions); + } + + [Theory] + [InlineData("0", 1)] + [InlineData("1", 1)] + [InlineData("0", 0)] + public async Task UpdateRow_OptimisticPreconditions_ControlAtomicMutation(string etag, int tableVersion) + { + var (fake, entry) = await CreateNativeTable(); + entry.Status = SiloStatus.Dead; + + var updated = await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, etag, new TableVersion(tableVersion + 1, tableVersion.ToString()), + TestContext.Current.CancellationToken); + + var expected = etag == "0" && tableVersion == 1; + Assert.Equal(expected, updated); + var result = await Read(fake); + Assert.Equal(expected ? SiloStatus.Dead : SiloStatus.Active, Assert.Single(result.Members).Item1.Status); + Assert.Equal(expected ? "1" : "0", result.Members[0].Item2); + Assert.Equal(expected ? 2 : 1, result.Version.Version); + if (expected) + { + Assert.Collection(Assert.Single(fake.Transactions), + operation => AssertSet(operation, "/", 1), + operation => AssertSet(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), 0), + operation => AssertSet(operation, ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0)); + } + else + { + Assert.Empty(fake.Transactions); + } + } + + [Fact] + public async Task UpdateRow_ConcurrentHeartbeat_RetriesAndPreservesMaximum() + { + var (fake, entry) = await CreateNativeTable(); + var heartbeat = CreateTimedEntry(); + heartbeat.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + fake.BeforeMulti = async _ => + { + fake.BeforeMulti = null; + await Heartbeat(fake, heartbeat); + }; + entry.Status = SiloStatus.Dead; + + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + + var result = await Read(fake); + var row = Assert.Single(result.Members); + Assert.Equal(SiloStatus.Dead, row.Item1.Status); + Assert.Equal(heartbeat.IAmAliveTime, row.Item1.IAmAliveTime); + Assert.Equal("1", row.Item2); + Assert.Equal(2, result.Version.Version); + Assert.Equal(2, fake.Transactions.Count); + AssertSet(fake.Transactions[0][2], ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0); + AssertSet(fake.Transactions[1][2], ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 1); + } + + [Fact] + public async Task UpdateRow_ProposedHeartbeat_AdvancesTimestamp() + { + var (fake, entry) = await CreateNativeTable(); + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + entry.Status = SiloStatus.Dead; + + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + + var result = await Read(fake); + var row = Assert.Single(result.Members); + Assert.Equal(entry.IAmAliveTime, row.Item1.IAmAliveTime); + Assert.Equal(SiloStatus.Dead, row.Item1.Status); + Assert.Equal("1", row.Item2); + Assert.Equal(2, result.Version.Version); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task UpdateRow_ConcurrentMembershipMutation_ReturnsConflictWithoutPartialWrite(bool sameRow) + { + var (fake, entry) = await CreateNativeTable(); + fake.BeforeMulti = async _ => + { + fake.BeforeMulti = null; + if (sameRow) + { + var winner = CreateTimedEntry(); + winner.Status = SiloStatus.ShuttingDown; + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, winner, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + } + else + { + Assert.True(await Insert(fake, CreateTimedEntry(12346), 1)); + } + }; + entry.Status = SiloStatus.Dead; + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(3); + + Assert.False(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + + var result = await Read(fake, entry.SiloAddress); + Assert.Equal(sameRow ? SiloStatus.ShuttingDown : SiloStatus.Active, Assert.Single(result.Members).Item1.Status); + Assert.Equal(DateTime.UnixEpoch.AddDays(1), result.Members[0].Item1.IAmAliveTime); + Assert.Equal(2, result.Version.Version); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task UpdateIAmAlive_DelayedHeartbeat_PreservesMaximumAndVersions(bool concurrentWrite) + { + var (fake, entry) = await CreateNativeTable(); + var older = CreateTimedEntry(); + older.IAmAliveTime = entry.IAmAliveTime.AddHours(1); + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + if (concurrentWrite) + { + fake.BeforeWrite = async _ => + { + fake.BeforeWrite = null; + await Heartbeat(fake, entry); + }; + } + else + { + await Heartbeat(fake, entry); + } + + await Heartbeat(fake, older); + + var result = await Read(fake); + Assert.Equal(entry.IAmAliveTime, Assert.Single(result.Members).Item1.IAmAliveTime); + Assert.Equal("0", result.Members[0].Item2); + Assert.Equal(1, result.Version.Version); + Assert.Equal(1, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Version); + Assert.Empty(fake.Transactions); + } + + [Fact] + public async Task UpdateIAmAlive_ConcurrentStatusUpdate_PreservesNewerTimestamp() + { + var (fake, entry) = await CreateNativeTable(); + var heartbeat = CreateTimedEntry(); + heartbeat.IAmAliveTime = entry.IAmAliveTime.AddHours(1); + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + entry.Status = SiloStatus.ShuttingDown; + fake.BeforeWrite = async _ => + { + fake.BeforeWrite = null; + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + }; + + await Heartbeat(fake, heartbeat); + + var result = await Read(fake); + var row = Assert.Single(result.Members); + Assert.Equal(entry.IAmAliveTime, row.Item1.IAmAliveTime); + Assert.Equal(SiloStatus.ShuttingDown, row.Item1.Status); + Assert.Equal("1", row.Item2); + Assert.Equal(2, result.Version.Version); + Assert.Equal(1, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Version); + } + + [Theory] + [InlineData("read")] + [InlineData("write")] + [InlineData("retry")] + public async Task UpdateIAmAlive_RetiredRow_CompletesWithoutRecreationOrVersionWrite(string phase) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + var cutoff = entry.IAmAliveTime.AddDays(2); + async Task Prune() + { + await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); + fake.Transactions.Clear(); + } + + if (phase == "read") + { + await Prune(); + } + else + { + fake.BeforeWrite = async _ => + { + fake.BeforeWrite = null; + await Prune(); + if (phase == "retry") + { + throw new KeeperException.BadVersionException(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)); + } + }; + } + + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + await Heartbeat(fake, entry); + + Assert.Equal("/", Assert.Single(fake.Nodes).Key); + Assert.Equal(1, fake.Nodes["/"].Version); + Assert.Empty(fake.Transactions); + } + + [Theory] + [InlineData(SiloStatus.None)] + [InlineData(SiloStatus.Created)] + [InlineData(SiloStatus.Joining)] + [InlineData(SiloStatus.Active)] + [InlineData(SiloStatus.ShuttingDown)] + [InlineData(SiloStatus.Stopping)] + [InlineData(SiloStatus.Dead)] + public async Task Cleanup_OnlyEligibleDeadRows_DeletesWithNativeVersionChecks(SiloStatus status) + { + var (fake, entry) = await CreateNativeTable(status); + var root = fake.Nodes["/"]; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + + Assert.Equal(root.Version, fake.Nodes["/"].Version); + if (status == SiloStatus.Dead) + { + Assert.Equal("/", Assert.Single(fake.Nodes).Key); + Assert.Collection(Assert.Single(fake.Transactions), + operation => AssertDelete(operation, ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0), + operation => AssertDelete(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), 0)); + } + else + { + Assert.Equal(3, fake.Nodes.Count); + Assert.Empty(fake.Transactions); + Assert.Equal(status, Assert.Single((await Read(fake)).Members).Item1.Status); + } + } + + [Theory] + [InlineData("start", -1)] + [InlineData("start", 0)] + [InlineData("start", 1)] + [InlineData("heartbeat", -1)] + [InlineData("heartbeat", 0)] + [InlineData("heartbeat", 1)] + [InlineData("vote", -1)] + [InlineData("vote", 0)] + [InlineData("vote", 1)] + public async Task Cleanup_UsesLatestEvidenceAndExactUtcCutoff(string evidence, int ticks) + { + var fake = new ZooKeeperNativeFake(); + var entry = CreateTimedEntry(); + entry.Status = SiloStatus.Dead; + var cutoff = new DateTimeOffset(1970, 1, 3, 7, 0, 0, TimeSpan.FromHours(7)); + var time = cutoff.UtcDateTime.AddTicks(ticks); + if (evidence == "start") + { + entry.StartTime = time; + } + else if (evidence == "heartbeat") + { + entry.IAmAliveTime = time; + } + else + { + entry.SuspectTimes = + [ + Tuple.Create(CreateSiloAddress(), time), + Tuple.Create(CreateTimedEntry(12346).SiloAddress, DateTime.UnixEpoch) + ]; + } + + Assert.True(await Insert(fake, entry, 0)); + fake.Transactions.Clear(); + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); + + Assert.Equal(ticks >= 0, fake.Nodes.ContainsKey(ZooKeeperNativeFake.RowPath(entry.SiloAddress))); + Assert.Equal(ticks >= 0, fake.Nodes.ContainsKey(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress))); + Assert.Equal(1, fake.Nodes["/"].Version); + Assert.Equal(ticks < 0 ? 1 : 0, fake.Transactions.Count); + } + + [Theory] + [InlineData("heartbeat")] + [InlineData("start")] + [InlineData("vote")] + [InlineData("status")] + public async Task Cleanup_ConcurrentEvidenceUpdate_PreservesBothNodes(string evidence) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + var cutoff = DateTime.UnixEpoch.AddDays(2); + fake.BeforeMulti = async _ => + { + fake.BeforeMulti = null; + if (evidence == "heartbeat") + { + entry.IAmAliveTime = cutoff; + await Heartbeat(fake, entry); + } + else + { + if (evidence == "start") + { + entry.StartTime = cutoff; + } + else if (evidence == "vote") + { + entry.SuspectTimes = [Tuple.Create(CreateSiloAddress(), cutoff)]; + } + else + { + entry.Status = SiloStatus.Active; + } + + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + } + }; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); + + Assert.Equal(3, fake.Nodes.Count); + Assert.Equal(evidence == "heartbeat" ? 1 : 2, fake.Nodes["/"].Version); + var row = Assert.Single((await Read(fake)).Members).Item1; + Assert.Equal(entry.Status, row.Status); + Assert.Equal(entry.StartTime, row.StartTime); + Assert.Equal(entry.IAmAliveTime, row.IAmAliveTime); + Assert.Equal(entry.SuspectTimes, row.SuspectTimes); + Assert.Collection(fake.Transactions[0], + operation => AssertDelete(operation, ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0), + operation => AssertDelete(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), 0)); + } + + [Fact] + public async Task Cleanup_ConcurrentCleaners_CompleteWithUnchangedTableVersion() + { + var (fake, _) = await CreateNativeTable(SiloStatus.Dead); + var cutoff = DateTime.UnixEpoch.AddDays(2); + fake.BeforeMulti = async _ => + { + fake.BeforeMulti = null; + await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); + }; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); + + Assert.Single(fake.Nodes); + Assert.Equal(1, fake.Nodes["/"].Version); + Assert.Equal(2, fake.Transactions.Count); + } + + [Theory] + [InlineData("read")] + [InlineData("point")] + [InlineData("insert")] + [InlineData("update")] + [InlineData("heartbeat")] + [InlineData("cleanup")] + public async Task MembershipOperations_MissingClusterHistory_PropagatesNoNode(string operation) + { + var fake = new ZooKeeperNativeFake(); + fake.Nodes.Clear(); + + var failure = await Assert.ThrowsAsync( + () => InvokeNative(operation, fake, CreateTimedEntry(), TestContext.Current.CancellationToken)); + + Assert.Equal("/", failure.getPath()); + Assert.Empty(fake.Nodes); + } + + [Theory] + [InlineData("read", "auth", false)] + [InlineData("point", "connection", false)] + [InlineData("insert", "session", true)] + [InlineData("update", "auth", false)] + [InlineData("update", "auth", true)] + [InlineData("update", "connection", true)] + [InlineData("update", "session", true)] + [InlineData("heartbeat", "connection", false)] + [InlineData("heartbeat", "auth", true)] + [InlineData("heartbeat", "connection", true)] + [InlineData("heartbeat", "session", true)] + [InlineData("cleanup", "session", false)] + [InlineData("cleanup", "auth", true)] + [InlineData("cleanup", "connection", true)] + [InlineData("cleanup", "session", true)] + public async Task MembershipOperations_InfrastructureFailure_PropagatesSameException(string operation, string kind, bool duringWrite) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + Exception failure = kind switch + { + "auth" => new KeeperException.NoAuthException(), + "session" => new KeeperException.SessionExpiredException(), + _ => new KeeperException.ConnectionLossException() + }; + if (!duringWrite) + { + fake.BeforeRead = _ => Task.FromException(failure); + } + + fake.BeforeMulti = _ => Task.FromException(failure); + fake.BeforeWrite = _ => Task.FromException(failure); + entry.IAmAliveTime = entry.IAmAliveTime.AddHours(1); + + var actual = await Record.ExceptionAsync(() => InvokeNative(operation, fake, entry, TestContext.Current.CancellationToken)); + + Assert.Same(failure, actual); + Assert.Equal(3, fake.Nodes.Count); + Assert.Equal(1, fake.Nodes["/"].Version); + } + + [Theory] + [InlineData("update")] + [InlineData("heartbeat")] + [InlineData("cleanup")] + public async Task MembershipOperations_ClusterRemovedDuringWrite_PropagatesNoNode(string operation) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + Task RemoveCluster() + { + fake.Nodes.Clear(); + return Task.CompletedTask; + } + + fake.BeforeMulti = _ => RemoveCluster(); + fake.BeforeWrite = _ => RemoveCluster(); + entry.IAmAliveTime = entry.IAmAliveTime.AddHours(1); + + var failure = await Assert.ThrowsAsync( + () => InvokeNative(operation, fake, entry, TestContext.Current.CancellationToken)); + + Assert.Equal("/", failure.getPath()); + Assert.Empty(fake.Nodes); + } + + [Theory] + [InlineData("update")] + [InlineData("heartbeat")] + [InlineData("cleanup")] + public async Task NativeOperations_CanceledConflict_StopsRetrying(string operation) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var requestCount = 0; + Task Conflict() + { + requestCount = fake.Calls.Count; + cancellation.Cancel(); + throw new KeeperException.BadVersionException(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)); + } + + fake.BeforeMulti = _ => Conflict(); + fake.BeforeWrite = _ => Conflict(); + entry.IAmAliveTime = entry.IAmAliveTime.AddHours(1); + + var failure = await Assert.ThrowsAnyAsync( + () => InvokeNative(operation, fake, entry, cancellation.Token)); + + Assert.Equal(cancellation.Token, failure.CancellationToken); + Assert.Equal(requestCount, fake.Calls.Count); + Assert.Equal(1, fake.Nodes["/"].Version); + Assert.Equal(3, fake.Nodes.Count); + } + + [Theory] + [InlineData("read")] + [InlineData("point")] + [InlineData("insert")] + [InlineData("update")] + [InlineData("heartbeat")] + [InlineData("cleanup")] + public async Task NativeOperations_PreCanceledToken_IssuesNoRequests(string operation) + { + var fake = new ZooKeeperNativeFake(); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + cancellation.Cancel(); + + var failure = await Assert.ThrowsAnyAsync( + () => InvokeNative(operation, fake, CreateTimedEntry(), cancellation.Token)); + + Assert.Equal(cancellation.Token, failure.CancellationToken); + Assert.Empty(fake.Calls); + } + + [Theory] + [InlineData("read")] + [InlineData("point")] + [InlineData("update")] + [InlineData("heartbeat")] + [InlineData("cleanup")] + public async Task NativeOperations_CanceledAfterRead_StartsNoFurtherRequests(string operation) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + fake.AfterRead = _ => + { + cancellation.Cancel(); + return Task.CompletedTask; + }; + entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); + + var failure = await Assert.ThrowsAnyAsync( + () => InvokeNative(operation, fake, entry, cancellation.Token)); + + Assert.Equal(cancellation.Token, failure.CancellationToken); + Assert.Equal(1, fake.Calls.Count(call => call.StartsWith("read ", StringComparison.Ordinal))); + Assert.Empty(fake.Transactions); + Assert.DoesNotContain(fake.Calls, call => call.StartsWith("write ", StringComparison.Ordinal)); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task AwaitOperation_CanceledCaller_PreservesRequestAndDisposalOwnership(bool fail) + { + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var request = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var disposing = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var disposal = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var failure = new KeeperException.ConnectionLossException(); + async Task NativeOperation() + { + try + { + await request.Task; + if (fail) + { + throw failure; + } + + return 42; + } + finally + { + disposing.SetResult(); + await disposal.Task; + } + } + + var operation = NativeOperation(); + var caller = ZooKeeperBasedMembershipTable.AwaitOperationAsync(operation, cancellation.Token); + cancellation.Cancel(); + var canceled = await Assert.ThrowsAnyAsync(() => caller); + Assert.Equal(cancellation.Token, canceled.CancellationToken); + Assert.False(operation.IsCompleted); + Assert.False(disposing.Task.IsCompleted); + + request.SetResult(); + await disposing.Task.WaitAsync(TestContext.Current.CancellationToken); + Assert.False(operation.IsCompleted); + disposal.SetResult(); + if (fail) + { + Assert.Same(failure, await Record.ExceptionAsync(() => operation)); + } + else + { + Assert.Equal(42, await operation); + } + } + + private static Task InvokeNative(string operation, ZooKeeperNativeFake fake, MembershipEntry entry, CancellationToken cancellationToken) => + operation switch + { + "read" => ZooKeeperBasedMembershipTable.ReadCoreAsync(fake.Operations, null, cancellationToken), + "point" => ZooKeeperBasedMembershipTable.ReadCoreAsync(fake.Operations, entry.SiloAddress, cancellationToken), + "insert" => ZooKeeperBasedMembershipTable.InsertRowCoreAsync(fake.Operations, entry, new TableVersion(2, "1"), cancellationToken), + "update" => ZooKeeperBasedMembershipTable.UpdateRowCoreAsync(fake.Operations, entry, "0", new TableVersion(2, "1"), cancellationToken), + "heartbeat" => ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.GetData, fake.SetData, cancellationToken), + "cleanup" => ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, DateTime.UnixEpoch.AddDays(3), cancellationToken), + _ => throw new ArgumentOutOfRangeException(nameof(operation)) + }; + + private static async Task<(ZooKeeperNativeFake Fake, MembershipEntry Entry)> CreateNativeTable(SiloStatus status = SiloStatus.Active) + { + var fake = new ZooKeeperNativeFake(); + var entry = CreateTimedEntry(); + entry.Status = status; + Assert.True(await Insert(fake, entry, 0)); + fake.Calls.Clear(); + fake.Transactions.Clear(); + return (fake, entry); + } + + private static MembershipEntry CreateTimedEntry(int generation = 12345) => + new() + { + SiloAddress = SiloAddress.New(new IPEndPoint(IPAddress.Loopback, 11111), generation), + HostName = "host-a", + SiloName = "silo-a", + Status = SiloStatus.Active, + StartTime = DateTime.UnixEpoch, + IAmAliveTime = DateTime.UnixEpoch.AddDays(1) + }; + + private static Task Insert(ZooKeeperNativeFake fake, MembershipEntry entry, int version) => + ZooKeeperBasedMembershipTable.InsertRowCoreAsync(fake.Operations, entry, + new TableVersion(version + 1, version.ToString()), TestContext.Current.CancellationToken); + + private static Task Read(ZooKeeperNativeFake fake, SiloAddress? address = null) => + ZooKeeperBasedMembershipTable.ReadCoreAsync(fake.Operations, address, TestContext.Current.CancellationToken); + + private static Task Heartbeat(ZooKeeperNativeFake fake, MembershipEntry entry) => + ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.GetData, fake.SetData, TestContext.Current.CancellationToken); + + private static void AssertSet(object operation, string path, int version) + { + var set = Assert.IsType(operation); + Assert.Equal(path, set.Path); + Assert.Equal(version, set.Version); + } + + private static void AssertCreate(object operation, string path) + { + var create = Assert.IsType(operation); + Assert.Equal(path, create.Path); + Assert.Equal(0, create.Flags); + Assert.Equal(ZooDefs.Ids.OPEN_ACL_UNSAFE, create.Acl); + } + + private static void AssertDelete(object operation, string path, int version) + { + var delete = Assert.IsType(operation); + Assert.Equal(path, delete.Path); + Assert.Equal(version, delete.Version); + } + private static ZooKeeperBasedMembershipTable CreateSut( string connectionString = "sentinel.invalid:2181", string clusterId = "cluster-a") => diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs new file mode 100644 index 00000000000..8d2d8854070 --- /dev/null +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs @@ -0,0 +1,186 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Reflection; +using System.Threading.Tasks; +using org.apache.zookeeper; +using org.apache.zookeeper.data; +using Orleans.Runtime.Membership; +using Xunit; + +namespace UnitTests.MembershipTests; + +internal sealed class ZooKeeperNativeFake +{ + internal sealed record Node(byte[] Data, int Version = 0, int ChildrenVersion = 0, int Flags = 0); + internal abstract record Request(string Path); + internal sealed record CreateRequest(string Path, byte[] Data, int Flags, List Acl) : Request(Path); + internal sealed record SetDataRequest(string Path, byte[]? Data, int Version) : Request(Path); + internal sealed record DeleteRequest(string Path, int Version) : Request(Path); + + internal Dictionary Nodes { get; private set; } = new(StringComparer.Ordinal) + { + ["/"] = new([]) + }; + + internal List Calls { get; } = []; + internal List> Transactions { get; } = []; + internal Func? BeforeRead { get; set; } + internal Func? AfterRead { get; set; } + internal Func, Task>? BeforeMulti { get; set; } + internal Func? BeforeWrite { get; set; } + internal Func? OnSync { get; set; } + + internal ZooKeeperBasedMembershipTable.NativeOperations Operations => new(GetData, GetChildren, Sync, Multi, SetData); + + internal async Task GetData(string path) + { + Calls.Add("read " + path); + if (BeforeRead is { } before) + { + await before(path); + } + + var node = GetNode(Nodes, path); + var result = CreateResult(node.Data, CreateStat(node)); + if (AfterRead is { } after) + { + await after(path); + } + + return result; + } + + private Task GetChildren(string path) + { + Calls.Add("children " + path); + var node = GetNode(Nodes, path); + var children = Nodes.Keys.Where(key => key != path && Parent(key) == path) + .Select(key => key[(key.LastIndexOf('/') + 1)..]).ToList(); + return Task.FromResult(CreateResult(children, CreateStat(node))); + } + + private async Task Sync(string path) + { + Calls.Add("sync " + path); + if (OnSync is { } sync) + { + await sync(); + } + } + + internal async Task Multi(List operations) + { + Calls.Add("multi"); + var requests = operations.Select(ToRequest).ToList(); + Transactions.Add(requests); + if (BeforeMulti is { } before) + { + await before(requests); + } + + // Apply against a copy so failed native preconditions roll back every earlier operation. + var next = new Dictionary(Nodes, StringComparer.Ordinal); + foreach (var request in requests) + { + switch (request) + { + case CreateRequest create: + var parent = Parent(create.Path); + var parentNode = GetNode(next, parent); + if (next.ContainsKey(create.Path)) + { + throw new KeeperException.NodeExistsException(create.Path); + } + + next.Add(create.Path, new(create.Data, Flags: create.Flags)); + next[parent] = parentNode with { ChildrenVersion = parentNode.ChildrenVersion + 1 }; + break; + case SetDataRequest set: + Set(next, set.Path, set.Data, set.Version); + break; + case DeleteRequest delete: + CheckVersion(next, delete.Path, delete.Version); + if (next.Keys.Any(key => key != delete.Path && Parent(key) == delete.Path)) + { + throw new KeeperException.NotEmptyException(delete.Path); + } + + next.Remove(delete.Path); + var parentPath = Parent(delete.Path); + var owner = GetNode(next, parentPath); + next[parentPath] = owner with { ChildrenVersion = owner.ChildrenVersion + 1 }; + break; + default: + throw new InvalidOperationException($"Unexpected native request: {request.GetType().Name}"); + } + } + + Nodes = next; + } + + internal async Task SetData(string path, byte[] data, int version) + { + Calls.Add("write " + path); + if (BeforeWrite is { } before) + { + await before(path); + } + + Set(Nodes, path, data, version); + return CreateStat(Nodes[path]); + } + + internal static string RowPath(Orleans.Runtime.SiloAddress address) => "/" + address.ToParsableString(); + internal static string HeartbeatPath(Orleans.Runtime.SiloAddress address) => RowPath(address) + "/IAmAlive"; + + private static Request ToRequest(Op operation) + { + var method = typeof(Op).GetMethod("toRequestRecord", BindingFlags.Instance | BindingFlags.NonPublic); + Assert.NotNull(method); + var result = method.Invoke(operation, null); + Assert.NotNull(result); + var path = GetValue(result, "getPath"); + return result.GetType().Name switch + { + "CreateRequest" => new CreateRequest(path, GetValue(result, "getData"), + GetValue(result, "getFlags"), GetValue>(result, "getAcl")), + "SetDataRequest" => new SetDataRequest(path, GetValue(result, "getData"), GetValue(result, "getVersion")), + "DeleteRequest" => new DeleteRequest(path, GetValue(result, "getVersion")), + _ => throw new InvalidOperationException($"Unexpected native request: {result.GetType().Name}") + }; + } + + private static T GetValue(object request, string methodName) + { + var method = request.GetType().GetMethod(methodName); + Assert.NotNull(method); + return (T)method.Invoke(request, null)!; + } + + private static T CreateResult(params object[] arguments) => + Assert.IsType(Activator.CreateInstance(typeof(T), BindingFlags.Instance | BindingFlags.NonPublic, + binder: null, args: arguments, culture: null)); + + private static Stat CreateStat(Node node) => new(0, 0, 0, 0, node.Version, node.ChildrenVersion, 0, 0, 0, 0, 0); + private static string Parent(string path) => path.LastIndexOf('/') is var index && index > 0 ? path[..index] : "/"; + private static Node GetNode(Dictionary nodes, string path) => + nodes.TryGetValue(path, out var node) ? node : throw new KeeperException.NoNodeException(path); + + private static Node CheckVersion(Dictionary nodes, string path, int version) + { + var node = GetNode(nodes, path); + if (version != -1 && node.Version != version) + { + throw new KeeperException.BadVersionException(path); + } + + return node; + } + + private static void Set(Dictionary nodes, string path, byte[]? data, int version) + { + var node = CheckVersion(nodes, path, version); + nodes[path] = node with { Data = data ?? [], Version = node.Version + 1 }; + } +} From ae69b1aab988a86ab63fb205ab3caaae2a738d96 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 07:11:38 -0700 Subject: [PATCH 3/8] fix(zookeeper): preserve read-only connections and clarify test diagnostics --- .../ZooKeeperBasedMembershipTable.cs | 8 ++--- .../ZooKeeperBasedMembershipTableUnitTests.cs | 14 ++++++++ .../ZooKeeperNativeFake.cs | 34 +++++++++++++------ 3 files changed, 42 insertions(+), 14 deletions(-) diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index 577d6302d81..08d2489c5cf 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -132,7 +132,7 @@ await UsingZookeeper(rootConnectionString, async zk => public Task ReadRowAsync(SiloAddress siloAddress, CancellationToken cancellationToken = default) { return UsingZookeeper(zk => ReadCoreAsync(zk, siloAddress, cancellationToken), - this.deploymentConnectionString, this.watcher, cancellationToken); + this.deploymentConnectionString, this.watcher, cancellationToken, canBeReadOnly: true); } /// @@ -154,7 +154,7 @@ public Task ReadAllAsync(CancellationToken cancellationToke internal static Task ReadAllAsync(string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken) { return UsingZookeeper(zk => ReadCoreAsync(zk, null, cancellationToken), - deploymentConnectionString, watcher, cancellationToken); + deploymentConnectionString, watcher, cancellationToken, canBeReadOnly: true); } internal static async Task ReadCoreAsync( @@ -494,7 +494,7 @@ internal sealed class NativeOperations( internal Func> SetData { get; } = setData; } - private static async Task UsingZookeeper(Func> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken) + private static async Task UsingZookeeper(Func> zkMethod, string deploymentConnectionString, ZooKeeperWatcher watcher, CancellationToken cancellationToken, bool canBeReadOnly = false) { cancellationToken.ThrowIfCancellationRequested(); var operation = ZooKeeper.Using(deploymentConnectionString, ZOOKEEPER_SESSION_TIMEOUT, watcher, zk => @@ -503,7 +503,7 @@ private static async Task UsingZookeeper(Func> z return zkMethod(new NativeOperations( path => zk.getDataAsync(path), path => zk.getChildrenAsync(path), zk.sync, operations => zk.multiAsync(operations), zk.setDataAsync)); - }); + }, canBeReadOnly); return await AwaitOperationAsync(operation, cancellationToken); } diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index 01c22398f05..f963eeef384 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -991,6 +991,20 @@ async Task NativeOperation() } } + [Theory] + [InlineData("MissingMethod", "expected a parameterless instance method")] + [InlineData("ToString", "expected System.Int32, received System.String")] + public void NativeFake_IncompatibleSdkMember_ReportsActionableDiagnostic(string methodName, string detail) + { + var failure = Assert.Throws(() => + ZooKeeperNativeFake.GetValue(new object(), methodName)); + + Assert.Contains($"System.Object.{methodName}", failure.Message, StringComparison.Ordinal); + Assert.Contains(detail, failure.Message, StringComparison.Ordinal); + Assert.Contains("Update the fake for the installed ZooKeeperNetEx API", failure.Message, StringComparison.Ordinal); + Assert.Contains(typeof(Op).Assembly.FullName!, failure.Message, StringComparison.Ordinal); + } + private static Task InvokeNative(string operation, ZooKeeperNativeFake fake, MembershipEntry entry, CancellationToken cancellationToken) => operation switch { diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs index 8d2d8854070..546b1318020 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs @@ -136,28 +136,42 @@ internal async Task SetData(string path, byte[] data, int version) private static Request ToRequest(Op operation) { - var method = typeof(Op).GetMethod("toRequestRecord", BindingFlags.Instance | BindingFlags.NonPublic); - Assert.NotNull(method); - var result = method.Invoke(operation, null); - Assert.NotNull(result); + // ZooKeeperNetEx exposes only an Op's type and path publicly; decoding its payload requires this SDK seam. + var result = GetValue(operation, "toRequestRecord", visibility: BindingFlags.NonPublic); var path = GetValue(result, "getPath"); return result.GetType().Name switch { "CreateRequest" => new CreateRequest(path, GetValue(result, "getData"), GetValue(result, "getFlags"), GetValue>(result, "getAcl")), - "SetDataRequest" => new SetDataRequest(path, GetValue(result, "getData"), GetValue(result, "getVersion")), + "SetDataRequest" => new SetDataRequest(path, GetValue(result, "getData", allowNull: true), GetValue(result, "getVersion")), "DeleteRequest" => new DeleteRequest(path, GetValue(result, "getVersion")), - _ => throw new InvalidOperationException($"Unexpected native request: {result.GetType().Name}") + _ => throw ApiMismatch(result.GetType(), "toRequestRecord", $"unsupported request type {result.GetType().FullName}") }; } - private static T GetValue(object request, string methodName) + internal static T GetValue(object request, string methodName, bool allowNull = false, BindingFlags visibility = BindingFlags.Public) { - var method = request.GetType().GetMethod(methodName); - Assert.NotNull(method); - return (T)method.Invoke(request, null)!; + var type = request.GetType(); + var method = type.GetMethod(methodName, BindingFlags.Instance | visibility, Type.EmptyTypes) + ?? throw ApiMismatch(type, methodName, "expected a parameterless instance method"); + var result = method.Invoke(request, null); + if (result is T value) + { + return value; + } + + if (result is null && allowNull && default(T) is null) + { + return default!; + } + + throw ApiMismatch(type, methodName, $"expected {typeof(T).FullName}, received {result?.GetType().FullName ?? "null"}"); } + private static InvalidOperationException ApiMismatch(Type type, string methodName, string detail) => + new($"ZooKeeperNativeFake cannot decode {type.FullName}.{methodName}: {detail}. " + + $"Update the fake for the installed ZooKeeperNetEx API ({typeof(Op).Assembly.FullName})."); + private static T CreateResult(params object[] arguments) => Assert.IsType(Activator.CreateInstance(typeof(T), BindingFlags.Instance | BindingFlags.NonPublic, binder: null, args: arguments, culture: null)); From 5747ebe320fd5660daec0c04ea57f4743b747879 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 08:51:00 -0700 Subject: [PATCH 4/8] fix(zookeeper): use a blind owner heartbeat write --- .../ZooKeeperBasedMembershipTable.cs | 48 +---- .../ZooKeeperBasedMembershipTableUnitTests.cs | 188 ++++++++++-------- .../ZooKeeperNativeFake.cs | 2 + 3 files changed, 119 insertions(+), 119 deletions(-) diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index 08d2489c5cf..aa0673e90f0 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -29,7 +29,7 @@ namespace Orleans.Runtime.Membership /// Every Orleans deployment has a node /UniqueDeploymentId /// Every Silo's state is saved in /UniqueDeploymentId/IP:Port@Gen /// Every Silo's IAmAlive is saved in /UniqueDeploymentId/IP:Port@Gen/IAmAlive - /// IAmAlive is saved in a separate node so compare-and-set heartbeat updates preserve the membership row's version. + /// IAmAlive is saved in a separate node so owner heartbeat writes preserve the membership row's version. /// /// a node's ZK version is its ETag: /// the table version is the version of /UniqueDeploymentId @@ -368,16 +368,10 @@ await zk.Multi( } /// - /// Updates the IAmAlive part (column) of the MembershipEntry for this silo. - /// This operation should only update the IAmAlive column and not change other columns. - /// This operation is a "dirty write" or "in place update" and is performed without etag validation. - /// With regards to eTags update: - /// This operation may automatically update the eTag associated with the given silo row, but it does not have to. It can also leave the etag not changed ("dirty write"). - /// With regards to TableVersion: - /// this operation should not change the TableVersion of the table. It should leave it untouched. - /// There is no scenario where this operation could fail due to table semantical reasons. It can only fail due to network problems or table unavailability. + /// Writes the owning silo's IAmAlive timestamp to its heartbeat node using one unconditional native update. + /// The membership row and table version are preserved. Native storage failures propagate to the caller. /// - /// The target MembershipEntry tp update + /// The owning silo's membership entry containing its heartbeat timestamp. /// Task representing the successful execution of this operation. /// is . [Obsolete("Use UpdateIAmAliveAsync instead.")] @@ -389,43 +383,17 @@ public Task UpdateIAmAliveAsync(MembershipEntry entry, CancellationToken cancell ArgumentNullException.ThrowIfNull(entry); cancellationToken.ThrowIfCancellationRequested(); - return UsingZookeeper(zk => UpdateIAmAliveCoreAsync(entry, zk.GetData, zk.SetData, cancellationToken), + return UsingZookeeper(zk => UpdateIAmAliveCoreAsync(entry, zk.SetData, cancellationToken), this.deploymentConnectionString, this.watcher, cancellationToken); } - internal static async Task UpdateIAmAliveCoreAsync( + internal static Task UpdateIAmAliveCoreAsync( MembershipEntry entry, - Func> getData, Func> setData, CancellationToken cancellationToken) { - var path = ConvertToRowIAmAlivePath(entry.SiloAddress); - while (true) - { - cancellationToken.ThrowIfCancellationRequested(); - try - { - var heartbeat = await getData(path); - if (Deserialize(heartbeat.Data) >= entry.IAmAliveTime) - { - return true; - } - - cancellationToken.ThrowIfCancellationRequested(); - await setData(path, Serialize(entry.IAmAliveTime), heartbeat.Stat.getVersion()); - return true; - } - catch (KeeperException.NoNodeException) - { - cancellationToken.ThrowIfCancellationRequested(); - await getData("/"); - return true; - } - catch (KeeperException.BadVersionException) - { - // Preserve the maximum when another heartbeat or status update wins the race. - } - } + cancellationToken.ThrowIfCancellationRequested(); + return setData(ConvertToRowIAmAlivePath(entry.SiloAddress), Serialize(entry.IAmAliveTime), -1); } /// diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index f963eeef384..f8974ea657d 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -539,100 +539,127 @@ public async Task UpdateRow_ConcurrentMembershipMutation_ReturnsConflictWithoutP } [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task UpdateIAmAlive_DelayedHeartbeat_PreservesMaximumAndVersions(bool concurrentWrite) + [InlineData(0)] + [InlineData(1)] + public async Task UpdateIAmAlive_OwnerTimestamp_PerformsOneBlindColumnWrite(int seconds) { var (fake, entry) = await CreateNativeTable(); - var older = CreateTimedEntry(); - older.IAmAliveTime = entry.IAmAliveTime.AddHours(1); - entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); - if (concurrentWrite) - { - fake.BeforeWrite = async _ => - { - fake.BeforeWrite = null; - await Heartbeat(fake, entry); - }; - } - else - { - await Heartbeat(fake, entry); - } + var rowPath = ZooKeeperNativeFake.RowPath(entry.SiloAddress); + var heartbeatPath = ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress); + var row = fake.Nodes[rowPath]; + var table = fake.Nodes["/"]; + var heartbeat = fake.Nodes[heartbeatPath]; + entry.IAmAliveTime = entry.IAmAliveTime.AddSeconds(seconds); + entry.Status = SiloStatus.Dead; + entry.HostName = "changed-host"; + entry.SiloName = "changed-silo"; + entry.ProxyPort = 123; + entry.StartTime = entry.StartTime.AddDays(1); + entry.SuspectTimes = [Tuple.Create(CreateTimedEntry(12346).SiloAddress, entry.IAmAliveTime)]; - await Heartbeat(fake, older); + await Heartbeat(fake, entry); - var result = await Read(fake); - Assert.Equal(entry.IAmAliveTime, Assert.Single(result.Members).Item1.IAmAliveTime); - Assert.Equal("0", result.Members[0].Item2); - Assert.Equal(1, result.Version.Version); - Assert.Equal(1, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Version); - Assert.Empty(fake.Transactions); + AssertSingleHeartbeatWrite(fake, entry); + Assert.Equal(3, fake.Nodes.Count); + Assert.Same(table, fake.Nodes["/"]); + Assert.Same(row, fake.Nodes[rowPath]); + Assert.Equal(entry.IAmAliveTime, ZooKeeperBasedMembershipTable.Deserialize(fake.Nodes[heartbeatPath].Data)); + Assert.Equal(heartbeat.Version + 1, fake.Nodes[heartbeatPath].Version); + Assert.Equal(heartbeat.ChildrenVersion, fake.Nodes[heartbeatPath].ChildrenVersion); + Assert.Equal(heartbeat.Flags, fake.Nodes[heartbeatPath].Flags); } - [Fact] - public async Task UpdateIAmAlive_ConcurrentStatusUpdate_PreservesNewerTimestamp() + [Theory] + [InlineData("auth")] + [InlineData("connection")] + [InlineData("session")] + [InlineData("version")] + public async Task UpdateIAmAlive_NativeFailure_PropagatesAfterOneBlindWrite(string kind) { var (fake, entry) = await CreateNativeTable(); - var heartbeat = CreateTimedEntry(); - heartbeat.IAmAliveTime = entry.IAmAliveTime.AddHours(1); - entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); - entry.Status = SiloStatus.ShuttingDown; - fake.BeforeWrite = async _ => + var original = fake.Nodes.ToArray(); + Exception failure = kind switch { - fake.BeforeWrite = null; - Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( - fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + "auth" => new KeeperException.NoAuthException(), + "connection" => new KeeperException.ConnectionLossException(), + "session" => new KeeperException.SessionExpiredException(), + "version" => new KeeperException.BadVersionException(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)), + _ => throw new ArgumentOutOfRangeException(nameof(kind)) }; + fake.BeforeWrite = _ => Task.FromException(failure); - await Heartbeat(fake, heartbeat); + var actual = await Record.ExceptionAsync(() => Heartbeat(fake, entry)); - var result = await Read(fake); - var row = Assert.Single(result.Members); - Assert.Equal(entry.IAmAliveTime, row.Item1.IAmAliveTime); - Assert.Equal(SiloStatus.ShuttingDown, row.Item1.Status); - Assert.Equal("1", row.Item2); - Assert.Equal(2, result.Version.Version); - Assert.Equal(1, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Version); + Assert.Same(failure, actual); + AssertSingleHeartbeatWrite(fake, entry); + Assert.Equal(original, fake.Nodes.ToArray()); } [Theory] - [InlineData("read")] - [InlineData("write")] - [InlineData("retry")] - public async Task UpdateIAmAlive_RetiredRow_CompletesWithoutRecreationOrVersionWrite(string phase) + [InlineData(false)] + [InlineData(true)] + public async Task UpdateIAmAlive_MissingNode_PropagatesNativeFailure(bool missingCluster) { var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); - var cutoff = entry.IAmAliveTime.AddDays(2); - async Task Prune() + if (missingCluster) { - await ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, cutoff, TestContext.Current.CancellationToken); - fake.Transactions.Clear(); + fake.Nodes.Clear(); } + else + { + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, entry.IAmAliveTime.AddDays(1), TestContext.Current.CancellationToken); + } + + fake.Calls.Clear(); + fake.Transactions.Clear(); + var original = fake.Nodes.ToArray(); + + var failure = await Assert.ThrowsAsync(() => Heartbeat(fake, entry)); + + Assert.Equal(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), failure.getPath()); + AssertSingleHeartbeatWrite(fake, entry); + Assert.Equal(original, fake.Nodes.ToArray()); + } - if (phase == "read") + [Fact] + public async Task UpdateIAmAlive_CanceledCaller_PreservesPendingNativeWrite() + { + var (fake, entry) = await CreateNativeTable(); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var writing = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var completeWrite = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + fake.BeforeWrite = _ => { - await Prune(); + writing.SetResult(); + return completeWrite.Task; + }; + var heartbeatPath = ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress); + var originalHeartbeat = fake.Nodes[heartbeatPath]; + entry.IAmAliveTime = entry.IAmAliveTime.AddSeconds(1); + var operation = ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.SetData, cancellation.Token); + var caller = ZooKeeperBasedMembershipTable.AwaitOperationAsync(operation, cancellation.Token); + try + { + await writing.Task.WaitAsync(TestContext.Current.CancellationToken); + cancellation.Cancel(); + + var failure = await Assert.ThrowsAnyAsync(() => caller); + + Assert.Equal(cancellation.Token, failure.CancellationToken); + AssertSingleHeartbeatWrite(fake, entry); + Assert.False(operation.IsCompleted); + Assert.Same(originalHeartbeat, fake.Nodes[heartbeatPath]); } - else + finally { - fake.BeforeWrite = async _ => - { - fake.BeforeWrite = null; - await Prune(); - if (phase == "retry") - { - throw new KeeperException.BadVersionException(ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)); - } - }; + completeWrite.SetResult(); + await operation; } - entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); - await Heartbeat(fake, entry); - - Assert.Equal("/", Assert.Single(fake.Nodes).Key); - Assert.Equal(1, fake.Nodes["/"].Version); - Assert.Empty(fake.Transactions); + AssertSingleHeartbeatWrite(fake, entry); + Assert.Equal(entry.IAmAliveTime, ZooKeeperBasedMembershipTable.Deserialize(fake.Nodes[heartbeatPath].Data)); + Assert.Equal(originalHeartbeat.Version + 1, fake.Nodes[heartbeatPath].Version); } [Theory] @@ -786,7 +813,6 @@ public async Task Cleanup_ConcurrentCleaners_CompleteWithUnchangedTableVersion() [InlineData("point")] [InlineData("insert")] [InlineData("update")] - [InlineData("heartbeat")] [InlineData("cleanup")] public async Task MembershipOperations_MissingClusterHistory_PropagatesNoNode(string operation) { @@ -808,10 +834,6 @@ public async Task MembershipOperations_MissingClusterHistory_PropagatesNoNode(st [InlineData("update", "auth", true)] [InlineData("update", "connection", true)] [InlineData("update", "session", true)] - [InlineData("heartbeat", "connection", false)] - [InlineData("heartbeat", "auth", true)] - [InlineData("heartbeat", "connection", true)] - [InlineData("heartbeat", "session", true)] [InlineData("cleanup", "session", false)] [InlineData("cleanup", "auth", true)] [InlineData("cleanup", "connection", true)] @@ -843,7 +865,6 @@ public async Task MembershipOperations_InfrastructureFailure_PropagatesSameExcep [Theory] [InlineData("update")] - [InlineData("heartbeat")] [InlineData("cleanup")] public async Task MembershipOperations_ClusterRemovedDuringWrite_PropagatesNoNode(string operation) { @@ -867,7 +888,6 @@ Task RemoveCluster() [Theory] [InlineData("update")] - [InlineData("heartbeat")] [InlineData("cleanup")] public async Task NativeOperations_CanceledConflict_StopsRetrying(string operation) { @@ -918,7 +938,6 @@ public async Task NativeOperations_PreCanceledToken_IssuesNoRequests(string oper [InlineData("read")] [InlineData("point")] [InlineData("update")] - [InlineData("heartbeat")] [InlineData("cleanup")] public async Task NativeOperations_CanceledAfterRead_StartsNoFurtherRequests(string operation) { @@ -1012,7 +1031,7 @@ private static Task InvokeNative(string operation, ZooKeeperNativeFake fake, Mem "point" => ZooKeeperBasedMembershipTable.ReadCoreAsync(fake.Operations, entry.SiloAddress, cancellationToken), "insert" => ZooKeeperBasedMembershipTable.InsertRowCoreAsync(fake.Operations, entry, new TableVersion(2, "1"), cancellationToken), "update" => ZooKeeperBasedMembershipTable.UpdateRowCoreAsync(fake.Operations, entry, "0", new TableVersion(2, "1"), cancellationToken), - "heartbeat" => ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.GetData, fake.SetData, cancellationToken), + "heartbeat" => ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.SetData, cancellationToken), "cleanup" => ZooKeeperBasedMembershipTable.CleanupCoreAsync(fake.Operations, DateTime.UnixEpoch.AddDays(3), cancellationToken), _ => throw new ArgumentOutOfRangeException(nameof(operation)) }; @@ -1047,7 +1066,18 @@ private static Task Read(ZooKeeperNativeFake fake, SiloAddr ZooKeeperBasedMembershipTable.ReadCoreAsync(fake.Operations, address, TestContext.Current.CancellationToken); private static Task Heartbeat(ZooKeeperNativeFake fake, MembershipEntry entry) => - ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.GetData, fake.SetData, TestContext.Current.CancellationToken); + ZooKeeperBasedMembershipTable.UpdateIAmAliveCoreAsync(entry, fake.SetData, TestContext.Current.CancellationToken); + + private static void AssertSingleHeartbeatWrite(ZooKeeperNativeFake fake, MembershipEntry entry) + { + var path = ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress); + Assert.Equal("write " + path, Assert.Single(fake.Calls)); + var write = Assert.Single(fake.Writes); + Assert.Equal(path, write.Path); + Assert.Equal(-1, write.Version); + Assert.Equal(ZooKeeperBasedMembershipTable.Serialize(entry.IAmAliveTime), write.Data); + Assert.Empty(fake.Transactions); + } private static void AssertSet(object operation, string path, int version) { diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs index 546b1318020..726092fc72d 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs @@ -25,6 +25,7 @@ internal sealed record DeleteRequest(string Path, int Version) : Request(Path); internal List Calls { get; } = []; internal List> Transactions { get; } = []; + internal List Writes { get; } = []; internal Func? BeforeRead { get; set; } internal Func? AfterRead { get; set; } internal Func, Task>? BeforeMulti { get; set; } @@ -122,6 +123,7 @@ internal async Task Multi(List operations) internal async Task SetData(string path, byte[] data, int version) { Calls.Add("write " + path); + Writes.Add(new(path, data, version)); if (BeforeWrite is { } before) { await before(path); From 2e8e0a2c386f9eef6bfaa90a236a5c2b863ac30f Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 09:04:51 -0700 Subject: [PATCH 5/8] fix(zookeeper): isolate heartbeat from membership concurrency --- .../ZooKeeperBasedMembershipTable.cs | 55 +++------ .../ZooKeeperBasedMembershipTableUnitTests.cs | 107 ++++++++++++------ 2 files changed, 93 insertions(+), 69 deletions(-) diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index aa0673e90f0..82fc565d562 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -30,6 +30,7 @@ namespace Orleans.Runtime.Membership /// Every Silo's state is saved in /UniqueDeploymentId/IP:Port@Gen /// Every Silo's IAmAlive is saved in /UniqueDeploymentId/IP:Port@Gen/IAmAlive /// IAmAlive is saved in a separate node so owner heartbeat writes preserve the membership row's version. + /// Membership updates atomically modify the row and deployment nodes, preserving the separate heartbeat node. /// /// a node's ZK version is its ETag: /// the table version is the version of /UniqueDeploymentId @@ -320,50 +321,30 @@ public Task UpdateRowAsync(MembershipEntry entry, string etag, TableVersio internal static async Task UpdateRowCoreAsync( NativeOperations zk, MembershipEntry entry, string etag, TableVersion tableVersion, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); string rowPath = ConvertToRowPath(entry.SiloAddress); - string rowIAmAlivePath = ConvertToRowIAmAlivePath(entry.SiloAddress); var newRowData = Serialize(entry); int expectedTableVersion = int.Parse(tableVersion.VersionEtag, CultureInfo.InvariantCulture); int expectedRowVersion = int.Parse(etag, CultureInfo.InvariantCulture); - while (true) + try + { + await zk.Multi( + [ + Op.setData("/", null, expectedTableVersion), + Op.setData(rowPath, newRowData, expectedRowVersion) + ]); + return true; + } + catch (KeeperException.BadVersionException) + { + return false; + } + catch (KeeperException.NoNodeException) { cancellationToken.ThrowIfCancellationRequested(); - if ((await zk.GetData("/")).Stat.getVersion() != expectedTableVersion) - { - return false; - } - - try - { - cancellationToken.ThrowIfCancellationRequested(); - if ((await zk.GetData(rowPath)).Stat.getVersion() != expectedRowVersion) - { - return false; - } - - cancellationToken.ThrowIfCancellationRequested(); - var heartbeat = await zk.GetData(rowIAmAlivePath); - var time = new DateTime(Math.Max(entry.IAmAliveTime.Ticks, Deserialize(heartbeat.Data).Ticks), DateTimeKind.Utc); - cancellationToken.ThrowIfCancellationRequested(); - await zk.Multi( - [ - Op.setData("/", null, expectedTableVersion), - Op.setData(rowPath, newRowData, expectedRowVersion), - Op.setData(rowIAmAlivePath, Serialize(time), heartbeat.Stat.getVersion()) - ]); - return true; - } - catch (KeeperException.BadVersionException) - { - // Retry heartbeat races while retaining the caller's row and table preconditions. - } - catch (KeeperException.NoNodeException) - { - cancellationToken.ThrowIfCancellationRequested(); - await zk.GetData("/"); - return false; - } + await zk.GetData("/"); + return false; } } diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index f8974ea657d..fe2d0f30572 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -341,14 +341,22 @@ public async Task Read_MissingPointRow_ReturnsEmptyViewWithCurrentVersion() } [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task Read_ConcurrentHeartbeat_UsesLatestSeparateTimestamp(bool pointRead) + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task Read_ConcurrentHeartbeat_PreservesCanonicalFence(bool pointRead, bool afterHeartbeatRead) { var (fake, entry) = await CreateNativeTable(); + var originalTime = entry.IAmAliveTime; + var root = fake.Nodes["/"]; + var row = fake.Nodes[ZooKeeperNativeFake.RowPath(entry.SiloAddress)]; + var interleavingPath = afterHeartbeatRead + ? ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress) + : ZooKeeperNativeFake.RowPath(entry.SiloAddress); fake.AfterRead = async path => { - if (path == ZooKeeperNativeFake.RowPath(entry.SiloAddress)) + if (path == interleavingPath) { fake.AfterRead = null; entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); @@ -358,9 +366,17 @@ public async Task Read_ConcurrentHeartbeat_UsesLatestSeparateTimestamp(bool poin var result = await Read(fake, pointRead ? entry.SiloAddress : null); - Assert.Equal(entry.IAmAliveTime, Assert.Single(result.Members).Item1.IAmAliveTime); + Assert.Equal(afterHeartbeatRead ? originalTime : entry.IAmAliveTime, Assert.Single(result.Members).Item1.IAmAliveTime); Assert.Equal(1, result.Version.Version); + Assert.Equal("1", result.Version.VersionEtag); Assert.Equal("0", result.Members[0].Item2); + Assert.Equal(1, fake.Calls.Count(call => call == "sync /")); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(entry.SiloAddress))); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress))); + Assert.Same(root, fake.Nodes["/"]); + Assert.Same(row, fake.Nodes[ZooKeeperNativeFake.RowPath(entry.SiloAddress)]); + Assert.Equal(entry.IAmAliveTime, ZooKeeperBasedMembershipTable.Deserialize( + fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Data)); } [Fact] @@ -435,6 +451,7 @@ public async Task Insert_DuplicateOrStaleVersion_RollsBackTransaction(bool dupli public async Task UpdateRow_OptimisticPreconditions_ControlAtomicMutation(string etag, int tableVersion) { var (fake, entry) = await CreateNativeTable(); + var original = fake.Nodes.ToArray(); entry.Status = SiloStatus.Dead; var updated = await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( @@ -443,63 +460,92 @@ public async Task UpdateRow_OptimisticPreconditions_ControlAtomicMutation(string var expected = etag == "0" && tableVersion == 1; Assert.Equal(expected, updated); + Assert.Equal("multi", Assert.Single(fake.Calls)); + Assert.Collection(Assert.Single(fake.Transactions), + operation => AssertSet(operation, "/", tableVersion), + operation => AssertSet(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), int.Parse(etag))); + if (!expected) + { + Assert.Equal(original, fake.Nodes.ToArray()); + } + var result = await Read(fake); Assert.Equal(expected ? SiloStatus.Dead : SiloStatus.Active, Assert.Single(result.Members).Item1.Status); Assert.Equal(expected ? "1" : "0", result.Members[0].Item2); Assert.Equal(expected ? 2 : 1, result.Version.Version); - if (expected) - { - Assert.Collection(Assert.Single(fake.Transactions), - operation => AssertSet(operation, "/", 1), - operation => AssertSet(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), 0), - operation => AssertSet(operation, ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0)); - } - else - { - Assert.Empty(fake.Transactions); - } } - [Fact] - public async Task UpdateRow_ConcurrentHeartbeat_RetriesAndPreservesMaximum() + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task UpdateRow_HeartbeatPreservesOriginalTokens_ForStatusAndVotes(bool voteUpdate, bool concurrentHeartbeat) { var (fake, entry) = await CreateNativeTable(); + var snapshot = await Read(fake); + var (update, etag) = Assert.Single(snapshot.Members); + var next = snapshot.Version.Next(); var heartbeat = CreateTimedEntry(); heartbeat.IAmAliveTime = entry.IAmAliveTime.AddDays(1); - fake.BeforeMulti = async _ => + if (concurrentHeartbeat) + { + fake.BeforeMulti = async _ => + { + fake.BeforeMulti = null; + await Heartbeat(fake, heartbeat); + }; + } + else { - fake.BeforeMulti = null; await Heartbeat(fake, heartbeat); - }; - entry.Status = SiloStatus.Dead; + } + + if (voteUpdate) + { + update.SuspectTimes = [Tuple.Create(CreateTimedEntry(12346).SiloAddress, heartbeat.IAmAliveTime)]; + } + else + { + update.Status = SiloStatus.Dead; + } + + fake.Calls.Clear(); Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( - fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + fake.Operations, update, etag, next, TestContext.Current.CancellationToken)); + + Assert.Equal(concurrentHeartbeat ? ["multi", "write " + ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)] : ["multi"], fake.Calls); + Assert.Collection(Assert.Single(fake.Transactions), + operation => AssertSet(operation, "/", snapshot.Version.Version), + operation => AssertSet(operation, ZooKeeperNativeFake.RowPath(entry.SiloAddress), int.Parse(etag))); + Assert.Equal(1, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)].Version); var result = await Read(fake); var row = Assert.Single(result.Members); - Assert.Equal(SiloStatus.Dead, row.Item1.Status); + Assert.Equal(update.Status, row.Item1.Status); + Assert.Equal(update.SuspectTimes, row.Item1.SuspectTimes); Assert.Equal(heartbeat.IAmAliveTime, row.Item1.IAmAliveTime); Assert.Equal("1", row.Item2); Assert.Equal(2, result.Version.Version); - Assert.Equal(2, fake.Transactions.Count); - AssertSet(fake.Transactions[0][2], ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 0); - AssertSet(fake.Transactions[1][2], ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress), 1); } [Fact] - public async Task UpdateRow_ProposedHeartbeat_AdvancesTimestamp() + public async Task UpdateRow_PreservesSeparateHeartbeatNode() { var (fake, entry) = await CreateNativeTable(); + var heartbeat = fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)]; entry.IAmAliveTime = entry.IAmAliveTime.AddDays(1); entry.Status = SiloStatus.Dead; Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( fake.Operations, entry, "0", new TableVersion(2, "1"), TestContext.Current.CancellationToken)); + Assert.Equal("multi", Assert.Single(fake.Calls)); + Assert.Same(heartbeat, fake.Nodes[ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress)]); var result = await Read(fake); var row = Assert.Single(result.Members); - Assert.Equal(entry.IAmAliveTime, row.Item1.IAmAliveTime); + Assert.Equal(ZooKeeperBasedMembershipTable.Deserialize(heartbeat.Data), row.Item1.IAmAliveTime); Assert.Equal(SiloStatus.Dead, row.Item1.Status); Assert.Equal("1", row.Item2); Assert.Equal(2, result.Version.Version); @@ -830,7 +876,6 @@ public async Task MembershipOperations_MissingClusterHistory_PropagatesNoNode(st [InlineData("read", "auth", false)] [InlineData("point", "connection", false)] [InlineData("insert", "session", true)] - [InlineData("update", "auth", false)] [InlineData("update", "auth", true)] [InlineData("update", "connection", true)] [InlineData("update", "session", true)] @@ -887,7 +932,6 @@ Task RemoveCluster() } [Theory] - [InlineData("update")] [InlineData("cleanup")] public async Task NativeOperations_CanceledConflict_StopsRetrying(string operation) { @@ -937,7 +981,6 @@ public async Task NativeOperations_PreCanceledToken_IssuesNoRequests(string oper [Theory] [InlineData("read")] [InlineData("point")] - [InlineData("update")] [InlineData("cleanup")] public async Task NativeOperations_CanceledAfterRead_StartsNoFurtherRequests(string operation) { From 59390e021a228769a9d3fe21b8e8159fd9f36171 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 09:22:08 -0700 Subject: [PATCH 6/8] test(zookeeper): clarify native result constructor diagnostics --- .../ZooKeeperBasedMembershipTableUnitTests.cs | 12 ++++++++++++ .../ZooKeeperNativeFake.cs | 13 +++++++++---- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index fe2d0f30572..3f3bf4de336 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -1067,6 +1067,18 @@ public void NativeFake_IncompatibleSdkMember_ReportsActionableDiagnostic(string Assert.Contains(typeof(Op).Assembly.FullName!, failure.Message, StringComparison.Ordinal); } + [Fact] + public void NativeFake_IncompatibleSdkConstructor_ReportsActionableDiagnostic() + { + var failure = Assert.Throws(() => + ZooKeeperNativeFake.CreateResult()); + + Assert.Contains($"{typeof(DataResult).FullName}..ctor", failure.Message, StringComparison.Ordinal); + Assert.Contains("expected a non-public instance constructor accepting ()", failure.Message, StringComparison.Ordinal); + Assert.Contains("Update the fake for the installed ZooKeeperNetEx API", failure.Message, StringComparison.Ordinal); + Assert.Contains(typeof(Op).Assembly.FullName!, failure.Message, StringComparison.Ordinal); + } + private static Task InvokeNative(string operation, ZooKeeperNativeFake fake, MembershipEntry entry, CancellationToken cancellationToken) => operation switch { diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs index 726092fc72d..ced2df00d5c 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperNativeFake.cs @@ -6,7 +6,6 @@ using org.apache.zookeeper; using org.apache.zookeeper.data; using Orleans.Runtime.Membership; -using Xunit; namespace UnitTests.MembershipTests; @@ -174,9 +173,15 @@ private static InvalidOperationException ApiMismatch(Type type, string methodNam new($"ZooKeeperNativeFake cannot decode {type.FullName}.{methodName}: {detail}. " + $"Update the fake for the installed ZooKeeperNetEx API ({typeof(Op).Assembly.FullName})."); - private static T CreateResult(params object[] arguments) => - Assert.IsType(Activator.CreateInstance(typeof(T), BindingFlags.Instance | BindingFlags.NonPublic, - binder: null, args: arguments, culture: null)); + internal static T CreateResult(params object[] arguments) + { + var parameterTypes = arguments.Select(argument => argument.GetType()).ToArray(); + var constructor = typeof(T).GetConstructor(BindingFlags.Instance | BindingFlags.NonPublic, + binder: null, types: parameterTypes, modifiers: null) + ?? throw ApiMismatch(typeof(T), ".ctor", + $"expected a non-public instance constructor accepting ({string.Join(", ", parameterTypes.Select(type => type.FullName))})"); + return (T)constructor.Invoke(arguments); + } private static Stat CreateStat(Node node) => new(0, 0, 0, 0, node.Version, node.ChildrenVersion, 0, 0, 0, 0, 0); private static string Parent(string path) => path.LastIndexOf('/') is var index && index > 0 ? path[..index] : "/"; From 385c0b9fcbcaa01a1db38a90ce845e2b50bc9ccd Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 10:23:11 -0700 Subject: [PATCH 7/8] perf(zookeeper): remove redundant cleanup coordination --- .../ZooKeeperBasedMembershipTable.cs | 91 +++++---- .../ZooKeeperBasedMembershipTableUnitTests.cs | 182 +++++++++++++++++- 2 files changed, 224 insertions(+), 49 deletions(-) diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index 82fc565d562..fdfdeae9b4c 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -106,8 +106,6 @@ await UsingZookeeper(rootConnectionString, async zk => try { await zk.createAsync(this.clusterPath, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); - cancellationToken.ThrowIfCancellationRequested(); - await zk.sync(this.clusterPath); //if we got here we know that we've just created the deployment path with version=0 LogInformationCreatedNewDeploymentPath(this.clusterPath); } @@ -161,10 +159,11 @@ internal static Task ReadAllAsync(string deploymentConnecti internal static async Task ReadCoreAsync( NativeOperations zk, SiloAddress? siloAddress, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + await zk.Sync("/"); + // Retries retain this session's ordered view. while (true) { - cancellationToken.ThrowIfCancellationRequested(); - await zk.Sync("/"); cancellationToken.ThrowIfCancellationRequested(); Stat before; IEnumerable addresses; @@ -274,12 +273,6 @@ await zk.Multi( { return false; } - catch (KeeperException.NoNodeException) - { - cancellationToken.ThrowIfCancellationRequested(); - await zk.GetData("/"); - return false; - } } /// @@ -387,12 +380,7 @@ internal static Task UpdateIAmAliveCoreAsync( public Task DeleteMembershipTableEntriesAsync(string clusterId, CancellationToken cancellationToken = default) { string pathToDelete = "/" + clusterId; - return UsingZookeeper(rootConnectionString, async zk => - { - await DeleteRecursive(zk, pathToDelete, cancellationToken); - cancellationToken.ThrowIfCancellationRequested(); - await zk.sync(pathToDelete); - }, cancellationToken); + return UsingZookeeper(rootConnectionString, zk => DeleteRecursive(zk, pathToDelete, cancellationToken), cancellationToken); } /// @@ -540,46 +528,53 @@ public Task CleanupDefunctSiloEntriesAsync(DateTimeOffset beforeDate, Cancellati internal static async Task CleanupCoreAsync(NativeOperations zk, DateTimeOffset beforeDate, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + var children = await zk.GetChildren("/"); var cutoff = beforeDate.UtcDateTime; + await Task.WhenAll(children.Children.Select(child => CleanupRowAsync(zk, "/" + child, cutoff, cancellationToken))); + return true; + } + + private static async Task CleanupRowAsync(NativeOperations zk, string rowPath, DateTime cutoff, CancellationToken cancellationToken) + { + var heartbeatPath = rowPath + "/IAmAlive"; while (true) { - var table = await ReadCoreAsync(zk, null, cancellationToken); - var candidates = table.Members.Where(row => row.Item1.Status == SiloStatus.Dead - && row.Item1.StartTime < cutoff && row.Item1.IAmAliveTime < cutoff - && row.Item1.SuspectTimes?.Any(vote => vote.Item2 >= cutoff) != true).ToList(); - if (candidates.Count == 0) - { - return true; - } - - foreach (var (entry, etag) in candidates) + cancellationToken.ThrowIfCancellationRequested(); + try { - cancellationToken.ThrowIfCancellationRequested(); - try + var row = await zk.GetData(rowPath); + var entry = Deserialize(row.Data); + if (entry.Status != SiloStatus.Dead || entry.StartTime >= cutoff + || entry.SuspectTimes?.Any(vote => vote.Item2 >= cutoff) == true) { - var heartbeatPath = ConvertToRowIAmAlivePath(entry.SiloAddress); - var heartbeat = await zk.GetData(heartbeatPath); - if (Deserialize(heartbeat.Data) >= cutoff) - { - continue; - } - - cancellationToken.ThrowIfCancellationRequested(); - await zk.Multi( - [ - Op.delete(heartbeatPath, heartbeat.Stat.getVersion()), - Op.delete(ConvertToRowPath(entry.SiloAddress), int.Parse(etag, CultureInfo.InvariantCulture)) - ]); + return; } - catch (KeeperException.BadVersionException) - { - // Re-evaluate eligibility after a concurrent row or heartbeat update. - } - catch (KeeperException.NoNodeException) + + cancellationToken.ThrowIfCancellationRequested(); + var heartbeat = await zk.GetData(heartbeatPath); + if (Deserialize(heartbeat.Data) >= cutoff) { - cancellationToken.ThrowIfCancellationRequested(); - await zk.GetData("/"); + return; } + + cancellationToken.ThrowIfCancellationRequested(); + await zk.Multi( + [ + Op.delete(heartbeatPath, heartbeat.Stat.getVersion()), + Op.delete(rowPath, row.Stat.getVersion()) + ]); + return; + } + catch (KeeperException.BadVersionException) + { + // Re-evaluate only this row: its version and heartbeat version guard eligibility. + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); + return; } } } diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index 3f3bf4de336..40ad8ac62b9 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -274,7 +274,8 @@ public async Task Read_ConcurrentStatusUpdate_RetriesWithCoherentVersion(bool po Assert.Equal("1", row.Item2); Assert.Equal(2, result.Version.Version); Assert.Equal("2", result.Version.VersionEtag); - Assert.Equal(2, fake.Calls.Count(call => call == "sync /")); + Assert.Equal(1, fake.Calls.Count(call => call == "sync /")); + Assert.Equal(2, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(entry.SiloAddress))); Assert.Equal("sync /", fake.Calls[0]); } @@ -736,6 +737,7 @@ await ZooKeeperBasedMembershipTable.CleanupCoreAsync( { Assert.Equal(3, fake.Nodes.Count); Assert.Empty(fake.Transactions); + Assert.Equal(["children /", "read " + ZooKeeperNativeFake.RowPath(entry.SiloAddress)], fake.Calls); Assert.Equal(status, Assert.Single((await Read(fake)).Members).Item1.Status); } } @@ -854,6 +856,180 @@ public async Task Cleanup_ConcurrentCleaners_CompleteWithUnchangedTableVersion() Assert.Equal(2, fake.Transactions.Count); } + [Fact] + public async Task Cleanup_SelectedBatch_ReadsEachCandidateOnce() + { + var (fake, first) = await CreateNativeTable(SiloStatus.Dead); + var second = CreateTimedEntry(12346); + second.Status = SiloStatus.Dead; + var active = CreateTimedEntry(12347); + Assert.True(await Insert(fake, second, 1)); + Assert.True(await Insert(fake, active, 2)); + fake.Calls.Clear(); + fake.Transactions.Clear(); + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + + Assert.Equal(8, fake.Calls.Count); + Assert.Equal(1, fake.Calls.Count(call => call == "children /")); + foreach (var entry in new[] { first, second, active }) + { + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(entry.SiloAddress))); + Assert.Equal(entry.Status == SiloStatus.Dead ? 1 : 0, + fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress))); + } + + Assert.Equal(2, fake.Transactions.Count); + Assert.All(fake.Transactions, transaction => Assert.Equal(2, transaction.Count)); + Assert.Equal(3, fake.Nodes.Count); + Assert.Equal(3, fake.Nodes["/"].Version); + Assert.True(fake.Nodes.ContainsKey(ZooKeeperNativeFake.RowPath(active.SiloAddress))); + Assert.True(fake.Nodes.ContainsKey(ZooKeeperNativeFake.HeartbeatPath(active.SiloAddress))); + } + + [Fact] + public async Task Cleanup_Conflict_RechecksOnlyContestedRow() + { + var (fake, first) = await CreateNativeTable(SiloStatus.Dead); + var second = CreateTimedEntry(12346); + second.Status = SiloStatus.Dead; + Assert.True(await Insert(fake, second, 1)); + fake.Calls.Clear(); + fake.Transactions.Clear(); + fake.BeforeMulti = async requests => + { + if (requests[^1].Path == ZooKeeperNativeFake.RowPath(first.SiloAddress)) + { + fake.BeforeMulti = null; + first.IAmAliveTime = first.IAmAliveTime.AddHours(1); + await Heartbeat(fake, first); + } + }; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + + Assert.Equal(1, fake.Calls.Count(call => call == "children /")); + Assert.Equal(2, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(first.SiloAddress))); + Assert.Equal(2, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.HeartbeatPath(first.SiloAddress))); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(second.SiloAddress))); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.HeartbeatPath(second.SiloAddress))); + Assert.Equal(3, fake.Transactions.Count); + Assert.Equal(11, fake.Calls.Count); + Assert.Single(fake.Nodes); + Assert.Equal(2, fake.Nodes["/"].Version); + } + + [Fact] + public async Task Cleanup_UnrelatedCanonicalUpdate_PreservesSelectedBatch() + { + var (fake, dead) = await CreateNativeTable(SiloStatus.Dead); + var active = CreateTimedEntry(12346); + Assert.True(await Insert(fake, active, 1)); + fake.Calls.Clear(); + fake.Transactions.Clear(); + fake.BeforeRead = async path => + { + if (path == ZooKeeperNativeFake.RowPath(dead.SiloAddress)) + { + fake.BeforeRead = null; + active.SuspectTimes = [Tuple.Create(dead.SiloAddress, DateTime.UnixEpoch.AddDays(2))]; + Assert.True(await ZooKeeperBasedMembershipTable.UpdateRowCoreAsync( + fake.Operations, active, "0", new TableVersion(3, "2"), TestContext.Current.CancellationToken)); + } + }; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + + Assert.Equal(6, fake.Calls.Count); + Assert.Equal(1, fake.Calls.Count(call => call == "children /")); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(dead.SiloAddress))); + Assert.Equal(1, fake.Calls.Count(call => call == "read " + ZooKeeperNativeFake.RowPath(active.SiloAddress))); + Assert.Equal(2, fake.Transactions.Count); + Assert.Equal(3, fake.Nodes.Count); + Assert.Equal(3, fake.Nodes["/"].Version); + var surviving = ZooKeeperBasedMembershipTable.Deserialize( + fake.Nodes[ZooKeeperNativeFake.RowPath(active.SiloAddress)].Data); + Assert.Equal(SiloStatus.Active, surviving.Status); + Assert.Equal(active.SuspectTimes, surviving.SuspectTimes); + } + + [Fact] + public async Task Cleanup_ParallelFailure_WaitsForOutstandingRequests() + { + var (fake, first) = await CreateNativeTable(); + var second = CreateTimedEntry(12346); + Assert.True(await Insert(fake, second, 1)); + var firstStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var failed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var completeFirst = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var failure = new KeeperException.NoAuthException(); + fake.BeforeRead = path => + { + if (path == ZooKeeperNativeFake.RowPath(first.SiloAddress)) + { + firstStarted.SetResult(); + return completeFirst.Task; + } + + if (path == ZooKeeperNativeFake.RowPath(second.SiloAddress)) + { + failed.SetResult(); + return Task.FromException(failure); + } + + return Task.CompletedTask; + }; + + var operation = ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + try + { + await Task.WhenAll(firstStarted.Task, failed.Task).WaitAsync(TestContext.Current.CancellationToken); + Assert.False(operation.IsCompleted); + } + finally + { + completeFirst.SetResult(); + } + + Assert.Same(failure, await Record.ExceptionAsync(() => operation)); + Assert.Equal(5, fake.Nodes.Count); + } + + [Fact] + public async Task Cleanup_ConcurrentRetirement_PreservesOtherRowsAuthorizationFailure() + { + var (fake, first) = await CreateNativeTable(SiloStatus.Dead); + var second = CreateTimedEntry(12346); + Assert.True(await Insert(fake, second, 1)); + var failure = new KeeperException.NoAuthException(); + fake.BeforeRead = path => + { + if (path == ZooKeeperNativeFake.RowPath(first.SiloAddress)) + { + fake.Nodes.Remove(ZooKeeperNativeFake.HeartbeatPath(first.SiloAddress)); + fake.Nodes.Remove(path); + fake.Nodes["/"] = fake.Nodes["/"] with { ChildrenVersion = fake.Nodes["/"].ChildrenVersion + 1 }; + } + else if (path == ZooKeeperNativeFake.RowPath(second.SiloAddress)) + { + return Task.FromException(failure); + } + + return Task.CompletedTask; + }; + + var actual = await Record.ExceptionAsync(() => ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken)); + + Assert.Same(failure, actual); + Assert.Equal(3, fake.Nodes.Count); + Assert.True(fake.Nodes.ContainsKey(ZooKeeperNativeFake.RowPath(second.SiloAddress))); + } + [Theory] [InlineData("read")] [InlineData("point")] @@ -870,6 +1046,10 @@ public async Task MembershipOperations_MissingClusterHistory_PropagatesNoNode(st Assert.Equal("/", failure.getPath()); Assert.Empty(fake.Nodes); + if (operation == "insert") + { + Assert.Equal("multi", Assert.Single(fake.Calls)); + } } [Theory] From 309ab40b8e1953102534cd8281e802605a8080d8 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 18 Sep 2026 10:36:23 -0700 Subject: [PATCH 8/8] fix(zookeeper): surface missing cleanup heartbeats --- .../ZooKeeperBasedMembershipTable.cs | 28 ++++++++- .../ZooKeeperBasedMembershipTableUnitTests.cs | 58 +++++++++++++++++++ 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs index fdfdeae9b4c..f81f18834fd 100644 --- a/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs +++ b/src/Orleans.Clustering.ZooKeeper/ZooKeeperBasedMembershipTable.cs @@ -541,9 +541,20 @@ private static async Task CleanupRowAsync(NativeOperations zk, string rowPath, D while (true) { cancellationToken.ThrowIfCancellationRequested(); + DataResult row; + try + { + row = await zk.GetData(rowPath); + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); + return; + } + try { - var row = await zk.GetData(rowPath); var entry = Deserialize(row.Data); if (entry.Status != SiloStatus.Dead || entry.StartTime >= cutoff || entry.SuspectTimes?.Any(vote => vote.Item2 >= cutoff) == true) @@ -573,8 +584,19 @@ await zk.Multi( catch (KeeperException.NoNodeException) { cancellationToken.ThrowIfCancellationRequested(); - await zk.GetData("/"); - return; + try + { + await zk.GetData(rowPath); + } + catch (KeeperException.NoNodeException) + { + cancellationToken.ThrowIfCancellationRequested(); + await zk.GetData("/"); + return; + } + + // Atomic retirement removes both nodes; a surviving row has a missing heartbeat. + throw; } } } diff --git a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs index 40ad8ac62b9..b33fc088b03 100644 --- a/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs +++ b/test/Extensions/Orleans.Clustering.ZooKeeper.Tests/ZooKeeperBasedMembershipTableUnitTests.cs @@ -856,6 +856,64 @@ public async Task Cleanup_ConcurrentCleaners_CompleteWithUnchangedTableVersion() Assert.Equal(2, fake.Transactions.Count); } + [Fact] + public async Task Cleanup_RowRetiredBeforeHeartbeatRead_CompletesWithUnchangedVersion() + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + fake.AfterRead = async path => + { + if (path == ZooKeeperNativeFake.RowPath(entry.SiloAddress)) + { + fake.AfterRead = null; + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + } + }; + + await ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken); + + Assert.Single(fake.Nodes); + Assert.Equal(1, fake.Nodes["/"].Version); + Assert.Single(fake.Transactions); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Cleanup_MissingHeartbeatWithSurvivingRow_PropagatesNoNode(bool duringMulti) + { + var (fake, entry) = await CreateNativeTable(SiloStatus.Dead); + var heartbeatPath = ZooKeeperNativeFake.HeartbeatPath(entry.SiloAddress); + var rowPath = ZooKeeperNativeFake.RowPath(entry.SiloAddress); + var row = fake.Nodes[rowPath]; + var root = fake.Nodes["/"]; + if (duringMulti) + { + fake.BeforeMulti = _ => + { + fake.Nodes.Remove(heartbeatPath); + return Task.CompletedTask; + }; + } + else + { + fake.Nodes.Remove(heartbeatPath); + } + + var failure = await Assert.ThrowsAsync(() => + ZooKeeperBasedMembershipTable.CleanupCoreAsync( + fake.Operations, DateTime.UnixEpoch.AddDays(2), TestContext.Current.CancellationToken)); + + Assert.Equal(heartbeatPath, failure.getPath()); + Assert.Same(row, fake.Nodes[rowPath]); + Assert.Same(root, fake.Nodes["/"]); + Assert.False(fake.Nodes.ContainsKey(heartbeatPath)); + Assert.Equal(2, fake.Calls.Count(call => call == "read " + rowPath)); + Assert.DoesNotContain("read /", fake.Calls); + Assert.Equal(duringMulti ? 1 : 0, fake.Transactions.Count); + } + [Fact] public async Task Cleanup_SelectedBatch_ReadsEachCandidateOnce() {