Add CancellationToken to IfState condition

PaGrom · Sep 6, 2020 · 756e6da · 756e6da
1 parent a0f1f3e
commit 756e6da
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/Telegrom.StateMachine/Builder/GeneratedState.cs b/Telegrom.StateMachine/Builder/GeneratedState.cs
@@ -34,7 +34,7 @@ public async Task OnEnter(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterOnEnter);
+            await MoveNextAsync(NextStateKind.AfterOnEnter, cancellationToken);
         }
 
         public async Task Handle(CancellationToken cancellationToken)
@@ -43,7 +43,7 @@ public async Task Handle(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterHandle);
+            await MoveNextAsync(NextStateKind.AfterHandle, cancellationToken);
         }
 
         public async Task OnExit(CancellationToken cancellationToken)
@@ -52,10 +52,10 @@ public async Task OnExit(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterOnExit);
+            await MoveNextAsync(NextStateKind.AfterOnExit, cancellationToken);
         }
 
-        private async Task MoveNextAsync(NextStateKind nextStateKind)
+        private async Task MoveNextAsync(NextStateKind nextStateKind, CancellationToken cancellationToken)
         {
             if (_stateNode.NextStateKind != nextStateKind)
             {
@@ -66,7 +66,7 @@ private async Task MoveNextAsync(NextStateKind nextStateKind)
 
             foreach (var ifState in _stateNode.IfStates)
             {
-                if (await ifState.Condition(_stateContext))
+                if (await ifState.Condition(_stateContext, cancellationToken))
                 {
                     _stateContext.StateMachineContext.MoveTo(ifState.StateNode.StateName);
                     break;

diff --git a/Telegrom.StateMachine/Builder/IfState.cs b/Telegrom.StateMachine/Builder/IfState.cs
@@ -1,14 +1,15 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Telegrom.StateMachine.Builder
 {
     public sealed class IfState
     {
-        internal Func<IStateContext, Task<bool>> Condition { get; }
+        internal Func<IStateContext, CancellationToken, Task<bool>> Condition { get; }
         internal StateNode StateNode { get; }
 
-        public IfState(Func<IStateContext, Task<bool>> condition, Type stateType, string stateName = null)
+        public IfState(Func<IStateContext, CancellationToken, Task<bool>> condition, Type stateType, string stateName = null)
         {
             if (!typeof(IState).IsAssignableFrom(stateType))
             {
@@ -25,17 +26,17 @@ public IfState(Func<IStateContext, Task<bool>> condition, Type stateType, string
         }
 
         public IfState(Func<IStateContext, bool> condition, Type stateType, string stateName = null)
-            : this(context => new Task<bool>(() => condition(context)), stateType, stateName)
+            : this((context, ctk) => Task.Run(() => condition(context), ctk), stateType, stateName)
         { }
 
-        public IfState(Func<IStateContext, Task<bool>> condition, StateNode stateNode)
+        public IfState(Func<IStateContext, CancellationToken, Task<bool>> condition, StateNode stateNode)
         {
             Condition = condition;
             StateNode = stateNode;
         }
 
         public IfState(Func<IStateContext, bool> condition, StateNode stateNode)
-            : this(context => new Task<bool>(() => condition(context)), stateNode)
+            : this((context, ctk) => Task.Run(() => condition(context), ctk), stateNode)
         { }
     }
 }