Add message-driven RestartExecutions to FlowsManager

Core/Cleipnir.ResilientFunctions.Tests/Messaging/TestTemplates/MessagesSubscriptionTests.cs

@@ -21,10 +21,12 @@ public abstract class MessagesSubscriptionTests
     // These tests hand-roll a QueueManager, which delegates message deletion to IMessageClearer. They don't
     // assert on store cleanup, so a no-op stub suffices.
     private static readonly IMessageClearer StubMessageClearer = new NoopMessageClearer();
+    private static readonly ClusterInfo StubClusterInfo = new(ReplicaId.NewId());
 
     private sealed class NoopMessageClearer : IMessageClearer
     {
         public Task Clear(IReadOnlyList<long> positions) => Task.CompletedTask;
+        public void ReopenPositions(IEnumerable<long> positions) { }
     }
 
     public abstract Task EventsSubscriptionSunshineScenario();
@@ -569,7 +571,7 @@ protected async Task QueueManagerFailsOnMessageDeserializationError(Task<IFuncti
             {
 
                 var flowTimeouts = new FlowTimeouts();
-                var flowsManager = new FlowsManager(functionStore);
+                var flowsManager = new FlowsManager(functionStore, StubMessageClearer, StubClusterInfo);
                 var flowState = flowsManager.CreateFlowState(workflow.StoredId, flowTimeouts, completed: ForeverTask.Instance);
                 var queueManager = new QueueManager(
                     workflow.FlowId,
@@ -633,7 +635,7 @@ protected async Task RegisteredTimeoutIsRemovedWhenPullingMessage(Task<IFunction
             {
                 storedId = workflow.StoredId;
                 var minimumTimeout = new FlowTimeouts();
-                var flowsManager = new FlowsManager(functionStore);
+                var flowsManager = new FlowsManager(functionStore, StubMessageClearer, StubClusterInfo);
                 var flowState = flowsManager.CreateFlowState(workflow.StoredId, minimumTimeout, completed: ForeverTask.Instance);
                 var queueManager = new QueueManager(
                     workflow.FlowId,
@@ -695,7 +697,7 @@ protected async Task PullEnvelopeReturnsEnvelopeWithReceiverAndSender(Task<IFunc
             {
 
                 var flowTimeouts = new FlowTimeouts();
-                var flowsManager = new FlowsManager(functionStore);
+                var flowsManager = new FlowsManager(functionStore, StubMessageClearer, StubClusterInfo);
                 var flowState = flowsManager.CreateFlowState(workflow.StoredId, flowTimeouts, completed: ForeverTask.Instance);
                 var queueManager = new QueueManager(
                     workflow.FlowId,

Core/Cleipnir.ResilientFunctions/CoreRuntime/FlowsManager.cs

@@ -2,6 +2,9 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Cleipnir.ResilientFunctions.CoreRuntime.Invocation;
+using Cleipnir.ResilientFunctions.CoreRuntime.Watchdogs;
+using Cleipnir.ResilientFunctions.Domain;
 using Cleipnir.ResilientFunctions.Messaging;
 using Cleipnir.ResilientFunctions.Storage;
 
@@ -11,9 +14,26 @@ public class FlowsManager
 {
     private readonly Dictionary<StoredId, FlowExecutionState> _dict = new();
     private readonly IFunctionStore _functionStore;
+    private readonly IMessageClearer _messageClearer;
+    private readonly ClusterInfo _clusterInfo;
+    private IFlowRestarter? _restarter;
     private readonly Lock _lock = new();
 
-    public FlowsManager(IFunctionStore functionStore) => _functionStore = functionStore;
+    internal FlowsManager(
+        IFunctionStore functionStore,
+        IMessageClearer messageClearer,
+        ClusterInfo clusterInfo)
+    {
+        _functionStore = functionStore;
+        _messageClearer = messageClearer;
+        _clusterInfo = clusterInfo;
+    }
+
+    /// <summary>
+    /// Supplies the per-type flow restarter (the Invoker). Late-bound because the Invoker is constructed after this
+    /// manager during registration.
+    /// </summary>
+    internal void SetRestarter(IFlowRestarter restarter) => _restarter = restarter;
 
     public FlowExecutionState CreateFlowState(StoredId id, FlowTimeouts timeouts, Task completed)
     {
@@ -53,6 +73,53 @@ public Task Push(IReadOnlyList<StoredMessages> messagesByFlow)
         return Task.WhenAll(tasks);
     }
 
+    /// <summary>
+    /// Restarts (claims for this replica) the targeted flows that are not already owned, then hands each restarted
+    /// flow - together with the in-hand messages - to the <see cref="ScheduleRestartFromWatchdog"/> delegate so it
+    /// resumes executing. Flows that could not be claimed have their positions reopened in the message clearer
+    /// (dropped from the ignore-set without deleting them from the store, since their actual owner still needs them).
+    /// </summary>
+    public async Task RestartExecutions(IEnumerable<StoredMessages> messages)
+    {
+        var groups = messages
+            .GroupBy(m => m.StoredId)
+            .ToDictionary(g => g.Key, g => g.ToList());
+
+        var results = await _functionStore
+            .RestartExecutionsWithoutMessages(groups.Keys.ToList(), _clusterInfo.ReplicaId);
+
+        // Flows that could not be restarted (already owned/running elsewhere, or no longer present) were never
+        // delivered to here, yet the MessageWatchdog optimistically marked their positions as pushed. Drop those
+        // positions from the clearer's ignore-set so they can be re-fetched - without deleting them from the store,
+        // since the actual owner still needs them.
+        var notRestartedPositions = groups
+            .Where(kv => !results.ContainsKey(kv.Key))
+            .SelectMany(kv => kv.Value)
+            .SelectMany(storedMessages => storedMessages.Messages)
+            .Select(message => message.Position)
+            .ToList();
+        _messageClearer.ReopenPositions(notRestartedPositions);
+
+        // Resume each restarted flow, supplying the messages we already hold so it does not re-fetch them. The
+        // claim + flow snapshot returned by RestartExecutionsWithoutMessages is everything the delegate needs, so
+        // no further store round-trip or re-claim is performed.
+        foreach (var (storedId, storedFlowWithEffects) in results)
+        {
+            var inHandMessages = groups[storedId]
+                .SelectMany(storedMessages => storedMessages.Messages)
+                .ToList();
+
+            var restartedFunction = new RestartedFunction(
+                storedFlowWithEffects.StoredFlow,
+                storedFlowWithEffects.Effects,
+                inHandMessages,
+                storedFlowWithEffects.StorageSession
+            );
+
+            await _restarter!.ScheduleRestart(storedId, restartedFunction, onCompletion: () => { });
+        }
+    }
+
     /// <summary>
     /// Wakes the target flow so it consumes the just-published message by applying the durable interrupt
     /// (interrupted flag + expires=0) - the suspend-race guard and watchdog backstop, so the message is never

Core/Cleipnir.ResilientFunctions/CoreRuntime/FlowsManagers.cs

@@ -2,6 +2,8 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Cleipnir.ResilientFunctions.CoreRuntime.Watchdogs;
+using Cleipnir.ResilientFunctions.Domain;
 using Cleipnir.ResilientFunctions.Messaging;
 using Cleipnir.ResilientFunctions.Storage;
 
@@ -16,10 +18,19 @@ namespace Cleipnir.ResilientFunctions.CoreRuntime;
 public class FlowsManagers
 {
     private readonly Dictionary<StoredType, FlowsManager> _managers = new();
+
     private readonly IFunctionStore _functionStore;
+    private readonly MessageClearer _messageClearer;
+    private readonly ClusterInfo _clusterInfo;
+
     private readonly Lock _lock = new();
 
-    public FlowsManagers(IFunctionStore functionStore) => _functionStore = functionStore;
+    internal FlowsManagers(IFunctionStore functionStore, MessageClearer messageClearer, ClusterInfo clusterInfo)
+    {
+        _functionStore = functionStore;
+        _messageClearer = messageClearer;
+        _clusterInfo = clusterInfo;
+    }
 
     public FlowsManager GetOrCreate(StoredType storedType)
     {
@@ -28,7 +39,7 @@ public FlowsManager GetOrCreate(StoredType storedType)
             if (_managers.TryGetValue(storedType, out var existing))
                 return existing;
 
-            return _managers[storedType] = new FlowsManager(_functionStore);
+            return _managers[storedType] = new FlowsManager(_functionStore, _messageClearer, _clusterInfo);
         }
     }
 
@@ -38,14 +49,24 @@ public FlowsManager GetOrCreate(StoredType storedType)
             return _managers.GetValueOrDefault(storedType);
     }
 
-    public Task Push(IReadOnlyList<StoredMessages> messagesByFlow)
+    public Task Push(IReadOnlyList<StoredMessages> messages)
     {
-        List<Task> tasks = new();
-        foreach (var group in messagesByFlow.GroupBy(m => m.StoredId.Type))
-            if (TryGet(group.Key) is { } manager)
-                tasks.Add(manager.Push(group.ToList()));
+        List<Task> messageDeliveries;
+        lock (_lock)
+        {
+            var notRunning = messages
+                .Where(msg => !_managers.ContainsKey(msg.StoredId.Type));
 
-        return Task.WhenAll(tasks);
+            var running = messages
+                .Where(msg => _managers.ContainsKey(msg.StoredId.Type))
+                .GroupBy(msg => msg.StoredId.Type)
+                .Select(g => _managers[g.Key].Push(g.ToList()))
+                .ToList();
+
+            messageDeliveries = running;
+        }
+
+        return Task.WhenAll(messageDeliveries);
     }
 
     public IReadOnlyList<StoredId> FilterOwned(IEnumerable<StoredId> ids)

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

@@ -0,0 +1,15 @@
+using System;
+using System.Threading.Tasks;
+using Cleipnir.ResilientFunctions.Storage;
+
+namespace Cleipnir.ResilientFunctions.CoreRuntime.Invocation;
+
+/// <summary>
+/// Resumes execution of an already-claimed flow from the snapshot returned by the store's restart call. Implemented
+/// by <see cref="Invoker{TParam,TReturn}"/> (which owns the user function and drives the execution loop) and held by
+/// the non-generic <see cref="FlowsManager"/> so it can restart flows without knowing their parameter/return types.
+/// </summary>
+internal interface IFlowRestarter
+{
+    Task ScheduleRestart(StoredId storedId, RestartedFunction restartedFunction, Action onCompletion);
+}

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

@@ -11,7 +11,7 @@
 
 namespace Cleipnir.ResilientFunctions.CoreRuntime.Invocation;
 
-public class Invoker<TParam, TReturn>
+public class Invoker<TParam, TReturn> : IFlowRestarter
 {
     private readonly FlowType _flowType;
     private readonly StoredType _storedType;
@@ -236,7 +236,10 @@ internal async Task ScheduleRestart(StoredId storedId, RestartedFunction rf, Act
             finally{ _flowsManager.RemoveFlow(storedId, flowState); onCompletion(); }
         });
     }
-
+
+    Task IFlowRestarter.ScheduleRestart(StoredId storedId, RestartedFunction restartedFunction, Action onCompletion)
+        => ScheduleRestart(storedId, restartedFunction, onCompletion);
+
     private async Task<PreparedInvocation> PrepareForInvocation(FlowId flowId, StoredId storedId, TParam param, StoredId? parent, InitialState? initialState, Task completed)
     {
         var disposables = new List<IDisposable>(capacity: 3);

Core/Cleipnir.ResilientFunctions/CoreRuntime/Watchdogs/IMessageClearer.cs

@@ -4,11 +4,14 @@
 namespace Cleipnir.ResilientFunctions.CoreRuntime.Watchdogs;
 
 /// <summary>
-/// The slice of <see cref="MessageClearer"/> a QueueManager depends on: deleting handled messages from the store
-/// and dropping their positions from the watchdog's ignore-set. Exists so tests that hand-roll a QueueManager can
-/// pass a no-op stub instead of a fully wired clearer.
+/// The slice of <see cref="MessageClearer"/> its collaborators depend on: the QueueManager deletes handled messages
+/// from the store (and drops their positions from the watchdog's ignore-set) via <see cref="Clear"/>, while the
+/// FlowsManager drops the positions of flows it could not restart back out of the ignore-set - without deleting them
+/// from the store - via <see cref="ReopenPositions"/>. Exists so tests can pass a no-op stub instead of a fully
+/// wired clearer.
 /// </summary>
 internal interface IMessageClearer
 {
     Task Clear(IReadOnlyList<long> positions);
+    void ReopenPositions(IEnumerable<long> positions);
 }

Core/Cleipnir.ResilientFunctions/CoreRuntime/Watchdogs/MessageClearer.cs

@@ -40,6 +40,13 @@ public void MarkPushed(IEnumerable<long> positions)
                 _pushedPositions.Add(position);
     }
 
+    public void ReopenPositions(IEnumerable<long> positions)
+    {
+        lock (_pushedPositionsLock)
+            foreach (var position in positions)
+                _pushedPositions.Remove(position);
+    }
+
     /// <summary>Snapshot of the not-yet-cleared positions, passed to the store as the fetch ignore-set.</summary>
     public IReadOnlyList<long> NonClearedPositions()
     {

Core/Cleipnir.ResilientFunctions/FunctionsRegistry.cs

@@ -43,10 +43,15 @@ public FunctionsRegistry(IFunctionStore functionStore, Settings? settings = null
         _shutdownCoordinator = new ShutdownCoordinator();
         _settings = SettingsWithDefaults.Default.Merge(settings);
         var utcNow = _settings.UtcNow;
-        _flowsManagers = new FlowsManagers(_functionStore);
-
+        _messageClearer = new MessageClearer(
+            _functionStore.MessageStore,
+            _settings.UnhandledExceptionHandler,
+            _settings.WatchdogCheckFrequency
+        );
         ClusterInfo = new ClusterInfo(ReplicaId.NewId());
-
+
+        _flowsManagers = new FlowsManagers(_functionStore, _messageClearer, ClusterInfo);
+
         _postponedWatchdog = new PostponedWatchdog(
             _functionStore,
             _shutdownCoordinator,
@@ -75,12 +80,6 @@ public FunctionsRegistry(IFunctionStore functionStore, Settings? settings = null
             utcNow
         );
 
-        _messageClearer = new MessageClearer(
-            _functionStore.MessageStore,
-            _settings.UnhandledExceptionHandler,
-            _settings.WatchdogCheckFrequency
-        );
-
         _messageWatchdog = new MessageWatchdog(
             _functionStore.MessageStore,
             _flowsManagers,
@@ -257,15 +256,17 @@ public FuncRegistration<TParam, TReturn> RegisterFunc<TParam, TReturn>(
                 settings?.ClearChildrenAfterCapture ?? true,
                 _messageClearer
             );
+            var flowsManager = _flowsManagers.GetOrCreate(storedType);
             var invoker = new Invoker<TParam, TReturn>(
                 flowType,
                 storedType,
                 inner,
                 invocationHelper,
                 settingsWithDefaults.UnhandledExceptionHandler,
                 ClusterInfo.ReplicaId,
-                _flowsManagers.GetOrCreate(storedType)
+                flowsManager
             );
+            flowsManager.SetRestarter(invoker);
 
             WatchDogsFactory.CreateAndStart(
                 flowType,
@@ -341,15 +342,17 @@ private ParamlessRegistration RegisterParamless(
                 settings?.ClearChildrenAfterCapture ?? true,
                 _messageClearer
             );
+            var flowsManager = _flowsManagers.GetOrCreate(storedType);
             var invoker = new Invoker<Unit, Unit>(
                 flowType,
                 storedType,
                 inner,
                 invocationHelper,
                 settingsWithDefaults.UnhandledExceptionHandler,
                 ClusterInfo.ReplicaId,
-                _flowsManagers.GetOrCreate(storedType)
+                flowsManager
             );
+            flowsManager.SetRestarter(invoker);
 
             WatchDogsFactory.CreateAndStart(
                 flowType,
@@ -425,15 +428,17 @@ public ActionRegistration<TParam> RegisterAction<TParam>(
                 settings?.ClearChildrenAfterCapture ?? true,
                 _messageClearer
             );
+            var flowsManager = _flowsManagers.GetOrCreate(storedType);
             var rActionInvoker = new Invoker<TParam, Unit>(
                 flowType,
                 storedType,
                 inner,
                 invocationHelper,
                 settingsWithDefaults.UnhandledExceptionHandler,
                 ClusterInfo.ReplicaId,
-                _flowsManagers.GetOrCreate(storedType)
+                flowsManager
             );
+            flowsManager.SetRestarter(rActionInvoker);
 
             WatchDogsFactory.CreateAndStart(
                 flowType,