Extract FlowState and centralize signal management (#123)

Author: stidsborg

Core/Cleipnir.ResilientFunctions/CoreRuntime/FlowState.cs

@@ -0,0 +1,65 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Cleipnir.ResilientFunctions.Helpers;
+using Cleipnir.ResilientFunctions.Queuing;
+using Cleipnir.ResilientFunctions.Storage;
+
+namespace Cleipnir.ResilientFunctions.CoreRuntime;
+
+internal class FlowState
+{
+    private readonly Lock _lock = new();
+
+    public StoredId Id { get; }
+    public Action Suspend { get; }
+    public QueueManager QueueManager { get; }
+    public int Threads { get; private set; }
+    public int SuspendedThreads { get; private set; }
+    public FlowTimeouts Timeouts { get; }
+    public AsyncSignal Signal { get; } = new();
+
+    public FlowState(
+        StoredId id,
+        Action suspend,
+        QueueManager queueManager,
+        int threads,
+        int suspendedThreads,
+        FlowTimeouts timeouts)
+    {
+        Id = id;
+        Suspend = suspend;
+        QueueManager = queueManager;
+        Threads = threads;
+        SuspendedThreads = suspendedThreads;
+        Timeouts = timeouts;
+    }
+
+    public void NewThreadStarted()
+    {
+        lock (_lock)
+            Threads++;
+    }
+
+    public void ThreadCompleted()
+    {
+        lock (_lock)
+            Threads--;
+    }
+
+    public async Task ThreadSuspended()
+    {
+        lock (_lock)
+            SuspendedThreads++;
+
+        await Signal.Wait();
+    }
+
+    public void Interrupt()
+    {
+        lock (_lock)
+            SuspendedThreads = 0;
+
+        Signal.Fire();
+    }
+}

Core/Cleipnir.ResilientFunctions/CoreRuntime/FlowsManager.cs

@@ -9,11 +9,9 @@
 
 namespace Cleipnir.ResilientFunctions.CoreRuntime;
 
-public record FlowStatus(StoredId Id, Action Suspend, QueueManager QueueManager, int Threads, int SuspendedThreads, FlowTimeouts Timeouts);
-
 public class FlowsManager : IDisposable
 {
-    private readonly Dictionary<StoredId, FlowStatus> _dict = new();
+    private readonly Dictionary<StoredId, FlowState> _dict = new();
     private readonly Lock _lock = new();
     private readonly UtcNow _utcNow;
     private volatile bool _disposed;
@@ -28,14 +26,14 @@ private async Task TimeoutCheckLoop()
     {
         while (!_disposed)
         {
-            var expiredStatuses = new List<FlowStatus>();
+            var expiredStates = new List<FlowState>();
             var now = _utcNow();
             lock (_lock)
                 foreach (var (_, status) in _dict)
                     if (status.Timeouts.HasExpiredTimeouts(now))
-                        expiredStatuses.Add(status);
+                        expiredStates.Add(status);
 
-            foreach (var status in expiredStatuses)
+            foreach (var status in expiredStates)
                 status.Timeouts.SignalExpiredTimeouts(now);
 
             await Task.Delay(10);
@@ -47,7 +45,7 @@ private async Task TimeoutCheckLoop()
     public void AddFlow(StoredId id, Action suspend, QueueManager queueManager, FlowTimeouts timeouts)
     {
         lock (_lock)
-            _dict[id] = new FlowStatus(id, suspend, queueManager, Threads: 1, SuspendedThreads: 0, timeouts);
+            _dict[id] = new FlowState(id, suspend, queueManager, threads: 1, suspendedThreads: 0, timeouts);
     }
 
     public void RemoveFlow(StoredId id)
@@ -62,62 +60,53 @@ public void Interrupt(IEnumerable<StoredId> ids)
         {
             foreach (var id in ids)
             {
-                if (!_dict.ContainsKey(id))
+                if (!_dict.TryGetValue(id, out var flowState))
                     continue;
-                
-                var queueManager = _dict[id].QueueManager;
-                InterruptThreads(id);
-                Task.Run(() => queueManager.FetchMessagesOnce());
+
+                flowState.Interrupt();
+                Task.Run(() => flowState.QueueManager.FetchMessagesOnce());
             }
+
         }
     }
 
     public void StartThread(StoredId id)
     {
         lock (_lock)
-        {
-            if (!_dict.ContainsKey(id))
-                return;
-
-            var status = _dict[id];
-            _dict[id] = status with { Threads = status.Threads + 1 };
-        }
+            if (_dict.TryGetValue(id, out var flowState))
+                flowState.NewThreadStarted();
     }
 
     public void CompleteThread(StoredId id)
     {
         lock (_lock)
-        {
-            if (!_dict.ContainsKey(id))
-                return;
-
-            var status = _dict[id];
-            _dict[id] = status with { Threads = status.Threads - 1 };
-        }
+            if (_dict.TryGetValue(id, out var flowState))
+                flowState.ThreadCompleted();
     }
 
-    public void InterruptThreads(StoredId id)
+    public async Task SuspendThread(StoredId id)
     {
+        FlowState? flowState;
         lock (_lock)
-        {
-            if (!_dict.ContainsKey(id))
-                return;
+            _dict.TryGetValue(id, out flowState);
 
-            var status = _dict[id];
-            _dict[id] = status with { SuspendedThreads = 0 };
-        }
+        if (flowState != null)
+            await flowState.ThreadSuspended();
     }
 
-    public void SuspendThread(StoredId id)
+    public Task GetInterruptedSignal(StoredId id)
     {
         lock (_lock)
-        {
-            if (!_dict.ContainsKey(id))
-                return;
+            return _dict.TryGetValue(id, out var flowState)
+                ? flowState.Signal.Wait()
+                : Task.CompletedTask;
+    }
 
-            var status = _dict[id];
-            _dict[id] = status with { SuspendedThreads = status.SuspendedThreads + 1 };
-        }
+    public void SignalInterrupt(StoredId id)
+    {
+        lock (_lock)
+            if (_dict.TryGetValue(id, out var flowState))
+                flowState.Signal.Fire();
     }
 
     [DoesNotReturn]

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

@@ -56,7 +56,7 @@ async Task Inner()
 
             var maxWait = suspend ? TimeSpan.Zero : (TimeSpan?)null;
             await Effect.FlowTimeouts.AddTimeout(timeoutId, expiry.ToDateTime(), maxWait);
-        
+            
             await Effect.Upsert(timeoutId, -1L, alias, flush: false);
             Effect.FlowTimeouts.RemoveTimeout(timeoutId);
         }

Core/Cleipnir.ResilientFunctions/Helpers/AsyncSignal.cs

@@ -0,0 +1,31 @@
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Cleipnir.ResilientFunctions.Helpers;
+
+public class AsyncSignal
+{
+    private TaskCompletionSource _tcs = new();
+    private readonly Lock _lock = new();
+    
+    public async Task Wait()
+    {
+        Task task;
+        lock (_lock)
+            task = _tcs.Task;
+
+        await task;
+    }
+
+    public void Fire()
+    {
+        TaskCompletionSource tcs;
+        lock (_lock)
+        {
+            tcs = _tcs;
+            _tcs = new();
+        }
+        
+        tcs.SetResult();
+    }
+}

Core/Cleipnir.ResilientFunctions/Queuing/QueueManager.cs

@@ -42,7 +42,6 @@ public class QueueManager(
     private volatile bool _disposed;
 
     private volatile Exception? _thrownException = null;
-    private TaskCompletionSource _pulse = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
     private async Task Initialize()
     {
@@ -152,24 +151,28 @@ public async Task FetchMessagesOnce()
         }
         finally
         {
-            PulseAll();
+            flowsManager.SignalInterrupt(storedId);
             _semaphoreSlim.Release();
         }
     }
 
     private (MessageData? Matched, int PositionToRemoveIndex, Task PulseTask) TryTakeMessage(MessagePredicate predicate)
     {
+        var interruptedSignal = flowsManager.GetInterruptedSignal(storedId);
+        
         lock (_lock)
         {
             for (var i = 0; i < _toDeliver.Count; i++)
                 if (predicate(_toDeliver[i].Envelope))
                 {
                     var matched = _toDeliver[i];
                     _toDeliver.RemoveAt(i);
-                    return (matched, _nextToRemoveIndex++, _pulse.Task);
+                    var positionToRemoveIndex = _nextToRemoveIndex++;
+                    effect.FlushlessUpsert(_toRemoveNextIndex, _nextToRemoveIndex, alias: null);
+                    return (matched, positionToRemoveIndex, interruptedSignal);
                 }
 
-            return (null, 0, _pulse.Task);
+            return (null, 0, interruptedSignal);
         }
     }
 
@@ -221,17 +224,6 @@ public async Task AfterFlush()
         }
     }
 
-    private void PulseAll()
-    {
-        TaskCompletionSource oldPulse;
-        lock (_lock)
-        {
-            oldPulse = _pulse;
-            _pulse = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-        }
-        oldPulse.TrySetResult();
-    }
-
     public async Task<Envelope?> Subscribe(
         MessagePredicate predicate,
         DateTime? timeout,
@@ -266,7 +258,6 @@ private void PulseAll()
                 effect.FlushlessUpserts(
                     new List<EffectResult>(
                     [
-                        new EffectResult(_toRemoveNextIndex, positionToRemoveIndex, Alias: null),
                         new EffectResult(toRemoveId, matched.Position, Alias: null),
                         new EffectResult(messageId, matched.MessageContentBytes, Alias: null),
                         new EffectResult(messageTypeId, matched.MessageTypeBytes, Alias: null),