Fix Cluster Sharding private replicator recovery

src/contrib/cluster/Akka.Cluster.Sharding.Tests/ClusterShardingReplicatorResiliencySpec.cs

@@ -0,0 +1,165 @@
+// -----------------------------------------------------------------------
+//  <copyright file="ClusterShardingReplicatorResiliencySpec.cs" company="Akka.NET Project">
+//      Copyright (C) 2009-2022 Lightbend Inc. <http://www.lightbend.com>
+//      Copyright (C) 2013-2025 .NET Foundation <https://github.com/akkadotnet/akka.net>
+//  </copyright>
+// -----------------------------------------------------------------------
+
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using Akka.Actor;
+using Akka.Configuration;
+using Akka.TestKit;
+using FluentAssertions;
+using FluentAssertions.Extensions;
+using Xunit;
+
+namespace Akka.Cluster.Sharding.Tests;
+
+public class ClusterShardingReplicatorResiliencySpec : AkkaSpec
+{
+    private sealed record ShardEnvelope(string EntityId, string Message);
+
+    private sealed class EntityActor : ReceiveActor
+    {
+        public EntityActor()
+        {
+            Receive<string>(message => Sender.Tell(message));
+        }
+    }
+
+    private static readonly HashCodeMessageExtractor MessageExtractor = HashCodeMessageExtractor.Create(
+        10,
+        message => message is ShardEnvelope envelope ? envelope.EntityId : null,
+        message => message is ShardEnvelope envelope ? envelope.Message : message);
+
+    private static Config SpecConfig =>
+        ConfigurationFactory.ParseString(@"
+            akka.loglevel = DEBUG
+            akka.actor.provider = cluster
+            akka.remote.dot-netty.tcp.port = 0
+
+            akka.cluster.sharding.state-store-mode = ddata
+            akka.cluster.sharding.remember-entities = on
+            akka.cluster.sharding.remember-entities-store = ddata
+            akka.cluster.sharding.distributed-data.majority-min-cap = 1
+            akka.cluster.sharding.distributed-data.durable.keys = []")
+            .WithFallback(ClusterSharding.DefaultConfig());
+
+    public ClusterShardingReplicatorResiliencySpec(ITestOutputHelper helper)
+        : base(SpecConfig, output: helper)
+    {
+    }
+
+    protected override void AtStartup()
+    {
+        var cluster = Cluster.Get(Sys);
+        cluster.Join(cluster.SelfAddress);
+        AwaitAssert(() =>
+            cluster.ReadView.Members.Count(member => member.Status == MemberStatus.Up).Should().Be(1));
+    }
+
+    [Fact]
+    public async Task Private_replicator_should_recover_at_the_same_path_without_restarting_consumers()
+    {
+        const string typeName = "replicator-resiliency";
+        const string replicatorPath = "/system/sharding/replicator";
+        const string firstEntityId = "entity-1";
+        var firstShardId = MessageExtractor.ShardId(firstEntityId);
+        var existingShardEntityId = Enumerable.Range(2, 100)
+            .Select(i => $"entity-{i}")
+            .First(id => MessageExtractor.ShardId(id) == firstShardId);
+        var newShardEntityId = Enumerable.Range(2, 100)
+            .Select(i => $"entity-{i}")
+            .First(id => MessageExtractor.ShardId(id) != firstShardId);
+        var region = ClusterSharding.Get(Sys).Start(
+            typeName,
+            Props.Create<EntityActor>(),
+            ClusterShardingSettings.Create(Sys),
+            MessageExtractor);
+
+        region.Tell(new ShardEnvelope(firstEntityId, "before"));
+        await ExpectMsgAsync("before");
+        var shard = LastSender.Path.Parent;
+
+        var coordinator = await Sys.ActorSelection(
+                $"/system/sharding/{typeName}Coordinator/singleton/coordinator")
+            .ResolveOne(3.Seconds());
+        var coordinatorWatcher = CreateTestProbe();
+        await coordinatorWatcher.WatchAsync(coordinator);
+        var shardWatcher = CreateTestProbe();
+        await shardWatcher.WatchAsync(await Sys.ActorSelection(shard).ResolveOne(3.Seconds()));
+
+        var firstReplicator = await Sys.ActorSelection(replicatorPath).ResolveOne(3.Seconds());
+        await WatchAsync(firstReplicator);
+        Sys.Stop(firstReplicator);
+        await ExpectTerminatedAsync(firstReplicator);
+
+        IActorRef replacement = null;
+        await AwaitAssertAsync(async () =>
+        {
+            replacement = await Sys.ActorSelection(replicatorPath).ResolveOne(1.Seconds());
+            replacement.Should().NotBe(firstReplicator);
+            replacement.Path.ToStringWithoutAddress().Should().Be(replicatorPath);
+        }, 10.Seconds());
+
+        // Existing shard: exercises its existing DData remember-entities store through the new replicator.
+        region.Tell(new ShardEnvelope(existingShardEntityId, "after-existing"));
+        await ExpectMsgAsync("after-existing");
+
+        // New shard: exercises the existing DData coordinator and the updated provider.
+        region.Tell(new ShardEnvelope(newShardEntityId, "after-new"));
+        await ExpectMsgAsync("after-new");
+
+        await coordinatorWatcher.ExpectNoMsgAsync(500.Milliseconds());
+        await shardWatcher.ExpectNoMsgAsync(500.Milliseconds());
+    }
+}
+
+public class PersistentShardingReplicatorCompatibilitySpec : AkkaSpec
+{
+    private sealed class NoOpMessageExtractor : IMessageExtractor
+    {
+        public string EntityId(object message) => null;
+        public object EntityMessage(object message) => message;
+        public string ShardId(object message) => null;
+        public string ShardId(string entityId, object messageHint = null) => "1";
+    }
+
+    private static Config SpecConfig =>
+        ConfigurationFactory.ParseString(@"
+            akka.actor.provider = cluster
+            akka.remote.dot-netty.tcp.port = 0
+
+            akka.cluster.sharding.state-store-mode = persistence
+            akka.cluster.sharding.remember-entities = on
+            akka.cluster.sharding.remember-entities-store = ddata")
+            .WithFallback(ClusterSharding.DefaultConfig());
+
+    public PersistentShardingReplicatorCompatibilitySpec(ITestOutputHelper helper)
+        : base(SpecConfig, output: helper)
+    {
+    }
+
+    protected override void AtStartup()
+    {
+        var cluster = Cluster.Get(Sys);
+        cluster.Join(cluster.SelfAddress);
+        AwaitAssert(() =>
+            cluster.ReadView.Members.Count(member => member.Status == MemberStatus.Up).Should().Be(1));
+    }
+
+    [Fact]
+    public async Task Persistence_with_DData_remember_entities_setting_should_not_create_a_replicator()
+    {
+        ClusterSharding.Get(Sys).Start(
+            "persistent-compatibility",
+            Props.Empty,
+            ClusterShardingSettings.Create(Sys),
+            new NoOpMessageExtractor());
+
+        await Assert.ThrowsAsync<ActorNotFoundException>(() =>
+            Sys.ActorSelection("/system/sharding/replicator").ResolveOne(500.Milliseconds()));
+    }
+}

src/contrib/cluster/Akka.Cluster.Sharding/ClusterShardingGuardian.cs

@@ -12,6 +12,7 @@
 using Akka.Cluster.Sharding.Internal;
 using Akka.Cluster.Tools.Singleton;
 using Akka.DistributedData;
+using Akka.Event;
 using Akka.Pattern;
 
 namespace Akka.Cluster.Sharding
@@ -23,6 +24,9 @@ internal sealed class ClusterShardingGuardian : ReceiveActor
     {
         #region messages
 
+        private sealed record ReplicatorTerminated(string Role, IActorRef Replicator)
+            : INoSerializationVerificationNeeded;
+
         /// <summary>
         /// TBD
         /// </summary>
@@ -148,9 +152,14 @@ public StartProxy(
 
         private readonly Cluster _cluster = Cluster.Get(Context.System);
         private readonly ClusterSharding _sharding = ClusterSharding.Get(Context.System);
+        private readonly ILoggingAdapter _log = Context.GetLogger();
 
         private readonly int _majorityMinCap = Context.System.Settings.Config.GetInt("akka.cluster.sharding.distributed-data.majority-min-cap", 0);
         private ImmutableDictionary<string, IActorRef> _replicatorsByRole = ImmutableDictionary<string, IActorRef>.Empty;
+        private ImmutableDictionary<string, ReplicatorSettings> _replicatorSettingsByRole = ImmutableDictionary<string, ReplicatorSettings>.Empty;
+        private ImmutableDictionary<string, ImmutableDictionary<string, DDataRememberEntitiesProvider>> _rememberEntitiesProvidersByRole =
+            ImmutableDictionary<string, ImmutableDictionary<string, DDataRememberEntitiesProvider>>.Empty;
+        private bool _clusterShuttingDown;
 
         private readonly ConcurrentDictionary<string, IActorRef> _regions;
         private readonly ConcurrentDictionary<string, IActorRef> _proxies;
@@ -165,6 +174,8 @@ public ClusterShardingGuardian(
         {
             _regions = regions;
             _proxies = proxies;
+            _cluster.Subscribe(Self, ClusterEvent.SubscriptionInitialStateMode.InitialStateAsEvents,
+                typeof(ClusterEvent.ClusterShuttingDown));
 
             Receive<Start>(start =>
             {
@@ -185,7 +196,14 @@ public ClusterShardingGuardian(
                         switch (rememberEntitiesProvider)
                         {
                             case RememberEntitiesStore.DData:
-                                rememberEntitiesStoreProvider = new DDataRememberEntitiesProvider(start.TypeName, settings, _majorityMinCap, replicator);
+                                var ddataProvider = new DDataRememberEntitiesProvider(start.TypeName, settings, _majorityMinCap, replicator);
+                                rememberEntitiesStoreProvider = ddataProvider;
+                                var role = settings.Role ?? string.Empty;
+                                var providersForRole = _rememberEntitiesProvidersByRole.TryGetValue(role, out var providers)
+                                    ? providers
+                                    : ImmutableDictionary<string, DDataRememberEntitiesProvider>.Empty;
+                                _rememberEntitiesProvidersByRole = _rememberEntitiesProvidersByRole.SetItem(
+                                    role, providersForRole.SetItem(start.TypeName, ddataProvider));
                                 break;
                             case RememberEntitiesStore.Eventsourced:
                                 rememberEntitiesStoreProvider = new EventSourcedRememberEntitiesProvider(start.TypeName, settings);
@@ -278,6 +296,8 @@ public ClusterShardingGuardian(
                 }
             });
 
+            Receive<ReplicatorTerminated>(HandleReplicatorTermination);
+
             Receive<Terminated>(msg =>
             {
                 if (!_typeLookup.TryGetValue(msg.ActorRef, out var typeName)) 
@@ -294,6 +314,8 @@ public ClusterShardingGuardian(
                 if(_proxies.TryGetValue(typeName, out var proxyActor) && proxyActor.Equals(msg.ActorRef))
                     _proxies.TryRemove(typeName, out _);
             });
+
+            Receive<ClusterEvent.ClusterShuttingDown>(_ => _clusterShuttingDown = true);
         }
 
         internal static ReplicatorSettings GetReplicatorSettings(ClusterShardingSettings shardingSettings)
@@ -314,20 +336,54 @@ private IActorRef Replicator(ClusterShardingSettings settings)
             {
                 // one replicator per role
                 var role = settings.Role ?? string.Empty;
-                if (_replicatorsByRole.TryGetValue(role, out var aref)) return aref;
-                else
-                {
-                    var name = string.IsNullOrEmpty(settings.Role) ? "replicator" : Uri.EscapeDataString(settings.Role) + "Replicator";
-                    var replicatorRef = Context.ActorOf(DistributedData.Replicator.Props(GetReplicatorSettings(settings)), name);
+                if (_replicatorsByRole.TryGetValue(role, out var replicator))
+                    return replicator;
 
-                    _replicatorsByRole = _replicatorsByRole.SetItem(role, replicatorRef);
-                    return replicatorRef;
-                }
+                var replicatorSettings = GetReplicatorSettings(settings);
+                _replicatorSettingsByRole = _replicatorSettingsByRole.SetItem(role, replicatorSettings);
+                return CreateReplicator(role, replicatorSettings);
             }
             else
                 return Context.System.DeadLetters;
         }
 
+        private IActorRef CreateReplicator(string role, ReplicatorSettings settings)
+        {
+            var replicator = Context.ActorOf(DistributedData.Replicator.Props(settings), ReplicatorName(role));
+            Context.WatchWith(replicator, new ReplicatorTerminated(role, replicator));
+            _replicatorsByRole = _replicatorsByRole.SetItem(role, replicator);
+            return replicator;
+        }
+
+        private void HandleReplicatorTermination(ReplicatorTerminated terminated)
+        {
+            if (!_replicatorsByRole.TryGetValue(terminated.Role, out var current)
+                || !current.Equals(terminated.Replicator))
+                return;
+
+            _replicatorsByRole = _replicatorsByRole.Remove(terminated.Role);
+            if (_clusterShuttingDown || _cluster.IsTerminated)
+                return;
+
+            _log.Error(
+                "Cluster Sharding DData replicator [{0}] for role [{1}] terminated; recreating it at the same path",
+                terminated.Replicator.Path, DisplayRole(terminated.Role));
+            var replacement = CreateReplicator(terminated.Role, _replicatorSettingsByRole[terminated.Role]);
+            if (_rememberEntitiesProvidersByRole.TryGetValue(terminated.Role, out var providers))
+            {
+                foreach (var provider in providers.Values)
+                    provider.ReplaceReplicator(terminated.Replicator, replacement);
+            }
+            Context.System.EventStream.Publish(new ReplicatorChanged(terminated.Replicator, replacement));
+            _log.Info("Recreated Cluster Sharding DData replicator [{0}] for role [{1}]",
+                replacement.Path, DisplayRole(terminated.Role));
+        }
+
+        private static string ReplicatorName(string role) =>
+            string.IsNullOrEmpty(role) ? "replicator" : Uri.EscapeDataString(role) + "Replicator";
+
+        private static string DisplayRole(string role) => string.IsNullOrEmpty(role) ? "<all>" : role;
+
         private string CoordinatorPath(string encName)
         {
             return (Self.Path / CoordinatorSingletonManagerName(encName) / "singleton" / "coordinator").ToStringWithoutAddress();
@@ -337,5 +393,11 @@ private static string CoordinatorSingletonManagerName(string encName)
         {
             return encName + "Coordinator";
         }
+
+        protected override void PostStop()
+        {
+            _cluster.Unsubscribe(Self);
+            base.PostStop();
+        }
     }
 }