Move suspension into FlowTimeouts.AddTimeout via maxWait parameter

Author: stidsborg

Core/Cleipnir.ResilientFunctions/CoreRuntime/FlowTimeouts.cs

@@ -4,6 +4,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Cleipnir.ResilientFunctions.Domain;
+using Cleipnir.ResilientFunctions.Domain.Exceptions.Commands;
 
 namespace Cleipnir.ResilientFunctions.CoreRuntime;
 
@@ -24,7 +25,7 @@ public DateTime? MinimumTimeout
     private DateTime? GetMinimumTimeout()
         => Timeouts.Values.Count != 0 ? Timeouts.Values.Min(t => t.Item1) : (DateTime?)null;
 
-    public Task AddTimeout(EffectId effectId, DateTime timeout)
+    public async Task AddTimeout(EffectId effectId, DateTime timeout, TimeSpan? maxWait = null)
     {
         TaskCompletionSource tcs;
         lock (_lock)
@@ -33,7 +34,15 @@ public Task AddTimeout(EffectId effectId, DateTime timeout)
             Timeouts[effectId] = Tuple.Create(timeout, tcs);
         }
 
-        return tcs.Task;
+        if (maxWait == null || timeout <= DateTime.UtcNow)
+        {
+            await tcs.Task;
+            return;
+        }
+
+        var completed = await Task.WhenAny(tcs.Task, Task.Delay(maxWait.Value));
+        if (completed != tcs.Task)
+            throw new SuspendInvocationException();
     }
 
     public void RemoveTimeout(EffectId effectId)

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

@@ -54,15 +54,8 @@ async Task Inner()
                 return;
             }
 
-            var timeoutTask = Effect.FlowTimeouts.AddTimeout(timeoutId, expiry.ToDateTime());
-            if (!suspend)
-                await timeoutTask;
-            else
-            {
-                var delay = (expiry.ToDateTime() - _utcNow()).RoundUpToZero();
-                if (delay > TimeSpan.Zero)
-                    await _flowsManager.Suspend(StoredId);
-            }
+            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/Domain/Effect.cs

@@ -24,8 +24,6 @@ public class Effect(EffectResults effectResults, UtcNow utcNow, FlowTimeouts flo
     internal IEnumerable<EffectId> EffectIds => effectResults.EffectIds;
     internal FlowTimeouts FlowTimeouts => flowTimeouts;
 
-    internal Task Suspend() => flowsManager.Suspend(storedId);
-
     internal WorkStatus? GetStatus(int id)
     {
         var effectId = CreateEffectId(id);

Core/Cleipnir.ResilientFunctions/Domain/RetryPolicy.cs

@@ -108,10 +108,7 @@ public async Task<T> Invoke<T>(Func<Task<T>> work, Effect effect, UtcNow utcNow,
         var hasDelayUntil = effect.TryGet<long>(delayUntilId, out var delayUntilValue);
         var delayUntil = hasDelayUntil ? delayUntilValue.ToDateTime() : DateTime.MinValue;
         if (hasDelayUntil && delayUntil > utcNow())
-        {
-            flowTimeouts.AddTimeout(delayUntilId, delayUntil);
-            await effect.Suspend();
-        }
+            await flowTimeouts.AddTimeout(delayUntilId, delayUntil, maxWait: TimeSpan.Zero);
 
         var iterationId = effect.CreateEffectId(1);
         var iteration = await effect.CreateOrGet(iterationId, 0, alias: null, flush: false);
@@ -149,12 +146,9 @@ await effect.Upserts(
                     throw;
 
                 if (delay >= suspendThreshold)
-                {
-                    flowTimeouts.AddTimeout(delayUntilId, delayUntil);
-                    await effect.Suspend();
-                }
-                
-                await Task.Delay(delay);
+                    await flowTimeouts.AddTimeout(delayUntilId, delayUntil, maxWait: TimeSpan.Zero);
+                else
+                    await Task.Delay(delay);
             }            
         }
     }