From dc46bfe1bb8d33a3e7a3907bce776f6bc407c70b Mon Sep 17 00:00:00 2001
From: Sipke Schoorstra <sipkeschoorstra@outlook.com>
Date: Sat, 25 Jun 2022 11:55:52 +0200
Subject: [PATCH] Update WorkflowStateStorageDriver.cs

---
 .../Implementations/WorkflowStateStorageDriver.cs  | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/modules/Elsa.Workflows.Core/Implementations/WorkflowStateStorageDriver.cs b/src/modules/Elsa.Workflows.Core/Implementations/WorkflowStateStorageDriver.cs
index 0713aaddd..45386255b 100644
--- a/src/modules/Elsa.Workflows.Core/Implementations/WorkflowStateStorageDriver.cs
+++ b/src/modules/Elsa.Workflows.Core/Implementations/WorkflowStateStorageDriver.cs
@@ -13,8 +13,7 @@ public class WorkflowStateStorageDriver : IStorageDriver
 
     public ValueTask WriteAsync(string id, object value, DataDriveContext context)
     {
-        var dictionary = GetVariablesDictionary(context);
-        dictionary[id] = value;
+        UpdateVariablesDictionary(context, dictionary => dictionary[id] = value);
         return ValueTask.CompletedTask;
     }
 
@@ -27,10 +26,17 @@ public class WorkflowStateStorageDriver : IStorageDriver
 
     public ValueTask DeleteAsync(string id, DataDriveContext context)
     {
-        var dictionary = GetVariablesDictionary(context);
-        dictionary.Remove(id);
+        UpdateVariablesDictionary(context, dictionary => dictionary.Remove(id));
         return ValueTask.CompletedTask;
     }
 
     private IDictionary<string, object> GetVariablesDictionary(DataDriveContext context) => context.WorkflowState.Properties.GetOrAdd(VariablesDictionaryStateKey, () => new Dictionary<string, object>());
+    private void SetVariablesDictionary(DataDriveContext context, IDictionary<string, object> dictionary) => context.WorkflowState.Properties[VariablesDictionaryStateKey] = dictionary;
+
+    private void UpdateVariablesDictionary(DataDriveContext context, Action<IDictionary<string, object>> update)
+    {
+        var dictionary = GetVariablesDictionary(context);
+        update(dictionary);
+        SetVariablesDictionary(context, dictionary);
+    }
 }
\ No newline at end of file