Pass MessageWatchdog to MessageWriter and MessageWriters instead of an Action

Author: stidsborg

Core/Cleipnir.ResilientFunctions/CoreRuntime/Invocation/InvocationHelper.cs

@@ -389,7 +389,7 @@ await _functionStore.BulkScheduleFunctions(
     }
 
     public MessageWriter CreateMessageWriter(StoredId storedId)
-        => new MessageWriter(storedId, _functionStore.MessageStore, Serializer, _replicaId, _messageWatchdog.Notify);
+        => new MessageWriter(storedId, _functionStore.MessageStore, Serializer, _replicaId, _messageWatchdog);
 
     public Effect CreateEffect(StoredId storedId, FlowId flowId, IReadOnlyList<StoredEffect> storedEffects, FlowTimeouts flowTimeouts, IStorageSession? storageSession, FlowExecutionState flowExecutionState)
     {

Core/Cleipnir.ResilientFunctions/FunctionsRegistry.cs

@@ -284,7 +284,7 @@ public FuncRegistration<TParam, TReturn> RegisterFunc<TParam, TReturn>(
                 _functionStore,
                 serializer,
                 ClusterInfo.ReplicaId,
-                _messageWatchdog.Notify
+                _messageWatchdog
             );
 
             var registration = new FuncRegistration<TParam, TReturn>(
@@ -372,7 +372,7 @@ private ParamlessRegistration RegisterParamless(
                 _functionStore,
                 serializer,
                 ClusterInfo.ReplicaId,
-                _messageWatchdog.Notify
+                _messageWatchdog
             );
 
             var registration = new ParamlessRegistration(
@@ -460,7 +460,7 @@ public ActionRegistration<TParam> RegisterAction<TParam>(
                 _functionStore,
                 serializer,
                 ClusterInfo.ReplicaId,
-                _messageWatchdog.Notify
+                _messageWatchdog
             );
             var registration = new ActionRegistration<TParam>(
                 flowType,

Core/Cleipnir.ResilientFunctions/Messaging/MessageWriter.cs

@@ -1,22 +1,37 @@
-using System;
 using System.Threading.Tasks;
 using Cleipnir.ResilientFunctions.CoreRuntime.Serialization;
+using Cleipnir.ResilientFunctions.CoreRuntime.Watchdogs;
 using Cleipnir.ResilientFunctions.Domain;
 using Cleipnir.ResilientFunctions.Storage;
 
 namespace Cleipnir.ResilientFunctions.Messaging;
 
-public class MessageWriter(StoredId storedIdId, IMessageStore messageStore, ISerializer eventSerializer, ReplicaId publisherReplica, Action? notifyDelivery = null)
+public class MessageWriter
 {
+    private readonly StoredId _storedId;
+    private readonly IMessageStore _messageStore;
+    private readonly ISerializer _eventSerializer;
+    private readonly ReplicaId _publisherReplica;
+    private readonly MessageWatchdog? _messageWatchdog;
+
+    internal MessageWriter(StoredId storedId, IMessageStore messageStore, ISerializer eventSerializer, ReplicaId publisherReplica, MessageWatchdog? messageWatchdog = null)
+    {
+        _storedId = storedId;
+        _messageStore = messageStore;
+        _eventSerializer = eventSerializer;
+        _publisherReplica = publisherReplica;
+        _messageWatchdog = messageWatchdog;
+    }
+
     public async Task AppendMessage<TMessage>(TMessage message, string? idempotencyKey = null, string? sender = null, string? receiver = null) where TMessage : class
     {
-        var eventJson = eventSerializer.Serialize(message, message.GetType());
-        var eventType = eventSerializer.SerializeType(message.GetType());
+        var eventJson = _eventSerializer.Serialize(message, message.GetType());
+        var eventType = _eventSerializer.SerializeType(message.GetType());
 
-        var storedMessage = new StoredMessage(eventJson, eventType, Position: 0, Replica: publisherReplica, IdempotencyKey: idempotencyKey, Sender: sender, Receiver: receiver);
-        await messageStore.AppendMessages([new StoredIdAndMessage(storedIdId, storedMessage)]);
+        var storedMessage = new StoredMessage(eventJson, eventType, Position: 0, Replica: _publisherReplica, IdempotencyKey: idempotencyKey, Sender: sender, Receiver: receiver);
+        await _messageStore.AppendMessages([new StoredIdAndMessage(_storedId, storedMessage)]);
 
         // Wake the MessageWatchdog so the appended message is delivered now rather than on the next poll.
-        notifyDelivery?.Invoke();
+        _messageWatchdog?.Notify();
     }
-}
+}

Core/Cleipnir.ResilientFunctions/Messaging/MessageWriters.cs

@@ -1,9 +1,9 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using Cleipnir.ResilientFunctions.CoreRuntime;
 using Cleipnir.ResilientFunctions.CoreRuntime.Invocation;
 using Cleipnir.ResilientFunctions.CoreRuntime.Serialization;
+using Cleipnir.ResilientFunctions.CoreRuntime.Watchdogs;
 using Cleipnir.ResilientFunctions.Domain;
 using Cleipnir.ResilientFunctions.Storage;
 
@@ -15,31 +15,31 @@ public class MessageWriters
     private readonly IFunctionStore _functionStore;
     private readonly ISerializer _serializer;
     private readonly ReplicaId _publisherReplica;
-    private readonly Action? _notifyDelivery;
+    private readonly MessageWatchdog? _messageWatchdog;
 
-    public MessageWriters(
+    internal MessageWriters(
         StoredType storedType,
         IFunctionStore functionStore,
         ISerializer serializer,
         ReplicaId publisherReplica,
-        Action? notifyDelivery = null)
+        MessageWatchdog? messageWatchdog = null)
     {
         _storedType = storedType;
         _functionStore = functionStore;
         _serializer = serializer;
         _publisherReplica = publisherReplica;
-        _notifyDelivery = notifyDelivery;
+        _messageWatchdog = messageWatchdog;
     }
 
     public MessageWriter For(FlowInstance instance)
     {
         var storedId = StoredId.Create(_storedType, instance.Value);
-        return new MessageWriter(storedId, _functionStore.MessageStore, _serializer, _publisherReplica, _notifyDelivery);
+        return new MessageWriter(storedId, _functionStore.MessageStore, _serializer, _publisherReplica, _messageWatchdog);
     }
 
     internal MessageWriter For(StoredId storedId)
     {
-        return new MessageWriter(storedId, _functionStore.MessageStore, _serializer, _publisherReplica, _notifyDelivery);
+        return new MessageWriter(storedId, _functionStore.MessageStore, _serializer, _publisherReplica, _messageWatchdog);
     }
 
     public async Task AppendMessages(IReadOnlyList<BatchedMessage> messages)
@@ -57,6 +57,6 @@ public async Task AppendMessages(IReadOnlyList<BatchedMessage> messages)
         await _functionStore.MessageStore.AppendMessages(storedIdAndMessages);
 
         // Wake the MessageWatchdog so the appended messages are delivered now rather than on the next poll.
-        _notifyDelivery?.Invoke();
+        _messageWatchdog?.Notify();
     }
 }