From f259552471ee0018d6b63f52b9de78224baa2528 Mon Sep 17 00:00:00 2001 From: Rhys Bevilaqua Date: Mon, 27 Jul 2026 16:20:14 +0800 Subject: [PATCH 1/2] Create upsert utility extension method --- .../PostgreSqlServiceControlDbContext.cs | 29 +++++++++ .../SqlServerServiceControlDbContext.cs | 13 ++++ .../DbContexts/ServiceControlDbContext.cs | 2 + .../EfCoreExtensions.cs | 52 +++++++++++++++ .../Implementation/MonitoringDataStore.cs | 61 ++++-------------- .../EFCore/EFCoreExtensionMethodTests.cs | 64 +++++++++++++++++++ 6 files changed, 174 insertions(+), 47 deletions(-) create mode 100644 src/ServiceControl.Persistence.EFCore/EfCoreExtensions.cs create mode 100644 src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs diff --git a/src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlServiceControlDbContext.cs b/src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlServiceControlDbContext.cs index 6164f58d2f..2e2daa87aa 100644 --- a/src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlServiceControlDbContext.cs +++ b/src/ServiceControl.Persistence.EFCore.PostgreSql/PostgreSqlServiceControlDbContext.cs @@ -1,6 +1,7 @@ namespace ServiceControl.Persistence.EFCore.PostgreSql; using Microsoft.EntityFrameworkCore; +using Npgsql; using ServiceControl.MessageFailures; using ServiceControl.Persistence.EFCore.DbContexts; using ServiceControl.Persistence.EFCore.Entities; @@ -23,4 +24,32 @@ protected override void OnModelCreating(ModelBuilder modelBuilder) .HasIndex(e => e.StatusChangedAt) .HasFilter($"status IN ({(int)FailedMessageStatus.Resolved}, {(int)FailedMessageStatus.Archived})"); } + + public override bool IsDuplicateKeyException(DbUpdateException exception) + { + var queue = new Queue([exception]); + while (queue.Count > 0) + { + var e = queue.Dequeue(); + if (e is PostgresException { SqlState: PostgresErrorCodes.UniqueViolation }) + { + return true; + } + + //it is unlikely, but there are cases where postgres EF throws aggregate exceptions + if (e is AggregateException aggregateException) + { + foreach (var inner in aggregateException.InnerExceptions) + { + queue.Enqueue(inner); + } + } + + if (e.InnerException != null) + { + queue.Enqueue(e.InnerException); + } + } + return false; + } } diff --git a/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs b/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs index 8c77803995..4c7ea9bb58 100644 --- a/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs +++ b/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs @@ -1,5 +1,6 @@ namespace ServiceControl.Persistence.EFCore.SqlServer; +using Microsoft.Data.SqlClient; using Microsoft.EntityFrameworkCore; using ServiceControl.MessageFailures; using ServiceControl.Persistence.EFCore.DbContexts; @@ -15,4 +16,16 @@ protected override void OnModelCreating(ModelBuilder modelBuilder) .HasIndex(e => e.StatusChangedAt) .HasFilter($"[Status] IN ({(int)FailedMessageStatus.Resolved}, {(int)FailedMessageStatus.Archived})"); } + + public override bool IsDuplicateKeyException(DbUpdateException exception) + { + for (var inner = exception.InnerException; inner != null; inner = inner.InnerException) + { + if (inner is SqlException { Number: 2601 or 2627 }) + { + return true; + } + } + return false; + } } diff --git a/src/ServiceControl.Persistence.EFCore/DbContexts/ServiceControlDbContext.cs b/src/ServiceControl.Persistence.EFCore/DbContexts/ServiceControlDbContext.cs index 49990d4363..e9c9a9a9ba 100644 --- a/src/ServiceControl.Persistence.EFCore/DbContexts/ServiceControlDbContext.cs +++ b/src/ServiceControl.Persistence.EFCore/DbContexts/ServiceControlDbContext.cs @@ -25,4 +25,6 @@ protected override void OnModelCreating(ModelBuilder modelBuilder) modelBuilder.ApplyConfiguration(new FailedMessageGroupConfiguration()); modelBuilder.ApplyConfiguration(new FailedMessageRetryConfiguration()); } + + public abstract bool IsDuplicateKeyException(DbUpdateException exception); } diff --git a/src/ServiceControl.Persistence.EFCore/EfCoreExtensions.cs b/src/ServiceControl.Persistence.EFCore/EfCoreExtensions.cs new file mode 100644 index 0000000000..7404e265e3 --- /dev/null +++ b/src/ServiceControl.Persistence.EFCore/EfCoreExtensions.cs @@ -0,0 +1,52 @@ +namespace ServiceControl.Persistence.EFCore; + +using DbContexts; +using Microsoft.EntityFrameworkCore; + +public static class EfCoreExtensions +{ + /// + /// Inserts a new entity when it does not exist, or updates the existing entity when it does. + /// + /// The EF Core database context used to query and persist the entity. + /// The key values used by to locate the entity. + /// Factory used to create a new entity instance when no matching entity is found. + /// Update action applied to the resolved entity before saving changes. + /// Token used to cancel database operations. + /// A implementation. + /// The entity type to insert or update. + /// + /// This extension method is intended to be a zero dependency utility and has been implemented + /// in a relatively naive way using standard EF building blocks. + /// If more performance critical uses for this are required please consider either implementing your + /// upserts using dialect specific implementations or a library such as https://github.com/artiomchi/FlexLabs.Upsert + /// + public static async Task UpsertAsync(this TContext context, + object?[] keys, + Func create, + Action update, + CancellationToken cancellationToken = default) + where TContext : ServiceControlDbContext + where TEntity : class + { + var entity = await context.FindAsync(keys, cancellationToken: cancellationToken); + if (entity == null) + { + entity = create(); + try + { + context.Add(entity); + await context.SaveChangesAsync(cancellationToken); + return; + } + catch (DbUpdateException e) when (context.IsDuplicateKeyException(e)) + { + //most likely an insert conflict, reload the object and fall through to the update logic + await context.Entry(entity).ReloadAsync(cancellationToken); + } + } + + update(entity); + await context.SaveChangesAsync(cancellationToken); + } +} \ No newline at end of file diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/MonitoringDataStore.cs b/src/ServiceControl.Persistence.EFCore/Implementation/MonitoringDataStore.cs index 1212338c56..38f458422d 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/MonitoringDataStore.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/MonitoringDataStore.cs @@ -7,9 +7,8 @@ namespace ServiceControl.Persistence.EFCore.Implementation; public class MonitoringDataStore(IServiceScopeFactory scopeFactory) : DataStoreBase(scopeFactory), IMonitoringDataStore { - public Task CreateIfNotExists(EndpointDetails endpoint) - { - return ExecuteWithDbContext(async dbContext => + public Task CreateIfNotExists(EndpointDetails endpoint) => + ExecuteWithDbContext(async dbContext => { var id = endpoint.GetDeterministicId(); @@ -45,58 +44,29 @@ public Task CreateIfNotExists(EndpointDetails endpoint) } } }); - } - public Task CreateOrUpdate(EndpointDetails endpoint, IEndpointInstanceMonitoring endpointInstanceMonitoring) - { - return ExecuteWithDbContext(async dbContext => + public Task CreateOrUpdate(EndpointDetails endpoint, IEndpointInstanceMonitoring endpointInstanceMonitoring) => + ExecuteWithDbContext(async dbContext => { var id = endpoint.GetDeterministicId(); - - var knownEndpoint = await dbContext.KnownEndpoints.FirstOrDefaultAsync(e => e.Id == id); - - if (knownEndpoint == null) - { - knownEndpoint = new KnownEndpointEntity + await dbContext.UpsertAsync([id], + () => new KnownEndpointEntity { Id = id, Name = endpoint.Name, HostId = endpoint.HostId, Host = endpoint.Host, Monitored = true - }; - dbContext.KnownEndpoints.Add(knownEndpoint); - - try - { - await dbContext.SaveChangesAsync(); - } - catch (DbUpdateException) + }, + knownEndpoint => { - // A concurrent insert with the same deterministic id may have won the race; - // fall back to the update path. Rethrow anything else. - dbContext.Entry(knownEndpoint).State = EntityState.Detached; - var monitored = endpointInstanceMonitoring.IsMonitored(id); - var updated = await dbContext.KnownEndpoints - .Where(e => e.Id == id) - .ExecuteUpdateAsync(setters => setters.SetProperty(e => e.Monitored, monitored)); - if (updated == 0) - { - throw; - } + knownEndpoint.Monitored = endpointInstanceMonitoring.IsMonitored(id); } - } - else - { - knownEndpoint.Monitored = endpointInstanceMonitoring.IsMonitored(id); - await dbContext.SaveChangesAsync(); - } + ); }); - } - public Task UpdateEndpointMonitoring(EndpointDetails endpoint, bool isMonitored) - { - return ExecuteWithDbContext(async dbContext => + public Task UpdateEndpointMonitoring(EndpointDetails endpoint, bool isMonitored) => + ExecuteWithDbContext(async dbContext => { var id = endpoint.GetDeterministicId(); @@ -104,11 +74,9 @@ await dbContext.KnownEndpoints .Where(e => e.Id == id) .ExecuteUpdateAsync(setters => setters.SetProperty(e => e.Monitored, isMonitored)); }); - } - public Task WarmupMonitoringFromPersistence(IEndpointInstanceMonitoring endpointInstanceMonitoring) - { - return ExecuteWithDbContext(async dbContext => + public Task WarmupMonitoringFromPersistence(IEndpointInstanceMonitoring endpointInstanceMonitoring) => + ExecuteWithDbContext(async dbContext => { await foreach (var endpoint in dbContext.KnownEndpoints.AsNoTracking().AsAsyncEnumerable()) { @@ -122,7 +90,6 @@ public Task WarmupMonitoringFromPersistence(IEndpointInstanceMonitoring endpoint endpointInstanceMonitoring.DetectEndpointFromPersistentStore(endpointDetails, endpoint.Monitored); } }); - } public Task Delete(Guid endpointId) { diff --git a/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs b/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs new file mode 100644 index 0000000000..f27da83180 --- /dev/null +++ b/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs @@ -0,0 +1,64 @@ +namespace DefaultNamespace; + +using System; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; +using NUnit.Framework; +using ServiceControl.Persistence.EFCore; +using ServiceControl.Persistence.EFCore.DbContexts; +using ServiceControl.Persistence.EFCore.Entities; +using ServiceControl.Persistence.Tests; + +public class EFCoreExtensionMethodTests : PersistenceTestBase +{ + [Test, CancelAfter(60_000)] + public async Task Insert_until_conflict_should_not_throw_errors(CancellationToken cancellationToken) + { + var i = 0; + while (!cancellationToken.IsCancellationRequested) + { + i++; + var key = Guid.NewGuid(); + var result = await Task.WhenAll(Write(key), Write(key), Write(key)); + if (result.Contains(Result.ConflictUpdate)) + { + Assert.Pass($"Test had concurrency failure ({i}): {string.Join(", ", result)}"); + } + } + + Assert.Fail("Test never had concurrency failure"); + + async Task Write(Guid key) + { + await using var scope = ServiceProvider.CreateAsyncScope(); + var context = scope.ServiceProvider.GetRequiredService(); + + var state = Result.Update; + await context.UpsertAsync([key], + () => + { + state = Result.Insert; + return new FailedMessageEntity + { + UniqueMessageId = key, + HeadersJson = "blah" + }; + }, + es => + { + state = state == Result.Insert ? Result.ConflictUpdate : Result.Update; + es.BodyText = Guid.NewGuid().ToString(); + }, + cancellationToken: cancellationToken); + return state; + } + } + + public enum Result + { + Insert, + Update, + ConflictUpdate + } +} \ No newline at end of file From 4e65f156166b248152d42d5c51ec7d939c4aa3cb Mon Sep 17 00:00:00 2001 From: Rhys Bevilaqua Date: Mon, 27 Jul 2026 16:36:58 +0800 Subject: [PATCH 2/2] remove whitespace --- .../SqlServerServiceControlDbContext.cs | 2 +- .../EFCore/EFCoreExtensionMethodTests.cs | 10 +++------- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs b/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs index 4c7ea9bb58..9fff389d1f 100644 --- a/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs +++ b/src/ServiceControl.Persistence.EFCore.SqlServer/SqlServerServiceControlDbContext.cs @@ -16,7 +16,7 @@ protected override void OnModelCreating(ModelBuilder modelBuilder) .HasIndex(e => e.StatusChangedAt) .HasFilter($"[Status] IN ({(int)FailedMessageStatus.Resolved}, {(int)FailedMessageStatus.Archived})"); } - + public override bool IsDuplicateKeyException(DbUpdateException exception) { for (var inner = exception.InnerException; inner != null; inner = inner.InnerException) diff --git a/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs b/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs index f27da83180..3600ddafd3 100644 --- a/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs +++ b/src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs @@ -39,11 +39,7 @@ await context.UpsertAsync([key], () => { state = Result.Insert; - return new FailedMessageEntity - { - UniqueMessageId = key, - HeadersJson = "blah" - }; + return new FailedMessageEntity { UniqueMessageId = key, HeadersJson = "blah" }; }, es => { @@ -54,8 +50,8 @@ await context.UpsertAsync([key], return state; } } - - public enum Result + + enum Result { Insert, Update,