From 4ab6edb8659f31e4470775bbd663fc1094341ce0 Mon Sep 17 00:00:00 2001
From: Sipke Schoorstra <sipkeschoorstra@outlook.com>
Date: Wed, 20 Apr 2022 14:37:20 +0200
Subject: [PATCH] Add behavior that schedules child completion callbacks

---
 src/core/Elsa.Core/Behaviors/BreakBehavior.cs |  6 ++--
 .../ScheduledChildCallbackBehavior.cs         | 34 +++++++++++++++++++
 src/core/Elsa.Core/Models/Activity.cs         | 17 ++--------
 3 files changed, 39 insertions(+), 18 deletions(-)
 create mode 100644 src/core/Elsa.Core/Behaviors/ScheduledChildCallbackBehavior.cs

diff --git a/src/core/Elsa.Core/Behaviors/BreakBehavior.cs b/src/core/Elsa.Core/Behaviors/BreakBehavior.cs
index b44d547cc..bf84e61e2 100644
--- a/src/core/Elsa.Core/Behaviors/BreakBehavior.cs
+++ b/src/core/Elsa.Core/Behaviors/BreakBehavior.cs
@@ -11,15 +11,15 @@ public class BreakBehavior : IBehavior
 {
     public async ValueTask HandleSignalAsync(object signal, SignalContext context)
     {
-        if (signal is not BreakSignal breakSignal)
+        if (signal is not BreakSignal)
             return;
 
-        await OnBreakAsync(breakSignal, context);
+        await OnBreakAsync(context);
     }
 
     public ValueTask ExecuteAsync(ActivityExecutionContext context) => ValueTask.CompletedTask;
 
-    private async ValueTask OnBreakAsync(BreakSignal signal, SignalContext context)
+    private async ValueTask OnBreakAsync(SignalContext context)
     {
         // Prevent bubbling.
         context.StopPropagation();
diff --git a/src/core/Elsa.Core/Behaviors/ScheduledChildCallbackBehavior.cs b/src/core/Elsa.Core/Behaviors/ScheduledChildCallbackBehavior.cs
new file mode 100644
index 000000000..ba5731777
--- /dev/null
+++ b/src/core/Elsa.Core/Behaviors/ScheduledChildCallbackBehavior.cs
@@ -0,0 +1,34 @@
+using Elsa.Contracts;
+using Elsa.Models;
+using Elsa.Signals;
+
+namespace Elsa.Behaviors;
+
+/// <summary>
+/// Implements a "break" behavior that handles the <see cref="BreakSignal"/> signal.
+/// </summary>
+public class ScheduledChildCallbackBehavior : IBehavior
+{
+    public async ValueTask HandleSignalAsync(object signal, SignalContext context)
+    {
+        if (signal is not ActivityCompleted)
+            return;
+
+        await OnChildActivityCompletedAsync(context);
+    }
+
+    public ValueTask ExecuteAsync(ActivityExecutionContext context) => ValueTask.CompletedTask;
+
+    private async ValueTask OnChildActivityCompletedAsync(SignalContext context)
+    {
+        var activityExecutionContext = context.ActivityExecutionContext;
+        var childActivityExecutionContext = context.SourceActivityExecutionContext;
+        var childActivity = childActivityExecutionContext.Activity;
+        var callbackEntry = activityExecutionContext.WorkflowExecutionContext.PopCompletionCallback(activityExecutionContext, childActivity);
+
+        if (callbackEntry == null)
+            return;
+
+        await callbackEntry(activityExecutionContext, childActivityExecutionContext);
+    }
+}
\ No newline at end of file
diff --git a/src/core/Elsa.Core/Models/Activity.cs b/src/core/Elsa.Core/Models/Activity.cs
index 0a39ff4fb..5b2dfd6c0 100644
--- a/src/core/Elsa.Core/Models/Activity.cs
+++ b/src/core/Elsa.Core/Models/Activity.cs
@@ -1,7 +1,7 @@
 using System.Linq.Expressions;
+using Elsa.Behaviors;
 using Elsa.Contracts;
 using Elsa.Helpers;
-using Elsa.Signals;
 
 namespace Elsa.Models;
 
@@ -12,7 +12,7 @@ public abstract class Activity : IActivity, ISignalHandler
     protected Activity()
     {
         TypeName = TypeNameHelper.GenerateTypeName(GetType());
-        OnSignalReceived<ActivityCompleted>(OnChildActivityCompletedAsync);
+        Behaviors.Add<ScheduledChildCallbackBehavior>();
     }
 
     protected Activity(string activityType) : this()
@@ -76,19 +76,6 @@ public abstract class Activity : IActivity, ISignalHandler
         });
     }
     
-    protected virtual async ValueTask OnChildActivityCompletedAsync(ActivityCompleted signal, SignalContext context)
-    {
-        var activityExecutionContext = context.ActivityExecutionContext;
-        var childActivityExecutionContext = context.SourceActivityExecutionContext;
-        var childActivity = childActivityExecutionContext.Activity;
-        var callbackEntry = activityExecutionContext.WorkflowExecutionContext.PopCompletionCallback(activityExecutionContext, childActivity);
-
-        if (callbackEntry == null)
-            return;
-
-        await callbackEntry(activityExecutionContext, childActivityExecutionContext);
-    }
-
     async ValueTask IActivity.ExecuteAsync(ActivityExecutionContext context)
     {
         await ExecuteAsync(context);