Fix duplicate-key races in endpoint throughput recording

src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlEndpointThroughputDialect.cs

@@ -0,0 +1,31 @@
+namespace ServiceControl.Persistence.EFCore.PostgreSql;
+
+using Particular.LicensingComponent.Contracts;
+using ServiceControl.Persistence.EFCore.DbContexts;
+using ServiceControl.Persistence.EFCore.Entities;
+using ServiceControl.Persistence.EFCore.Infrastructure;
+
+class PostgreSqlEndpointThroughputDialect : PostgreSqlDialect, IEndpointThroughputDialect
+{
+    public async Task RecordEndpointThroughput(ServiceControlDbContext dbContext, string normalizedName, ThroughputSource throughputSource, IReadOnlyList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken = default)
+    {
+        foreach (var (date, messageCount) in throughput)
+        {
+            // INSERT ... ON CONFLICT is the atomic add-or-insert: the insert arm creates the day, the
+            // conflict arm adds to the day's total. PostgreSQL serializes concurrent inserts of the
+            // same key inside the uniqueness check, so there is no window in which a duplicate key
+            // can surface to the application at all.
+            var table = Table<LicensingEndpointThroughputEntity>(dbContext);
+            await Execute(
+                dbContext,
+                $"""
+                 INSERT INTO {table} (normalized_name, throughput_source, date_utc, message_count)
+                 VALUES (@p0, @p1, @p2, @p3)
+                 ON CONFLICT (normalized_name, throughput_source, date_utc)
+                 DO UPDATE SET message_count = {table}.message_count + EXCLUDED.message_count
+                 """,
+                [[normalizedName, (int)throughputSource, date, messageCount]],
+                cancellationToken);
+        }
+    }
+}

src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlPersistence.cs

@@ -20,6 +20,7 @@ public void AddPersistence(IServiceCollection services)
         services.AddSingleton<IFailedMessageIngestionSqlDialect, PostgreSqlFailedMessageIngestionSqlDialect>();
         services.AddSingleton<IRetryBatchSqlDialect, PostgreSqlRetryBatchSqlDialect>();
         services.AddSingleton<IFullTextSearchDialect, PostgreSqlFullTextSearchDialect>();
+        services.AddSingleton<IEndpointThroughputDialect, PostgreSqlEndpointThroughputDialect>();
         services.AddSingleton<IDatabaseHostingProbe, PostgreSqlDatabaseHostingProbe>();
     }
 

src/ServiceControl.Persistence.EFCore.PostgreSql/ServiceControl.Persistence.EFCore.PostgreSql.csproj

@@ -15,12 +15,14 @@
     <ProjectReference Include="../ServiceControl.Configuration/ServiceControl.Configuration.csproj" Private="false" ExcludeAssets="runtime" />
     <ProjectReference Include="../ServiceControl.Infrastructure/ServiceControl.Infrastructure.csproj" Private="false" ExcludeAssets="runtime" />
     <ProjectReference Include="../ServiceControl.Persistence/ServiceControl.Persistence.csproj" Private="false" ExcludeAssets="runtime" />
+    <ProjectReference Include="../Particular.LicensingComponent.Contracts/Particular.LicensingComponent.Contracts.csproj" Private="false" ExcludeAssets="runtime" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(Configuration)' != 'Release'">
     <ProjectReference Include="../ServiceControl.Configuration/ServiceControl.Configuration.csproj" />
     <ProjectReference Include="../ServiceControl.Infrastructure/ServiceControl.Infrastructure.csproj" />
     <ProjectReference Include="../ServiceControl.Persistence/ServiceControl.Persistence.csproj" />
+    <ProjectReference Include="../Particular.LicensingComponent.Contracts/Particular.LicensingComponent.Contracts.csproj" />
   </ItemGroup>
 
   <ItemGroup>

src/ServiceControl.Persistence.EFCore.SqlServer/ServiceControl.Persistence.EFCore.SqlServer.csproj

@@ -15,12 +15,14 @@
     <ProjectReference Include="../ServiceControl.Configuration/ServiceControl.Configuration.csproj" Private="false" ExcludeAssets="runtime" />
     <ProjectReference Include="../ServiceControl.Infrastructure/ServiceControl.Infrastructure.csproj" Private="false" ExcludeAssets="runtime" />
     <ProjectReference Include="../ServiceControl.Persistence/ServiceControl.Persistence.csproj" Private="false" ExcludeAssets="runtime" />
+    <ProjectReference Include="../Particular.LicensingComponent.Contracts/Particular.LicensingComponent.Contracts.csproj" Private="false" ExcludeAssets="runtime" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(Configuration)' != 'Release'">
     <ProjectReference Include="../ServiceControl.Configuration/ServiceControl.Configuration.csproj" />
     <ProjectReference Include="../ServiceControl.Infrastructure/ServiceControl.Infrastructure.csproj" />
     <ProjectReference Include="../ServiceControl.Persistence/ServiceControl.Persistence.csproj" />
+    <ProjectReference Include="../Particular.LicensingComponent.Contracts/Particular.LicensingComponent.Contracts.csproj" />
   </ItemGroup>
 
   <ItemGroup>

src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerEndpointThroughputDialect.cs

@@ -0,0 +1,36 @@
+namespace ServiceControl.Persistence.EFCore.SqlServer;
+
+using System.Data;
+using System.Text;
+using Microsoft.EntityFrameworkCore;
+using Microsoft.EntityFrameworkCore.Storage;
+using Particular.LicensingComponent.Contracts;
+using ServiceControl.Persistence.EFCore.DbContexts;
+using ServiceControl.Persistence.EFCore.Entities;
+using ServiceControl.Persistence.EFCore.Infrastructure;
+
+class SqlServerEndpointThroughputDialect : SqlServerDialect, IEndpointThroughputDialect
+{
+    public async Task RecordEndpointThroughput(ServiceControlDbContext dbContext, string normalizedName, ThroughputSource throughputSource, IReadOnlyList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken = default)
+    {
+        foreach (var (date, messageCount) in throughput)
+        {
+            // MERGE is the atomic add-or-insert: the WHEN MATCHED arm adds to the day's total, the
+            // WHEN NOT MATCHED arm creates the day. HOLDLOCK serializes concurrent recorders on the
+            // same day's key range, so there is no window in which two inserts of the same day can
+            // both pass the match test.
+            await Execute(
+                dbContext,
+                $"""
+                 MERGE {Table<LicensingEndpointThroughputEntity>(dbContext)} WITH (HOLDLOCK) AS t
+                 USING (VALUES (@p0, @p1, @p2, @p3)) AS s (NormalizedEndpointName, Source, Day, Count)
+                 ON t.NormalizedName = s.NormalizedEndpointName AND t.ThroughputSource = s.Source AND t.DateUtc = s.Day
+                 WHEN MATCHED THEN UPDATE SET t.MessageCount = t.MessageCount + s.Count
+                 WHEN NOT MATCHED THEN INSERT (NormalizedName, ThroughputSource, DateUtc, MessageCount)
+                 VALUES (s.NormalizedEndpointName, s.Source, s.Day, s.Count);
+                 """,
+                [[normalizedName, (int)throughputSource, date, messageCount]],
+                cancellationToken);
+        }
+    }
+}

src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerPersistence.cs

@@ -20,6 +20,7 @@ public void AddPersistence(IServiceCollection services)
         services.AddSingleton<IFailedMessageIngestionSqlDialect, SqlServerFailedMessageIngestionSqlDialect>();
         services.AddSingleton<IRetryBatchSqlDialect, SqlServerRetryBatchSqlDialect>();
         services.AddSingleton<IFullTextSearchDialect, SqlServerFullTextSearchDialect>();
+        services.AddSingleton<IEndpointThroughputDialect, SqlServerEndpointThroughputDialect>();
         services.AddSingleton<IDatabaseHostingProbe, SqlServerDatabaseHostingProbe>();
     }
 

src/ServiceControl.Persistence.EFCore/Implementation/LicensingDataStore.cs

@@ -8,10 +8,8 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 using ServiceControl.Persistence.EFCore.Entities;
 using ServiceControl.Persistence.EFCore.Infrastructure;
 
-class LicensingDataStore(IServiceScopeFactory scopeFactory, TimeProvider timeProvider) : DataStoreBase(scopeFactory), ILicensingDataStore
+class LicensingDataStore(IServiceScopeFactory scopeFactory, TimeProvider timeProvider, IEndpointThroughputDialect throughputDialect) : DataStoreBase(scopeFactory), ILicensingDataStore
 {
-    const int MaxRecordAttempts = 5;
-
     static readonly string PlatformEndpointIndicator = EndpointIndicator.PlatformEndpoint.ToString();
     static readonly AuditServiceMetadata DefaultAuditServiceMetadata = new([], []);
     static readonly BrokerMetadata DefaultBrokerMetadata = new(null, []);
@@ -162,89 +160,46 @@ public Task<IDictionary<string, IEnumerable<ThroughputData>>> GetEndpointThrough
             return (IDictionary<string, IEnumerable<ThroughputData>>)results;
         }, cancellationToken);
 
-    public async Task RecordEndpointThroughput(string endpointName, ThroughputSource throughputSource, IList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken = default)
+    public Task RecordEndpointThroughput(string endpointName, ThroughputSource throughputSource, IList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken = default)
     {
         if (throughput.Count == 0)
         {
-            return;
-        }
-
-        // Only the first recording of a day can lose a race, and once its row exists every later
-        // writer takes the update path, so a couple of attempts is enough.
-        for (var attempt = 1; attempt <= MaxRecordAttempts; attempt++)
-        {
-            var recorded = await ExecuteWithDbContext((context, token) =>
-                TryRecordEndpointThroughput(context, endpointName, throughputSource, throughput, token), cancellationToken);
-
-            if (recorded)
-            {
-                return;
-            }
+            return Task.CompletedTask;
         }
 
-        throw new InvalidOperationException(
-            $"Could not record throughput for {endpointName} from {throughputSource} after {MaxRecordAttempts} attempts because of concurrent updates.");
-    }
-
-    static async Task<bool> TryRecordEndpointThroughput(ServiceControlDbContext context, string endpointName, ThroughputSource throughputSource, IList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken)
-    {
-        var normalizedName = Normalize(endpointName);
-
-        // Ordered so that concurrent calls covering overlapping days take the row locks in the same
-        // order and cannot deadlock against each other.
-        var dailyThroughput = throughput.OrderBy(entry => entry.DateUTC).ToList();
-
-        var strategy = context.Database.CreateExecutionStrategy();
-
-        return await strategy.ExecuteAsync(async () =>
+        return ExecuteWithDbContext(async (context, token) =>
         {
-            await using var transaction = await context.Database.BeginTransactionAsync(cancellationToken);
+            var normalizedName = Normalize(endpointName);
 
-            var endpointExists = await context.LicensingEndpoints
-                .AnyAsync(endpoint => endpoint.NormalizedName == normalizedName && endpoint.ThroughputSource == throughputSource, cancellationToken);
+            // Ordered so that concurrent calls covering overlapping days take the row locks in the
+            // same order and cannot deadlock against each other.
+            var dailyThroughput = throughput.OrderBy(entry => entry.DateUTC).ToList();
 
-            if (!endpointExists)
-            {
-                throw new InvalidOperationException($"Endpoint {endpointName} from {throughputSource} does not exist ");
-            }
+            var strategy = context.Database.CreateExecutionStrategy();
 
-            foreach (var (date, messageCount) in dailyThroughput)
+            await strategy.ExecuteAsync(async () =>
             {
-                // Recording adds to the day's total, because a source can report the same day
-                // repeatedly. The addition happens in the database, so concurrent writers queue up on
-                // the row instead of overwriting each other's totals.
-                var updated = await context.LicensingEndpointThroughput
-                    .Where(row => row.NormalizedName == normalizedName && row.ThroughputSource == throughputSource && row.DateUtc == date)
-                    .ExecuteUpdateAsync(row => row.SetProperty(p => p.MessageCount, p => p.MessageCount + messageCount), cancellationToken);
+                await using var transaction = await context.Database.BeginTransactionAsync(token);
 
-                if (updated > 0)
-                {
-                    continue;
-                }
+                var endpointExists = await context.LicensingEndpoints
+                    .AnyAsync(endpoint => endpoint.NormalizedName == normalizedName && endpoint.ThroughputSource == throughputSource, token);
 
-                context.LicensingEndpointThroughput.Add(new LicensingEndpointThroughputEntity
+                if (!endpointExists)
                 {
-                    NormalizedName = normalizedName,
-                    ThroughputSource = throughputSource,
-                    DateUtc = date,
-                    MessageCount = messageCount
-                });
-
-                try
-                {
-                    await context.SaveChangesAsync(cancellationToken);
-                }
-                catch (DbUpdateException exception) when (context.IsDuplicateKeyException(exception))
-                {
-                    // Another writer created the day's row first, so this call has to add to it
-                    // instead. The failed insert leaves the transaction unusable, hence the retry.
-                    return false;
+                    throw new InvalidOperationException($"Endpoint {endpointName} from {throughputSource} does not exist ");
                 }
-            }
 
-            await transaction.CommitAsync(cancellationToken);
-            return true;
-        });
+                // Recording adds to the day's total, because a source can report the same day
+                // repeatedly, and it has to tolerate the same day being recorded concurrently:
+                // monitoring throughput messages are processed with high concurrency, so two
+                // recorders racing on one endpoint/day is routine. The add and the insert happen as
+                // one atomic statement in the database, so there is no check-then-insert window to
+                // race in and no duplicate-key failure left for either recorder to see.
+                await throughputDialect.RecordEndpointThroughput(context, normalizedName, throughputSource, dailyThroughput, token);
+
+                await transaction.CommitAsync(token);
+            });
+        }, cancellationToken);
     }
 
     public Task RemoveEndpoints(EndpointIdentifier[] endpointIds, CancellationToken cancellationToken = default) =>

src/ServiceControl.Persistence.EFCore/Infrastructure/IEndpointThroughputDialect.cs

@@ -0,0 +1,22 @@
+namespace ServiceControl.Persistence.EFCore.Infrastructure;
+
+using Particular.LicensingComponent.Contracts;
+using ServiceControl.Persistence.EFCore.DbContexts;
+
+/// <summary>
+/// The provider-specific SQL of recording days of endpoint throughput. Recording adds to a day's
+/// total and creates the day's row when absent, and it must do so as one atomic statement: the same
+/// endpoint/day is recorded concurrently (monitoring throughput messages are processed with high
+/// concurrency, and more than one collector can report the same source), so a check-then-insert in
+/// application code races and any insert conflict it leaves behind surfaces as a duplicate-key
+/// failure to the collectors. Implementations run on the DbContext connection inside the
+/// transaction the caller has already opened, like the ingestion dialects.
+/// </summary>
+public interface IEndpointThroughputDialect
+{
+    /// <summary>
+    /// One atomic add-or-insert per day, in the order given (the caller orders by date so concurrent
+    /// calls covering overlapping days take row locks in the same order and cannot deadlock).
+    /// </summary>
+    Task RecordEndpointThroughput(ServiceControlDbContext dbContext, string normalizedName, ThroughputSource throughputSource, IReadOnlyList<EndpointDailyThroughput> throughput, CancellationToken cancellationToken = default);
+}