Fixed postgres saga store concurrency issue (#10085)

Author: PascalSenn

src/Mocha/src/Mocha.EntityFrameworkCore.Postgres/Sagas/PostgresSagaStore.cs

@@ -1,3 +1,4 @@
+using System.Collections.Concurrent;
 using Microsoft.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.Extensions.DependencyInjection;
@@ -13,6 +14,8 @@ internal sealed class PostgresSagaStore(DbContext context, PostgresSagaStoreQuer
     , IDisposable
 {
     private readonly object _lock = new();
+    private readonly ConcurrentDictionary<SagaStateKey, Guid> _versions = new();
+    private Guid? _transactionId = context.Database.CurrentTransaction?.TransactionId;
     private PooledArrayWriter? _arrayWriter;
 
     /// <summary>
@@ -28,7 +31,7 @@ public static PostgresSagaStore Create(Type contextType, string optionsName, ISe
         var dbContext = (DbContext)services.GetRequiredService(contextType);
         var optionsMonitor = services.GetRequiredService<IOptionsMonitor<PostgresSagaStoreOptions>>();
         var options = optionsMonitor.Get(optionsName);
-        var timeProvider = services.GetRequiredService<TimeProvider>();
+        var timeProvider = services.GetService<TimeProvider>() ?? TimeProvider.System;
 
         return new PostgresSagaStore(dbContext, options.Queries, timeProvider);
     }
@@ -55,8 +58,9 @@ public async Task<ISagaTransaction> StartTransactionAsync(CancellationToken canc
     }
 
     /// <summary>
-    /// Persists the saga state using raw SQL against Postgres, inserting a new record or updating
-    /// the existing one with optimistic concurrency control via the version column.
+    /// Persists the saga state using raw SQL against Postgres. States loaded by this store are
+    /// updated with optimistic concurrency control using the version captured at load time. States
+    /// that were not loaded by this store are inserted and conflict if a row already exists.
     /// </summary>
     /// <typeparam name="T">The saga state type derived from <see cref="SagaStateBase"/>.</typeparam>
     /// <param name="saga">The saga definition providing name and serialization metadata.</param>
@@ -67,6 +71,8 @@ public async Task<ISagaTransaction> StartTransactionAsync(CancellationToken canc
     /// </exception>
     public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellationToken) where T : SagaStateBase
     {
+        EnsureCurrentTransactionGeneration();
+
         var connection = (NpgsqlConnection)context.Database.GetDbConnection();
         var transaction = context.Database.CurrentTransaction?.GetDbTransaction() as NpgsqlTransaction;
 
@@ -75,22 +81,13 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
             await connection.OpenAsync(cancellationToken);
         }
 
-        // Check if record exists
-        var existingVersion = await GetExistingVersionAsync(
-            connection,
-            transaction,
-            saga.Name,
-            state.Id,
-            cancellationToken);
-
-        // Serialize state to JSON
         var jsonData = SerializeState(saga, state);
         var newVersion = NewVersion();
         var now = timeProvider.GetUtcNow();
+        var key = new SagaStateKey(saga.Name, state.Id);
 
-        if (existingVersion is null)
+        if (!_versions.TryGetValue(key, out var version))
         {
-            // Insert new record
             await using var cmd = connection.CreateCommand();
             cmd.CommandText = queries.InsertState;
             cmd.Transaction = transaction;
@@ -102,11 +99,21 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
             cmd.Parameters.AddWithValue("@version", newVersion);
             await cmd.PrepareAsync(cancellationToken);
 
-            await cmd.ExecuteNonQueryAsync(cancellationToken);
+            try
+            {
+                await cmd.ExecuteNonQueryAsync(cancellationToken);
+            }
+            catch (PostgresException ex) when (ex.SqlState == PostgresErrorCodes.UniqueViolation)
+            {
+                throw new DbUpdateConcurrencyException(
+                    "The saga state was concurrently created or already exists.",
+                    ex);
+            }
+
+            _versions[key] = newVersion;
         }
         else
         {
-            // Update existing record with optimistic concurrency
             await using var cmd = connection.CreateCommand();
             cmd.CommandText = queries.UpdateState;
             cmd.Transaction = transaction;
@@ -115,14 +122,16 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
             cmd.Parameters.AddWithValue("@newVersion", newVersion);
             cmd.Parameters.AddWithValue("@id", state.Id);
             cmd.Parameters.AddWithValue("@sagaName", saga.Name);
-            cmd.Parameters.AddWithValue("@oldVersion", existingVersion.Value);
+            cmd.Parameters.AddWithValue("@oldVersion", version);
             await cmd.PrepareAsync(cancellationToken);
 
             var rowsAffected = await cmd.ExecuteNonQueryAsync(cancellationToken);
             if (rowsAffected == 0)
             {
                 throw new DbUpdateConcurrencyException("The saga state was modified by another process.");
             }
+
+            _versions[key] = newVersion;
         }
     }
 
@@ -134,6 +143,8 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
     /// <param name="cancellationToken">A token to observe for cancellation.</param>
     public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellationToken)
     {
+        EnsureCurrentTransactionGeneration();
+
         var connection = (NpgsqlConnection)context.Database.GetDbConnection();
 
         if (connection.State != System.Data.ConnectionState.Open)
@@ -148,6 +159,7 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation
         await cmd.PrepareAsync(cancellationToken);
 
         await cmd.ExecuteNonQueryAsync(cancellationToken);
+        _versions.TryRemove(new SagaStateKey(saga.Name, id), out _);
     }
 
     /// <summary>
@@ -160,6 +172,8 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation
     /// <returns>The deserialized saga state, or <c>default</c> if no state is found for the given identifier.</returns>
     public async Task<T?> LoadAsync<T>(Saga saga, Guid id, CancellationToken cancellationToken)
     {
+        EnsureCurrentTransactionGeneration();
+
         var connection = (NpgsqlConnection)context.Database.GetDbConnection();
         var transaction = context.Database.CurrentTransaction?.GetDbTransaction() as NpgsqlTransaction;
 
@@ -182,25 +196,26 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation
         }
 
         var stateJson = reader.GetFieldValue<ReadOnlyMemory<byte>>(0);
-        return saga.StateSerializer.Deserialize<T>(stateJson);
+        var version = reader.GetFieldValue<Guid>(1);
+        var state = saga.StateSerializer.Deserialize<T>(stateJson);
+
+        if (state is SagaStateBase)
+        {
+            _versions[new SagaStateKey(saga.Name, id)] = version;
+        }
+
+        return state;
     }
 
-    private async Task<Guid?> GetExistingVersionAsync(
-        NpgsqlConnection connection,
-        NpgsqlTransaction? transaction,
-        string sagaName,
-        Guid id,
-        CancellationToken cancellationToken)
+    private void EnsureCurrentTransactionGeneration()
     {
-        await using var cmd = connection.CreateCommand();
-        cmd.CommandText = queries.SelectVersion;
-        cmd.Transaction = transaction;
-        cmd.Parameters.AddWithValue("@id", id);
-        cmd.Parameters.AddWithValue("@sagaName", sagaName);
-        await cmd.PrepareAsync(cancellationToken);
+        var transactionId = context.Database.CurrentTransaction?.TransactionId;
 
-        var result = await cmd.ExecuteScalarAsync(cancellationToken);
-        return result is Guid version ? version : null;
+        if (_transactionId != transactionId)
+        {
+            _versions.Clear();
+            _transactionId = transactionId;
+        }
     }
 
     private byte[] SerializeState(Saga saga, SagaStateBase state)
@@ -232,4 +247,6 @@ public void Dispose()
     {
         _arrayWriter?.Dispose();
     }
+
+    private readonly record struct SagaStateKey(string SagaName, Guid Id);
 }

src/Mocha/src/Mocha.EntityFrameworkCore.Postgres/Sagas/PostgresSagaStoreQueries.cs

@@ -9,15 +9,10 @@ namespace Mocha.Sagas.EfCore;
 internal sealed class PostgresSagaStoreQueries
 {
     /// <summary>
-    /// Gets or sets the SQL query to select the serialized saga state by id and saga name.
+    /// Gets or sets the SQL query to select the serialized saga state and concurrency version by id and saga name.
     /// </summary>
     public string SelectState { get; set; } = null!;
 
-    /// <summary>
-    /// Gets or sets the SQL query to select the concurrency version by id and saga name.
-    /// </summary>
-    public string SelectVersion { get; set; } = null!;
-
     /// <summary>
     /// Gets or sets the SQL statement to insert a new saga state record.
     /// </summary>
@@ -43,13 +38,7 @@ public static PostgresSagaStoreQueries From(SagaStateTableInfo t)
         return new PostgresSagaStoreQueries
         {
             SelectState = $"""
-                SELECT "{t.State}"
-                FROM {t.QualifiedTableName}
-                WHERE "{t.Id}" = @id AND "{t.SagaName}" = @sagaName;
-                """,
-
-            SelectVersion = $"""
-                SELECT "{t.Version}"
+                SELECT "{t.State}", "{t.Version}"
                 FROM {t.QualifiedTableName}
                 WHERE "{t.Id}" = @id AND "{t.SagaName}" = @sagaName;
                 """,