Add grain-side status compare-and-set helper (#70)

* Add grain-side status compare-and-set helper

* Update ManagedCode.Communication.Orleans/Grains/CommandIdempotencyGrain.cs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: KSemenenko

ManagedCode.Communication.Orleans/Grains/CommandIdempotencyGrain.cs

@@ -27,27 +27,26 @@ public Task<CommandExecutionStatus> GetStatusAsync()
     public async Task<bool> TryStartProcessingAsync()
     {
         // Reject concurrent executions
-        if (state.State.Status is CommandExecutionStatus.InProgress or CommandExecutionStatus.Processing)
+        switch (state.State.Status)
         {
-            return false;
-        }
-
-        if (state.State.Status is CommandExecutionStatus.Completed)
-        {
-            return false;
-        }
-
-        // Allow retries after failures by clearing the previous outcome
-        if (state.State.Status is CommandExecutionStatus.Failed)
-        {
-            state.State.Result = null;
-            state.State.ErrorMessage = null;
-            state.State.CompletedAt = null;
-            state.State.FailedAt = null;
-        }
-        else if (state.State.Status is not CommandExecutionStatus.NotFound and not CommandExecutionStatus.NotStarted)
-        {
-            return false;
+            case CommandExecutionStatus.InProgress:
+            case CommandExecutionStatus.Processing:
+            case CommandExecutionStatus.Completed:
+                return false;
+
+            case CommandExecutionStatus.Failed:
+                state.State.Result = null;
+                state.State.ErrorMessage = null;
+                state.State.CompletedAt = null;
+                state.State.FailedAt = null;
+                break;
+
+            case CommandExecutionStatus.NotFound:
+            case CommandExecutionStatus.NotStarted:
+                break;
+
+            default:
+                return false;
         }
 
         state.State.Status = CommandExecutionStatus.Processing;
@@ -58,6 +57,47 @@ public async Task<bool> TryStartProcessingAsync()
         return true;
     }
 
+    public async Task<bool> TrySetStatusAsync(CommandExecutionStatus expectedStatus, CommandExecutionStatus newStatus)
+    {
+        if (state.State.Status != expectedStatus)
+        {
+            return false;
+        }
+
+        switch (newStatus)
+        {
+            case CommandExecutionStatus.InProgress:
+            case CommandExecutionStatus.Processing:
+                return await TryStartProcessingAsync();
+
+            case CommandExecutionStatus.Completed:
+                await MarkCompletedAsync(state.State.Result);
+                return true;
+
+            case CommandExecutionStatus.Failed:
+                await MarkFailedAsync(state.State.ErrorMessage ?? "Status set to failed");
+                return true;
+
+            case CommandExecutionStatus.NotFound:
+                await ClearAsync();
+                return true;
+
+            case CommandExecutionStatus.NotStarted:
+                state.State.Status = CommandExecutionStatus.NotStarted;
+                state.State.Result = null;
+                state.State.ErrorMessage = null;
+                state.State.StartedAt = null;
+                state.State.CompletedAt = null;
+                state.State.FailedAt = null;
+                state.State.ExpiresAt = null;
+                await state.WriteStateAsync();
+                return true;
+
+            default:
+                return false;
+        }
+    }
+
     public async Task MarkCompletedAsync<TResult>(TResult result)
     {
         state.State.Status = CommandExecutionStatus.Completed;

ManagedCode.Communication.Orleans/Grains/ICommandIdempotencyGrain.cs

@@ -42,4 +42,12 @@ public interface ICommandIdempotencyGrain : IGrainWithStringKey
     /// Clears the command state from the grain.
     /// </summary>
     Task ClearAsync();
-}
+
+    /// <summary>
+    /// Attempts to transition the command to a new status when the current status matches the expected value.
+    /// </summary>
+    /// <param name="expectedStatus">The status the caller believes the command currently has.</param>
+    /// <param name="newStatus">The desired status to transition to.</param>
+    /// <returns><c>true</c> when the transition succeeds, otherwise <c>false</c>.</returns>
+    Task<bool> TrySetStatusAsync(CommandExecutionStatus expectedStatus, CommandExecutionStatus newStatus);
+}

ManagedCode.Communication.Orleans/Stores/OrleansCommandIdempotencyStore.cs

@@ -88,14 +88,12 @@ public async Task RemoveCommandAsync(string commandId, CancellationToken cancell
     public async Task<bool> TrySetCommandStatusAsync(string commandId, CommandExecutionStatus expectedStatus, CommandExecutionStatus newStatus, CancellationToken cancellationToken = default)
     {
         var grain = _grainFactory.GetGrain<ICommandIdempotencyGrain>(commandId);
-        var currentStatus = await grain.GetStatusAsync();
-        
-        if (currentStatus == expectedStatus)
+
+        if (await grain.TrySetStatusAsync(expectedStatus, newStatus))
         {
-            await SetCommandStatusAsync(commandId, newStatus, cancellationToken);
             return true;
         }
-        
+
         return false;
     }
 

ManagedCode.Communication/Commands/Stores/MemoryCacheCommandIdempotencyStore.cs

@@ -107,9 +107,9 @@ private void ReleaseLockReference(string commandId, CommandLock commandLock)
     {
         if (Interlocked.Decrement(ref commandLock.RefCount) == 0)
         {
-            if (_commandLocks.TryRemove(commandId, out var existingLock) && !ReferenceEquals(existingLock, commandLock))
+            if (_commandLocks.TryGetValue(commandId, out var existingLock) && ReferenceEquals(existingLock, commandLock))
             {
-                _commandLocks.TryAdd(commandId, existingLock);
+                _commandLocks.TryRemove(new KeyValuePair<string, CommandLock>(commandId, commandLock));
             }
 
             commandLock.Semaphore.Dispose();