From 90651b4735b7e43c4d4e8a19441fe9518f27fd9a Mon Sep 17 00:00:00 2001
From: Raymond den Haan <raymond.den.haan@nexxbiz.io>
Date: Mon, 13 May 2024 12:52:12 +0200
Subject: [PATCH] Updated local handling of workflow definition events for
 updating activity registry

---
 .../Handlers/RefreshActivityRegistry.cs       | 60 ++++++++++++-------
 1 file changed, 39 insertions(+), 21 deletions(-)

diff --git a/src/modules/Elsa.Workflows.Management/Handlers/RefreshActivityRegistry.cs b/src/modules/Elsa.Workflows.Management/Handlers/RefreshActivityRegistry.cs
index b558a388c..a8f0e0a79 100644
--- a/src/modules/Elsa.Workflows.Management/Handlers/RefreshActivityRegistry.cs
+++ b/src/modules/Elsa.Workflows.Management/Handlers/RefreshActivityRegistry.cs
@@ -14,7 +14,6 @@ namespace Elsa.Workflows.Management.Handlers;
 [PublicAPI]
 public class RefreshActivityRegistry(IActivityRegistryPopulator activityRegistryPopulator) :
     INotificationHandler<WorkflowDefinitionPublished>,
-    INotificationHandler<WorkflowDefinitionRetracted>,
     INotificationHandler<WorkflowDefinitionDeleted>,
     INotificationHandler<WorkflowDefinitionsDeleted>,
     INotificationHandler<WorkflowDefinitionCreated>,
@@ -22,31 +21,50 @@ public class RefreshActivityRegistry(IActivityRegistryPopulator activityRegistry
     INotificationHandler<WorkflowDefinitionVersionsDeleted>
 {
     /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionPublished notification, CancellationToken cancellationToken)
+    public Task HandleAsync(WorkflowDefinitionPublished notification, CancellationToken cancellationToken)
     {
-        await activityRegistryPopulator.AddToRegistry(typeof(WorkflowDefinitionActivityProvider), notification.WorkflowDefinition.Id);
+        return activityRegistryPopulator.AddToRegistry(typeof(WorkflowDefinitionActivityProvider), notification.WorkflowDefinition.Id, cancellationToken);
     }
 
     /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionRetracted notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
+    public Task HandleAsync(WorkflowDefinitionDeleted notification, CancellationToken cancellationToken)
+    { 
+        activityRegistryPopulator.RemoveDefinitionFromRegistry(typeof(WorkflowDefinitionActivityProvider), notification.DefinitionId, cancellationToken);
+        return Task.CompletedTask;
+    }
+    
     /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionDeleted notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
-    /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionsDeleted notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
-    /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionCreated notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
-    /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionVersionDeleted notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
-    /// <inheritdoc />
-    public async Task HandleAsync(WorkflowDefinitionVersionsDeleted notification, CancellationToken cancellationToken) => await RefreshAsync(cancellationToken);
-
-    private async Task RefreshAsync(CancellationToken cancellationToken)
+    public Task HandleAsync(WorkflowDefinitionsDeleted notification, CancellationToken cancellationToken)
     {
-        await activityRegistryPopulator.PopulateRegistryAsync(typeof(WorkflowDefinitionActivityProvider), cancellationToken);
+        foreach (string id in notification.DefinitionIds)
+        {
+            activityRegistryPopulator.RemoveDefinitionFromRegistry(typeof(WorkflowDefinitionActivityProvider), id, cancellationToken);
+        }
+
+        return Task.CompletedTask;
+    }
+
+    /// <inheritdoc />
+    public Task HandleAsync(WorkflowDefinitionCreated notification, CancellationToken cancellationToken)
+    {
+        return activityRegistryPopulator.AddToRegistry(typeof(WorkflowDefinitionActivityProvider), notification.WorkflowDefinition.Id, cancellationToken);
+    }
+
+    /// <inheritdoc />
+    public Task HandleAsync(WorkflowDefinitionVersionDeleted notification, CancellationToken cancellationToken)
+    { 
+        activityRegistryPopulator.RemoveDefinitionFromRegistry(typeof(WorkflowDefinitionActivityProvider), notification.WorkflowDefinition.Id, cancellationToken);
+        return Task.CompletedTask;
+    }
+
+    /// <inheritdoc />
+    public Task HandleAsync(WorkflowDefinitionVersionsDeleted notification, CancellationToken cancellationToken)
+    {
+        foreach (string id in notification.Ids)
+        {
+            activityRegistryPopulator.RemoveDefinitionVersionFromRegistry(typeof(WorkflowDefinitionActivityProvider), id, cancellationToken);
+        }
+
+        return Task.CompletedTask;
     }
 }
\ No newline at end of file