fix: serialize InitializeDb across concurrent instances

README.md

@@ -142,7 +142,9 @@ builder.Services.AddDataProtection()
 `InitializeTable` is `true`, creates the table and the friendly-name unique index if they are
 missing. The index is also created on an existing table that lacks it; if that table already
 contains duplicate non-null friendly names, initialization fails and the duplicates must be
-removed before retrying. If you set
+removed before retrying. Initialization is safe when several instances start at the same time:
+it runs in a transaction that holds a database-wide lock (`sp_getapplock` on SQL Server,
+`pg_advisory_xact_lock` on PostgreSQL), so concurrent calls wait for each other. If you set
 `InitializeTable = false` you are responsible for creating the schema (see below) before the
 app stores keys.
 

src/AspNetCore.DataProtection.CustomStorage.Dapper.PostgreSQL/PostgreSQLDataProtectionRepository.cs

@@ -11,6 +11,8 @@ namespace AspNetCore.DataProtection.CustomStorage.Dapper.PostgreSQL;
 /// <inheritdoc />
 public class PostgreSQLDataProtectionRepository : IDbDataProtectionStorage
 {
+    // Database-wide: initializations of different tables are serialized too, which is harmless.
+    private const string InitializeDbLockName = "AspNetCore.DataProtection.CustomStorage.InitializeDb";
     private readonly NpgsqlDataSource _dataSource;
     private readonly DapperDataProtectionConfig _config;
     private string GetAllQuery() => 
@@ -104,10 +106,19 @@ private static List<int> GetIds(IReadOnlyCollection<DataProtectionKey> keys)
     }
 
     /// <inheritdoc />
+    /// <remarks>
+    /// Safe to call concurrently from several instances: the initialization is serialized
+    /// with a transaction-level advisory lock.
+    /// </remarks>
     public void InitializeDb()
     {
-        using var connection = _dataSource.CreateConnection();
+        using var connection = _dataSource.OpenConnection();
+        using var transaction = connection.BeginTransaction();
+        // Without the lock, concurrent IF NOT EXISTS statements can all create the same relation
+        // and fail with a unique violation in the system catalogs.
         connection.Execute($"""
+             SELECT pg_advisory_xact_lock(hashtext(@LockName));
+
              CREATE TABLE IF NOT EXISTS {_config.SchemaName}.{_config.TableName}
               (
                  id INTEGER generated always as identity,
@@ -121,6 +132,9 @@ CREATE UNIQUE INDEX IF NOT EXISTS ix_{_config.SchemaName}_{_config.TableName}_fr
               ON {_config.SchemaName}.{_config.TableName} USING btree
               (friendly_name COLLATE pg_catalog."default" ASC NULLS LAST)
               TABLESPACE pg_default;
-             """);
+             """,
+            new { LockName = InitializeDbLockName },
+            transaction);
+        transaction.Commit();
     }
 }

src/AspNetCore.DataProtection.CustomStorage.Dapper.SQLServer/SQLServerDataProtectionRepository.cs

@@ -10,6 +10,8 @@ namespace AspNetCore.DataProtection.CustomStorage.Dapper.SQLServer;
 /// <inheritdoc />
 public class SQLServerDataProtectionRepository : IDbDataProtectionStorage
 {
+    // Database-wide: initializations of different tables are serialized too, which is harmless.
+    private const string InitializeDbLockResource = "AspNetCore.DataProtection.CustomStorage.InitializeDb";
     private readonly SqlConnectionFactory _connectionFactory;
     private readonly DapperDataProtectionConfig _config;
 
@@ -104,14 +106,25 @@ private static List<int> GetIds(IReadOnlyCollection<DataProtectionKey> keys)
     }
 
     /// <inheritdoc />
+    /// <remarks>
+    /// Safe to call concurrently from several instances: the initialization is serialized
+    /// with an exclusive application lock held until the transaction ends.
+    /// </remarks>
     public void InitializeDb()
     {
         var indexName = $"IX_{_config.TableName}_{nameof(DataProtectionKeyEntity.FriendlyName)}";
         using var connection = _connectionFactory.CreateConnection();
+        connection.Open();
+        using var transaction = connection.BeginTransaction();
+        // Without the lock, concurrent existence checks can all pass and the later CREATE statements fail.
         // The unique index is ensured independently of the table, so that an existing table
         // without the index (created manually, or with the index dropped) gets it restored.
         connection.Execute(
             $"""
+               DECLARE @lockResult int;
+               EXEC @lockResult = sp_getapplock @Resource = @LockResource, @LockMode = N'Exclusive', @LockOwner = N'Transaction';
+               IF @lockResult < 0
+                  THROW 50000, N'Cannot acquire the application lock to initialize the data protection keys table.', 1;
                IF NOT EXISTS(SELECT t.[object_id]
                FROM [sys].[tables] t
                   JOIN sys.schemas s ON s.schema_id=t.schema_id
@@ -137,6 +150,9 @@ HAVING COUNT(*) > 1)
                      THROW 50000, N'Cannot create unique index [{indexName}] on [{_config.SchemaName}].[{_config.TableName}]: the table contains duplicate non-null [{nameof(DataProtectionKeyEntity.FriendlyName)}] values. Remove the duplicates and initialize again.', 1;
                   CREATE UNIQUE NONCLUSTERED INDEX [{indexName}] ON [{_config.SchemaName}].[{_config.TableName}]([{nameof(DataProtectionKeyEntity.FriendlyName)}] ASC) WHERE [{nameof(DataProtectionKeyEntity.FriendlyName)}] IS NOT NULL ON [PRIMARY];
                END
-               """);
+               """,
+            new { LockResource = InitializeDbLockResource },
+            transaction);
+        transaction.Commit();
     }
 }

tests/AspNetCore.DataProtection.CustomStorage.Tests/PostgreSQL/PostgreSqlDataProtectionRepositoryTests.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using AspNetCore.DataProtection.CustomStorage.Dapper;
 using AspNetCore.DataProtection.CustomStorage.Dapper.PostgreSQL;
@@ -143,6 +144,45 @@ public void InitializeDb_CustomSchemaAndTable_PersistsKeysInConfiguredTable()
         _sut.GetAll().Should().NotContain(x => x.FriendlyName == key.FriendlyName);
     }
 
+    [Fact]
+    public async Task InitializeDb_ConcurrentCallsOnNewTable_AllSucceed()
+    {
+        const int callCount = 10;
+        // Pooled connections make the calls overlap; with cold connections the race rarely shows.
+        var connections = Enumerable.Range(0, callCount).Select(_ => _fixture.NpgsqlDataSource.OpenConnection()).ToList();
+        connections.ForEach(connection => connection.Dispose());
+
+        // Several rounds, because a single round does not always hit the race.
+        for (var round = 0; round < 5; round++)
+        {
+            var tableName = $"keys_{Guid.NewGuid().ToString("N")[..12]}";
+            var options = Substitute.For<IOptionsSnapshot<DapperDataProtectionConfig>>();
+            options.Value.Returns(new DapperDataProtectionConfig
+            {
+                SchemaName = Constants.DefaultSchema,
+                TableName = tableName
+            });
+            // One repository per call, as on separate nodes starting at the same time.
+            var suts = Enumerable.Range(0, callCount)
+                .Select(_ => new PostgreSQLDataProtectionRepository(_fixture.NpgsqlDataSource, options))
+                .ToList();
+            using var barrier = new Barrier(callCount);
+
+            var calls = suts.Select(sut => Task.Factory.StartNew(() =>
+            {
+                barrier.SignalAndWait();
+                sut.InitializeDb();
+            }, TaskCreationOptions.LongRunning)).ToList();
+
+            await FluentActions.Awaiting(() => Task.WhenAll(calls)).Should().NotThrowAsync();
+            using var connection = _fixture.NpgsqlDataSource.OpenConnection();
+            connection.ExecuteScalar<long>(
+                    "SELECT COUNT(*) FROM pg_indexes WHERE schemaname = @Schema AND indexname = @Index;",
+                    new { Schema = Constants.DefaultSchema, Index = $"ix_{Constants.DefaultSchema}_{tableName}_friendly_name" })
+                .Should().Be(1);
+        }
+    }
+
     [Fact]
     public void Delete_NullKeys_ThrowsArgumentNullException()
     {

tests/AspNetCore.DataProtection.CustomStorage.Tests/SQLServer/SqlServerDataProtectionRepositoryTests.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using AspNetCore.DataProtection.CustomStorage.Dapper;
 using AspNetCore.DataProtection.CustomStorage.Dapper.SQLServer;
@@ -278,6 +279,34 @@ public void InitializeDb_CustomSchemaAndTable_PersistsKeysInConfiguredTable()
         _sut.GetAll().Should().NotContain(x => x.FriendlyName == key.FriendlyName);
     }
 
+    [Fact]
+    public async Task InitializeDb_ConcurrentCallsOnNewTable_AllSucceed()
+    {
+        const int callCount = 10;
+        // Pooled connections make the calls overlap; with cold connections the race rarely shows.
+        var connections = Enumerable.Range(0, callCount).Select(_ => _fixture.SqlConnectionFactory.CreateConnection()).ToList();
+        connections.ForEach(connection => connection.Open());
+        connections.ForEach(connection => connection.Dispose());
+
+        // Several rounds, because a single round does not always hit the race.
+        for (var round = 0; round < 5; round++)
+        {
+            var tableName = NewTableName();
+            // One repository per call, as on separate nodes starting at the same time.
+            var suts = Enumerable.Range(0, callCount).Select(_ => CreateSut(tableName)).ToList();
+            using var barrier = new Barrier(callCount);
+
+            var calls = suts.Select(sut => Task.Factory.StartNew(() =>
+            {
+                barrier.SignalAndWait();
+                sut.InitializeDb();
+            }, TaskCreationOptions.LongRunning)).ToList();
+
+            await FluentActions.Awaiting(() => Task.WhenAll(calls)).Should().NotThrowAsync();
+            FriendlyNameIndexExists(tableName).Should().BeTrue();
+        }
+    }
+
     private static string NewTableName() => $"Keys_{Guid.NewGuid():N}";
 
     private SQLServerDataProtectionRepository CreateSut(string tableName, string schemaName = Constants.DefaultSchema)