Add wait-for-subscribers feature

src/contrib/cluster/Akka.Cluster.Tools/PublishSubscribe/DistributedMessages.cs

@@ -353,7 +353,7 @@ public override string ToString()
     /// TBD
     /// </summary>
     [Serializable]
-    public sealed class Publish : IDistributedPubSubMessage, IEquatable<Publish>, IWrappedMessage
+    public class Publish : IDistributedPubSubMessage, IEquatable<Publish>, IWrappedMessage
     {
         /// <summary>
         /// TBD
@@ -415,6 +415,20 @@ public override string ToString()
             return $"Publish<topic:{Topic}, sendOneToEachGroup:{SendOneMessageToEachGroup}, message:{Message}>";
         }
     }
+
+    public sealed class PublishWithAck : Publish
+    {
+        public PublishWithAck(string topic, object message, TimeSpan timeout, bool sendOneMessageToEachGroup = false) : base(topic, message, sendOneMessageToEachGroup)
+        {
+            Timeout = timeout;
+        }
+
+        public TimeSpan Timeout { get; }
+    }
+
+    public sealed record PublishFailed(PublishWithAck Message);
+
+    public sealed record PublishSucceeded(PublishWithAck Message);
 
     /// <summary>
     /// TBD

src/contrib/cluster/Akka.Cluster.Tools/PublishSubscribe/DistributedPubSubMediator.cs

@@ -14,6 +14,7 @@
 using Akka.Cluster.Tools.PublishSubscribe.Internal;
 using Akka.Event;
 using Akka.Pattern;
+using Akka.Remote;
 using Akka.Routing;
 using Akka.Util;
 using Status = Akka.Cluster.Tools.PublishSubscribe.Internal.Status;
@@ -102,6 +103,7 @@ public class DistributedPubSubMediator : ReceiveActor, IWithTimers
     {
         private const string GossipTimerKey = "GossipTimer";
         private const string PruneTimerKey = "PruneTimer";
+        private const string PruneBufferTimerKey = "PruneBufferTimer";
 
         /// <summary>
         /// TBD
@@ -127,20 +129,19 @@ public static Props Props(DistributedPubSubSettings settings)
         private readonly string _topicPrefix;
         private readonly PubSubCache _cache;
 
+        private readonly int _maxBufferPerTopic;
+        private readonly TimeSpan _bufferedMessageTimeoutCheckInterval;
+        private readonly Dictionary<string, List<BufferedMessage>> _bufferedMessages = new();
+
         public ITimerScheduler Timers { get; set; }
 
         /// <summary>
         /// TBD
         /// </summary>
         public IImmutableDictionary<Address, long> OwnVersions
-        {
-            get
-            {
-                return _registry
-                    .Select(entry => new KeyValuePair<Address, long>(entry.Key, entry.Value.Version))
-                    .ToImmutableDictionary(kv => kv.Key, kv => kv.Value);
-            }
-        }
+            => _registry
+                .Select(entry => new KeyValuePair<Address, long>(entry.Key, entry.Value.Version))
+                .ToImmutableDictionary(kv => kv.Key, kv => kv.Value);
 
         /// <summary>
         /// TBD
@@ -161,6 +162,9 @@ public DistributedPubSubMediator(DistributedPubSubSettings settings)
             _log = Context.GetLogger();
 
             _role = settings.Role;
+            _maxBufferPerTopic = settings.MaxBufferedMessagePerTopic;
+            _bufferedMessageTimeoutCheckInterval = settings.BufferedMessageTimeoutCheckInterval;
+
             _pruneInterval = new TimeSpan(_settings.RemovedTimeToLive.Ticks / 2);
             _buffer = new PerGroupingBuffer();
 
@@ -195,7 +199,7 @@ public DistributedPubSubMediator(DistributedPubSubSettings settings)
                     new Router(_settings.RoutingLogic, routees.ToArray()).Route(
                         Internal.Utils.WrapIfNeeded(send.Message), Sender);
                 else
-                    IgnoreOrSendToDeadLetters(send);
+                    IgnoreOrSendToDeadLetters(send, Sender);
             });
             Receive<SendToAll>(sendToAll =>
             {
@@ -320,14 +324,54 @@ public DistributedPubSubMediator(DistributedPubSubSettings settings)
                             if (!_registry.TryGetValue(bucket.Owner, out var myBucket))
                                 myBucket = new Bucket(bucket.Owner);
 
-                            if (bucket.Version > myBucket.Version)
-                                _registry[bucket.Owner] = new Bucket(myBucket.Owner, bucket.Version, myBucket.Content.SetItems(bucket.Content));
+                            if (bucket.Version <= myBucket.Version)
+                                continue;
+
+                            // Create a bucket diff, we're only interested in new items being added
+                            var keys = new HashSet<string>(bucket.Content.Keys);
+                            keys.ExceptWith(new HashSet<string>(myBucket.Content.Keys));
+                            if (keys.Count > 0)
+                            {
+                                // Send the diff to ourselves
+                                Self.Tell(new NewBucketKeysAdded(keys));
+                            }
+
+                            // Merge remote bucket with ours
+                            var newBucket = new Bucket(myBucket.Owner, bucket.Version, myBucket.Content.SetItems(bucket.Content));
+                            _registry[bucket.Owner] = newBucket;
                         }
                     }
                 }
             });
             Receive<GossipTick>(_ => HandleGossip());
             Receive<Prune>(_ => HandlePrune());
+            Receive<NewBucketKeysAdded>(subs =>
+            {
+                foreach (var key in subs.Topics)
+                {
+                    if (!_bufferedMessages.TryGetValue(key, out var buffer)) 
+                        continue;
+
+                    foreach (var bufferedMessage in buffer)
+                    {
+                        Self.Tell(bufferedMessage.Message, bufferedMessage.Sender);
+                    }
+
+                    _bufferedMessages.Remove(key);
+                }
+            });
+            Receive<PruneBufferTick>(_ =>
+            {
+                foreach (var buffer in _bufferedMessages.Values)
+                {
+                    var removed = buffer.Where(bufferedMessage => bufferedMessage.Deadline.IsOverdue).ToArray();
+                    foreach (var removedMsg in removed)
+                    {
+                        buffer.Remove(removedMsg);
+                        IgnoreOrSendToDeadLetters(removedMsg.Message, removedMsg.Sender);
+                    }
+                }
+            });
             Receive<Terminated>(terminated =>
             {
                 var key = Internal.Utils.MakeKey(terminated.ActorRef);
@@ -417,6 +461,32 @@ public DistributedPubSubMediator(DistributedPubSubSettings settings)
             });
         }
 
+        private void BufferMessageOrDeadLetter(PublishWithAck message, IActorRef sender)
+        {
+            var topic = message.Topic;
+            var encodedTopic = _cache.EncodeName(topic);
+            var key = _cache.MakeKey(Self.Path, encodedTopic);
+
+            if (!_bufferedMessages.TryGetValue(key, out var buffer))
+            {
+                buffer = [];
+                _bufferedMessages[key] = buffer;
+            }
+
+            if(buffer.Count >= _maxBufferPerTopic)
+            {
+                _log.Warning("PublishWithAck buffer overflowed for topic [{0}]. New message inserted: {1}", topic, message);
+                while (buffer.Count >= _maxBufferPerTopic)
+                {
+                    var removed = buffer[0];
+                    buffer.RemoveAt(0);
+                    IgnoreOrSendToDeadLetters(removed.Message, removed.Sender);
+                }
+            }
+
+            buffer.Add(new BufferedMessage(message, new Deadline(DateTime.UtcNow + message.Timeout), sender));
+        }
+
         private bool OtherHasNewerVersions(IImmutableDictionary<Address, long> versions)
         {
             return versions.Any(entry =>
@@ -438,11 +508,8 @@ private IEnumerable<Bucket> CollectDelta(IImmutableDictionary<Address, long> ver
             }
 
             var count = 0;
-            foreach (var entry in filledOtherVersions)
+            foreach (var (owner, v) in filledOtherVersions)
             {
-                var owner = entry.Key;
-                var v = entry.Value;
-
                 if (!_registry.TryGetValue(owner, out var bucket))
                     bucket = new Bucket(owner);
 
@@ -493,15 +560,21 @@ private void HandleRegisterTopic(IActorRef actorRef)
         private void PutToRegistry(string key, IActorRef value)
         {
             var v = NextVersion();
-            if (!_registry.TryGetValue(_cluster.SelfAddress, out var bucket))
-                _registry.Add(_cluster.SelfAddress,
-                    new Bucket(_cluster.SelfAddress, v, ImmutableDictionary<string, ValueHolder>.Empty.Add(key, new ValueHolder(v, value))));
-            else
-                _registry[_cluster.SelfAddress] = new Bucket(bucket.Owner, v, bucket.Content.SetItem(key, new ValueHolder(v, value)));
+            var newBucket = _registry.TryGetValue(_cluster.SelfAddress, out var bucket)
+                ? new Bucket(bucket.Owner, v, bucket.Content.SetItem(key, new ValueHolder(v, value)))
+                : new Bucket(_cluster.SelfAddress, v, ImmutableDictionary<string, ValueHolder>.Empty.Add(key, new ValueHolder(v, value)));
+
+            _registry[_cluster.SelfAddress] = newBucket;
+            Self.Tell(new NewBucketKeysAdded([key]));
         }
 
-        private void IgnoreOrSendToDeadLetters(IWrappedMessage message)
+        private void IgnoreOrSendToDeadLetters(object message, IActorRef sender)
         {
+            if (message is PublishWithAck needAck)
+            {
+                sender.Tell(new PublishFailed(needAck));
+            }
+
             if (_settings.SendToDeadLettersWhenNoSubscribers)
             {
                 var topic = message switch
@@ -514,7 +587,7 @@ private void IgnoreOrSendToDeadLetters(IWrappedMessage message)
                 // Use the specialized DeadLetterWithNoSubscribers class to clearly indicate
                 // that the message was not delivered because there were no subscribers,
                 // not because the mediator itself is dead.
-                var deadLetter = new DeadLetterWithNoSubscribers(message, topic, Sender, Context.Self);
+                var deadLetter = new DeadLetterWithNoSubscribers(message, topic, sender, Context.Self);
                 Context.System.DeadLetters.Tell(deadLetter);
             }
         }
@@ -527,9 +600,17 @@ private void PublishMessage(string path, IWrappedMessage publish, bool allButSel
                 if (r == null) continue;
                 r.Forward(publish.Message);
                 counter++;
+                if(publish is PublishWithAck needAck)
+                    Sender.Tell(new PublishSucceeded(needAck));
             }
 
-            if (counter == 0) IgnoreOrSendToDeadLetters(publish);
+            if (counter == 0)
+            {
+                if(publish is PublishWithAck needAck)
+                    BufferMessageOrDeadLetter(needAck, Sender);
+                else
+                    IgnoreOrSendToDeadLetters(publish, Sender);
+            }
             return;
 
             IEnumerable<IActorRef> Refs()
@@ -555,7 +636,10 @@ private void PublishToEachGroup(string path, Publish publish)
 
             if (groups.Count == 0)
             {
-                IgnoreOrSendToDeadLetters(publish);
+                if(publish is PublishWithAck needAck)
+                    BufferMessageOrDeadLetter(needAck, Sender);
+                else
+                    IgnoreOrSendToDeadLetters(publish, Sender);
             }
             else
             {
@@ -565,6 +649,9 @@ private void PublishToEachGroup(string path, Publish publish)
                     if (routees.Length != 0)
                         new Router(_settings.RoutingLogic, routees).Route(wrappedMessage, Sender);
                 }
+
+                if(publish is PublishWithAck needAck)
+                    Sender.Tell(new PublishSucceeded(needAck));
             }
         }
 
@@ -582,8 +669,8 @@ private IEnumerable<KeyValuePair<string, Routee>> ExtractGroups(string prefix, s
 
         private void HandlePrune()
         {
-            var modifications = new Dictionary<Address, Bucket>();
-            foreach (var (owner, bucket) in _registry)
+            var modifications = new List<Bucket>();
+            foreach (var bucket in _registry.Values)
             {
                 var oldRemoved = bucket.Content
                     .Where(kv => kv.Value.Ref.IsNobody() && (bucket.Version - kv.Value.Version) > _settings.RemovedTimeToLive.TotalMilliseconds)
@@ -592,13 +679,13 @@ private void HandlePrune()
 
                 if (oldRemoved.Length > 0)
                 {
-                    modifications.Add(owner, new Bucket(bucket.Owner, bucket.Version, bucket.Content.RemoveRange(oldRemoved)));
+                    modifications.Add(new Bucket(bucket.Owner, bucket.Version, bucket.Content.RemoveRange(oldRemoved)));
                 }
             }
 
             foreach (var entry in modifications)
             {
-                _registry[entry.Key] = entry.Value;
+                _registry[entry.Owner] = entry;
             }
         }
 
@@ -634,6 +721,7 @@ protected override void PreStart()
             //Start periodic gossip to random nodes in cluster
             Timers.StartPeriodicTimer(GossipTimerKey, GossipTick.Instance, _settings.GossipInterval, _settings.GossipInterval, Self);
             Timers.StartPeriodicTimer(PruneTimerKey, Prune.Instance, _pruneInterval, _pruneInterval, Self);
+            Timers.StartPeriodicTimer(PruneBufferTimerKey, PruneBufferTick.Instance, _bufferedMessageTimeoutCheckInterval, _bufferedMessageTimeoutCheckInterval, Self);
         }
 
         /// <summary>