Made known_endpoint table insert only

Author: johnsimons

src/ServiceControl.Audit.Persistence.PostgreSQL/PostgreSQLPersistence.cs

@@ -25,5 +25,6 @@ public void AddPersistence(IServiceCollection services)
         services.AddSingleton<IBodyStorage, PostgreSQLAttachmentsBodyStorage>();
         services.AddSingleton<PostgreSQLConnectionFactory>();
         services.AddHostedService<RetentionCleanupService>();
+        services.AddHostedService<UpdateKnownEndpointTable>();
     }
 }

src/ServiceControl.Audit.Persistence.PostgreSQL/PostgreSQLPersistenceInstaller.cs

@@ -186,15 +186,28 @@ CREATE INDEX IF NOT EXISTS idx_saga_snapshots_saga_id ON saga_snapshots (
             await cmd.ExecuteNonQueryAsync(cancellationToken);
         }
 
+        // Create known_endpoints table
+        await using (var cmd = new NpgsqlCommand(@"
+            CREATE TABLE IF NOT EXISTS known_endpoints_insert (
+                id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+                endpoint_id TEXT,
+                name TEXT,
+                host_id UUID,
+                host TEXT,
+                last_seen TIMESTAMPTZ
+            );", connection))
+        {
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+        }
+
         // Create known_endpoints table
         await using (var cmd = new NpgsqlCommand(@"
             CREATE TABLE IF NOT EXISTS known_endpoints (
                 id TEXT PRIMARY KEY,
                 name TEXT,
                 host_id UUID,
                 host TEXT,
-                last_seen TIMESTAMPTZ,
-                updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
+                last_seen TIMESTAMPTZ
             );", connection))
         {
             await cmd.ExecuteNonQueryAsync(cancellationToken);
@@ -212,11 +225,6 @@ CREATE OR REPLACE FUNCTION update_updated_at_column() RETURNS trigger AS $$
             DROP TRIGGER IF EXISTS saga_snapshots_updated_at_trigger ON saga_snapshots;
             CREATE TRIGGER saga_snapshots_updated_at_trigger
             BEFORE UPDATE ON saga_snapshots
-            FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
-
-            DROP TRIGGER IF EXISTS known_endpoints_updated_at_trigger ON known_endpoints;
-            CREATE TRIGGER known_endpoints_updated_at_trigger
-            BEFORE UPDATE ON known_endpoints
             FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();", connection))
         {
             await cmd.ExecuteNonQueryAsync(cancellationToken);

src/ServiceControl.Audit.Persistence.PostgreSQL/RetentionCleanupService.cs

@@ -49,7 +49,7 @@ async Task CleanupOldMessages(CancellationToken cancellationToken)
         await CleanupTable("saga_snapshots", "updated_at", cutoffDate, conn, cancellationToken);
 
         // Cleanup known endpoints
-        await CleanupTable("known_endpoints", "updated_at", cutoffDate, conn, cancellationToken);
+        await CleanupTable("known_endpoints", "last_seen", cutoffDate, conn, cancellationToken);
     }
 
     async Task CleanupTable(string tableName, string dateColumn, DateTime cutoffDate, NpgsqlConnection conn, CancellationToken cancellationToken)
@@ -59,10 +59,11 @@ async Task CleanupTable(string tableName, string dateColumn, DateTime cutoffDate
 
         while (!cancellationToken.IsCancellationRequested)
         {
-            // Delete in batches
+            // Delete in batches - skip if another process is already cleaning this table
             var sql = $@"
                 DELETE FROM {tableName}
-                WHERE ctid IN (
+                WHERE pg_try_advisory_xact_lock(hashtext('{tableName}'))
+                AND ctid IN (
                     SELECT ctid FROM {tableName}
                     WHERE {dateColumn} < @cutoff
                     LIMIT {batchSize}

src/ServiceControl.Audit.Persistence.PostgreSQL/UnitOfWork/PostgreSQLAuditIngestionUnitOfWork.cs

@@ -129,15 +129,14 @@ public Task RecordKnownEndpoint(KnownEndpoint knownEndpoint, CancellationToken c
         // Insert KnownEndpoint into known_endpoints table
         var cmd = batch.CreateBatchCommand();
         cmd.CommandText = @"
-                INSERT INTO known_endpoints (
-                    id, name, host_id, host, last_seen
+                
+                INSERT INTO known_endpoints_insert (
+                    endpoint_id, name, host_id, host, last_seen
                 ) VALUES (
-                    @id, @name, @host_id, @host, @last_seen
-                )
-                ON CONFLICT (id) DO UPDATE SET
-                    last_seen = GREATEST(known_endpoints.last_seen, EXCLUDED.last_seen);";
+                    @endpoint_id, @name, @host_id, @host, @last_seen
+                );";
 
-        cmd.Parameters.AddWithValue("id", knownEndpoint.Id);
+        cmd.Parameters.AddWithValue("endpoint_id", knownEndpoint.Id);
         cmd.Parameters.AddWithValue("name", knownEndpoint.Name);
         cmd.Parameters.AddWithValue("host_id", knownEndpoint.HostId);
         cmd.Parameters.AddWithValue("host", knownEndpoint.Host);

src/ServiceControl.Audit.Persistence.PostgreSQL/UpdateKnownEndpointTable.cs

@@ -0,0 +1,61 @@
+namespace ServiceControl.Audit.Persistence.PostgreSQL;
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+using Npgsql;
+
+class UpdateKnownEndpointTable(
+    ILogger<RetentionCleanupService> logger,
+        PostgreSQLConnectionFactory connectionFactory) : BackgroundService
+{
+    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+    {
+        logger.LogInformation($"{nameof(UpdateKnownEndpointTable)} started.");
+
+        while (!stoppingToken.IsCancellationRequested)
+        {
+            try
+            {
+                await Task.Delay(TimeSpan.FromSeconds(30), stoppingToken);
+
+                await UpdateTable(stoppingToken);
+            }
+            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
+            {
+                logger.LogInformation($"{nameof(UpdateKnownEndpointTable)} stopped.");
+                break;
+            }
+            catch (Exception ex)
+            {
+                logger.LogError(ex, "Error during update known_endpoints table.");
+            }
+        }
+    }
+
+    async Task UpdateTable(CancellationToken stoppingToken)
+    {
+        await using var conn = await connectionFactory.OpenConnection(stoppingToken);
+
+        var sql = @"
+            DO $$
+            BEGIN
+                IF pg_try_advisory_xact_lock(hashtext('known_endpoints_sync')) THEN
+                    INSERT INTO known_endpoints (id, name, host_id, host, last_seen)
+                    SELECT DISTINCT ON (endpoint_id) endpoint_id, name, host_id, host, last_seen
+                    FROM known_endpoints_insert
+                    ORDER BY endpoint_id, last_seen DESC
+                    ON CONFLICT (id) DO UPDATE SET
+                        last_seen = GREATEST(known_endpoints.last_seen, EXCLUDED.last_seen);
+                    
+                    DELETE FROM known_endpoints_insert;
+                END IF;
+            END $$;
+        ";
+
+        await using var cmd = new NpgsqlCommand(sql, conn);
+        await cmd.ExecuteNonQueryAsync(stoppingToken);
+    }
+}