Create upsert utility extension method

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>([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;
+    }
 }

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;
+    }
 }

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);
 }

src/ServiceControl.Persistence.EFCore/EfCoreExtensions.cs

@@ -0,0 +1,52 @@
+namespace ServiceControl.Persistence.EFCore;
+
+using DbContexts;
+using Microsoft.EntityFrameworkCore;
+
+public static class EfCoreExtensions
+{
+    /// <summary>
+    /// Inserts a new entity when it does not exist, or updates the existing entity when it does.
+    /// </summary>
+    /// <param name="context">The EF Core database context used to query and persist the entity.</param>
+    /// <param name="keys">The key values used by <see cref="DbContext.FindAsync{TEntity}(object?[], CancellationToken)"/> to locate the entity.</param>
+    /// <param name="create">Factory used to create a new entity instance when no matching entity is found.</param>
+    /// <param name="update">Update action applied to the resolved entity before saving changes.</param>
+    /// <param name="cancellationToken">Token used to cancel database operations.</param>
+    /// <typeparam name="TContext">A <see cref="ServiceControlDbContext"/> implementation.</typeparam>
+    /// <typeparam name="TEntity">The entity type to insert or update.</typeparam>
+    /// <remarks>
+    /// 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
+    /// </remarks>
+    public static async Task UpsertAsync<TContext, TEntity>(this TContext context,
+        object?[] keys,
+        Func<TEntity> create,
+        Action<TEntity> update,
+        CancellationToken cancellationToken = default)
+        where TContext : ServiceControlDbContext
+        where TEntity : class
+    {
+        var entity = await context.FindAsync<TEntity>(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);
+    }
+}

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,70 +44,39 @@ 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();
 
             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)
     {

src/ServiceControl.Persistence.Tests/EFCore/EFCoreExtensionMethodTests.cs

@@ -0,0 +1,60 @@
+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<Result> Write(Guid key)
+        {
+            await using var scope = ServiceProvider.CreateAsyncScope();
+            var context = scope.ServiceProvider.GetRequiredService<ServiceControlDbContext>();
+
+            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;
+        }
+    }
+
+    enum Result
+    {
+        Insert,
+        Update,
+        ConflictUpdate
+    }
+}