From 6a91eb5130ae7b0f144ee00914d20e3ece137d1f Mon Sep 17 00:00:00 2001
From: Sverre Winkelmans <sverre.winkelmans@gmail.com>
Date: Sat, 14 Dec 2024 12:19:21 +0800
Subject: [PATCH] Allow to define cleanup strategy for workflow instances

---
 .../DeleteWorkflowInstanceStrategy.cs         | 20 +++++++++++++++++++
 .../Feature/RetentionFeature.cs               |  2 ++
 src/modules/Elsa.Retention/Jobs/CleanupJob.cs | 17 ++++++++++++----
 3 files changed, 35 insertions(+), 4 deletions(-)
 create mode 100644 src/modules/Elsa.Retention/CleanupStrategies/DeleteWorkflowInstanceStrategy.cs

diff --git a/src/modules/Elsa.Retention/CleanupStrategies/DeleteWorkflowInstanceStrategy.cs b/src/modules/Elsa.Retention/CleanupStrategies/DeleteWorkflowInstanceStrategy.cs
new file mode 100644
index 000000000..da3b272e5
--- /dev/null
+++ b/src/modules/Elsa.Retention/CleanupStrategies/DeleteWorkflowInstanceStrategy.cs
@@ -0,0 +1,20 @@
+using Elsa.Retention.Contracts;
+using Elsa.Workflows.Management;
+using Elsa.Workflows.Management.Entities;
+using Elsa.Workflows.Management.Filters;
+
+namespace Elsa.Retention.CleanupStrategies;
+
+/// <summary>
+///     Deletes the workflow instance.
+/// </summary>
+public class DeleteWorkflowInstanceStrategy(IWorkflowInstanceStore store) : IDeletionCleanupStrategy<WorkflowInstance>
+{
+    public async Task Cleanup(ICollection<WorkflowInstance> collection)
+    {
+        await store.DeleteAsync(new WorkflowInstanceFilter
+        {
+            Ids = collection.Select(x => x.Id).ToArray()
+        });
+    }
+}
\ No newline at end of file
diff --git a/src/modules/Elsa.Retention/Feature/RetentionFeature.cs b/src/modules/Elsa.Retention/Feature/RetentionFeature.cs
index f29922909..cd3783d83 100644
--- a/src/modules/Elsa.Retention/Feature/RetentionFeature.cs
+++ b/src/modules/Elsa.Retention/Feature/RetentionFeature.cs
@@ -7,6 +7,7 @@ using Elsa.Retention.Contracts;
 using Elsa.Retention.Extensions;
 using Elsa.Retention.Jobs;
 using Elsa.Retention.Options;
+using Elsa.Workflows.Management.Entities;
 using Elsa.Workflows.Runtime.Entities;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -44,6 +45,7 @@ public class RetentionFeature : FeatureBase
         Services.AddScoped<IDeletionCleanupStrategy<StoredBookmark>, DeleteBookmarkStrategy>();
         Services.AddScoped<IDeletionCleanupStrategy<ActivityExecutionRecord>, DeleteActivityExecutionRecordStrategy>();
         Services.AddScoped<IDeletionCleanupStrategy<WorkflowExecutionLogRecord>, DeleteWorkflowExecutionRecordStrategy>();
+        Services.AddScoped<IDeletionCleanupStrategy<WorkflowInstance>, DeleteWorkflowInstanceStrategy>();
 
         Services.AddScoped<IRelatedEntityCollector, BookmarkCollector>();
         Services.AddScoped<IRelatedEntityCollector, ActivityExecutionRecordCollector>();
diff --git a/src/modules/Elsa.Retention/Jobs/CleanupJob.cs b/src/modules/Elsa.Retention/Jobs/CleanupJob.cs
index 701ab32f3..330e9d41e 100644
--- a/src/modules/Elsa.Retention/Jobs/CleanupJob.cs
+++ b/src/modules/Elsa.Retention/Jobs/CleanupJob.cs
@@ -31,6 +31,7 @@ public class CleanupJob(
     /// <param name="cancellationToken"></param>
     public async Task ExecuteAsync(CancellationToken cancellationToken = default)
     {
+        Console.WriteLine(DateTime.Now.ToLongTimeString());
         var collectors = GetServices(typeof(IRelatedEntityCollector), typeof(IRelatedEntityCollector<>));
         var deletedWorkflowInstances = 0L;
         
@@ -51,6 +52,10 @@ public class CleanupJob(
                 foreach (var collectorService in collectors)
                 {
                     var cleanupStrategyConcreteType = policy.CleanupStrategy.MakeGenericType(collectorService.Key);
+                    
+                    if(cleanupStrategyConcreteType == typeof(WorkflowInstance))
+                        continue;
+                    
                     var collector = collectorService.Value as IRelatedEntityCollector;
                     var cleanupService = serviceProvider.GetService(cleanupStrategyConcreteType) as ICleanupStrategy;
 
@@ -71,11 +76,15 @@ public class CleanupJob(
                         await cleanupService.Cleanup(entities);
                     }
                 }
+                
+                var cleanupWorkflowInstances = policy.CleanupStrategy.MakeGenericType(typeof(WorkflowInstance));
+                var workflowInstanceCleaner = serviceProvider.GetService(cleanupWorkflowInstances) as ICleanupStrategy<WorkflowInstance>;
 
-                deletedWorkflowInstances += await workflowInstanceStore.DeleteAsync(new WorkflowInstanceFilter
-                {
-                    Ids = page.Items.Select(x => x.Id).ToArray()
-                }, cancellationToken);
+                if (workflowInstanceCleaner == null)
+                    throw new Exception($"{policy.CleanupStrategy} has no strategy to clean WorkflowInstances");
+                
+                await workflowInstanceCleaner.Cleanup(page.Items);
+                deletedWorkflowInstances += page.Items.Count;
 
                 if (page.TotalCount <= page.Items.Count + pageArgs.Offset)
                 {