From 735861365324be88ddeea34cb7aeb65941da15d5 Mon Sep 17 00:00:00 2001
From: Sipke Schoorstra <sipkeschoorstra@outlook.com>
Date: Sat, 26 Aug 2023 02:02:02 +0200
Subject: [PATCH] Fix orphaned completion callback handling

---
 src/modules/Elsa.Scheduling/Activities/Delay.cs    | 14 +++++++-------
 .../Contexts/ActivityExecutionContext.cs           |  2 +-
 .../Services/WorkflowExecutionContextMapper.cs     |  7 ++++++-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/modules/Elsa.Scheduling/Activities/Delay.cs b/src/modules/Elsa.Scheduling/Activities/Delay.cs
index 464833214..56dac44cf 100644
--- a/src/modules/Elsa.Scheduling/Activities/Delay.cs
+++ b/src/modules/Elsa.Scheduling/Activities/Delay.cs
@@ -15,7 +15,7 @@ namespace Elsa.Scheduling.Activities;
 /// <summary>
 /// Delay execution for the specified amount of time.
 /// </summary>
-[Activity( "Elsa", "Scheduling", "Delay execution for the specified amount of time.")]
+[Activity("Elsa", "Scheduling", "Delay execution for the specified amount of time.")]
 public class Delay : Activity, IActivityPropertyDefaultValueProvider
 {
     /// <inheritdoc />
@@ -67,7 +67,7 @@ public class Delay : Activity, IActivityPropertyDefaultValueProvider
     [Input(
         Description = "The timespan to delay workflow execution.",
         DefaultValueProvider = typeof(Delay)
-    )] 
+    )]
     public Input<TimeSpan> TimeSpan { get; set; } = default!;
 
     /// <inheritdoc />
@@ -81,35 +81,35 @@ public class Delay : Activity, IActivityPropertyDefaultValueProvider
         context.JournalData.Add("ResumeAt", resumeAt);
         context.CreateBookmark(payload);
     }
-    
+
     /// <summary>
     /// Creates a new <see cref="Delay"/> from the specified number of milliseconds.
     /// </summary>
     public static Delay FromMilliseconds(
         double value,
         [CallerFilePath] string? source = default, [CallerLineNumber] int? line = default) => new(System.TimeSpan.FromMilliseconds(value), source, line);
-    
+
     /// <summary>
     /// Creates a new <see cref="Delay"/> from the specified number of seconds.
     /// </summary>
     public static Delay FromSeconds(
         double value,
         [CallerFilePath] string? source = default, [CallerLineNumber] int? line = default) => new(System.TimeSpan.FromSeconds(value), source, line);
-    
+
     /// <summary>
     /// Creates a new <see cref="Delay"/> from the specified number of minutes.
     /// </summary>
     public static Delay FromMinutes(
         double value,
         [CallerFilePath] string? source = default, [CallerLineNumber] int? line = default) => new(System.TimeSpan.FromMinutes(value), source, line);
-    
+
     /// <summary>
     /// Creates a new <see cref="Delay"/> from the specified number of hours.
     /// </summary>
     public static Delay FromHours(
         double value,
         [CallerFilePath] string? source = default, [CallerLineNumber] int? line = default) => new(System.TimeSpan.FromHours(value), source, line);
-    
+
     /// <summary>
     /// Creates a new <see cref="Delay"/> from the specified number of days.
     /// </summary>
diff --git a/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs b/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
index 74f848603..8e8e5a406 100644
--- a/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
+++ b/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
@@ -530,7 +530,7 @@ public class ActivityExecutionContext : IExecutionContext
     /// </summary>
     public void ClearCompletionCallbacks()
     {
-        var entriesToRemove = WorkflowExecutionContext.CompletionCallbacks.Where(x => x.Owner == this || x.Child.Activity == Activity);
+        var entriesToRemove = WorkflowExecutionContext.CompletionCallbacks.Where(x => x.Owner == this);
         WorkflowExecutionContext.RemoveCompletionCallbacks(entriesToRemove);
     }
 
diff --git a/src/modules/Elsa.Workflows.Core/Services/WorkflowExecutionContextMapper.cs b/src/modules/Elsa.Workflows.Core/Services/WorkflowExecutionContextMapper.cs
index ee5f2e82b..2b1df33bf 100644
--- a/src/modules/Elsa.Workflows.Core/Services/WorkflowExecutionContextMapper.cs
+++ b/src/modules/Elsa.Workflows.Core/Services/WorkflowExecutionContextMapper.cs
@@ -65,7 +65,12 @@ public class WorkflowExecutionContextMapper : IWorkflowExecutionContextMapper
         foreach (var completionCallbackEntry in state.CompletionCallbacks)
         {
             var ownerActivityExecutionContext = workflowExecutionContext.ActiveActivityExecutionContexts.First(x => x.Id == completionCallbackEntry.OwnerInstanceId);
-            var childNode = workflowExecutionContext.ActiveActivityExecutionContexts.First(x => x.NodeId == completionCallbackEntry.ChildNodeId).ActivityNode;
+            var childNode = workflowExecutionContext.ActiveActivityExecutionContexts.FirstOrDefault(x => x.NodeId == completionCallbackEntry.ChildNodeId)?.ActivityNode;
+            
+            // If the child node is null, it means the completion callback was registered for an activity instance that has already completed or was canceled. 
+            if(childNode == null)
+                continue;
+            
             var callbackName = completionCallbackEntry.MethodName;
             var callbackDelegate = !string.IsNullOrEmpty(callbackName) ? ownerActivityExecutionContext.Activity.GetActivityCompletionCallback(callbackName) : default;
             workflowExecutionContext.AddCompletionCallback(ownerActivityExecutionContext, childNode, callbackDelegate);