Reuse one query plan per SQL Server ingestion statement (#6011)

* Add ColumnTypes method for parameter size handling and implement ingestion plan reuse tests

* Enhance BatchInsert to handle unbounded text parameters and add ingestion plan reuse tests

Author: warwickschroeder

src/ServiceControl.Audit.Persistence.EFCore/Infrastructure/BatchInsert.cs

@@ -35,7 +35,16 @@ public static async Task Rows<TEntity>(DbContext dbContext, IReadOnlyCollection<
             {
                 foreach (var column in table.Columns)
                 {
-                    command.Parameters.Add(column.TypeMapping.CreateParameter(command, $"@p{index++}", column.Read(row), column.IsNullable));
+                    var parameter = column.TypeMapping.CreateParameter(command, $"@p{index++}", column.Read(row), column.IsNullable);
+
+                    // SQL Server reuses a plan only when the declarations match, and EF declares unbounded text as
+                    // nvarchar(4000) or nvarchar(max) by the value's length. Npgsql ignores a size of -1.
+                    if (column.IsUnboundedText)
+                    {
+                        parameter.Size = -1;
+                    }
+
+                    command.Parameters.Add(parameter);
                 }
             }
 
@@ -124,5 +133,8 @@ .. entityType.GetProperties()
         static readonly ConcurrentDictionary<(IModel Model, Type EntityType), InsertTable> cache = new();
     }
 
-    sealed record InsertColumn(string Name, RelationalTypeMapping TypeMapping, bool IsNullable, Func<object, object?> Read);
+    sealed record InsertColumn(string Name, RelationalTypeMapping TypeMapping, bool IsNullable, Func<object, object?> Read)
+    {
+        public bool IsUnboundedText { get; } = (TypeMapping.Converter?.ProviderClrType ?? TypeMapping.ClrType) == typeof(string) && TypeMapping.Size is null;
+    }
 }

src/ServiceControl.Audit.Persistence.Tests.SqlServer/IngestionPlanReuseTests.cs

@@ -0,0 +1,82 @@
+namespace ServiceControl.Audit.Persistence.Tests
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Text;
+    using System.Threading.Tasks;
+    using Microsoft.Data.SqlClient;
+    using NUnit.Framework;
+    using ServiceControl.SagaAudit;
+
+    // SQL Server reuses a cached plan only when the parameter declarations match too, so these prove
+    // that an unbounded text column is always declared nvarchar(max), whether its value is null, short or long.
+    class IngestionPlanReuseTests : EFPersistenceTestFixture
+    {
+        [Test]
+        public async Task Messages_with_no_body_a_short_body_and_a_long_body_share_one_plan()
+        {
+            await Ingest(MakeMessage(contentType: "text/plain"));
+            await Ingest(MakeMessage(contentType: "text/plain"), Encoding.UTF8.GetBytes(new string('x', 10)));
+            await Ingest(MakeMessage(contentType: "text/plain"), Encoding.UTF8.GetBytes(new string('x', 4001)));
+
+            var plans = await CachedPlans("AuditMessages");
+
+            Assert.That(plans, Has.Count.EqualTo(1));
+            Assert.That(plans[0].UseCount, Is.EqualTo(3));
+            Assert.That(plans[0].Text, Does.Contain("nvarchar(max)").And.Not.Contain("nvarchar(4000)"));
+        }
+
+        [Test]
+        public async Task Saga_snapshots_with_no_state_a_short_state_and_a_long_state_share_one_plan()
+        {
+            await RecordSnapshot(state: null);
+            await RecordSnapshot(state: new string('x', 10));
+            await RecordSnapshot(state: new string('x', 4001));
+
+            var plans = await CachedPlans("SagaSnapshots");
+
+            Assert.That(plans, Has.Count.EqualTo(1));
+            Assert.That(plans[0].UseCount, Is.EqualTo(3));
+            Assert.That(plans[0].Text, Does.Contain("nvarchar(max)").And.Not.Contain("nvarchar(4000)"));
+        }
+
+        async Task RecordSnapshot(string state)
+        {
+            await using var unitOfWork = await StartAuditUnitOfWork(1);
+            await unitOfWork.RecordSagaSnapshot(new SagaSnapshot
+            {
+                SagaId = Guid.NewGuid(),
+                SagaType = "Shipping.ShippingPolicy",
+                Status = SagaStateChangeStatus.Updated,
+                StartTime = TimeSent,
+                FinishTime = TimeSent,
+                StateAfterChange = state
+            });
+            await unitOfWork.Complete();
+        }
+
+        async Task<List<(string Text, int UseCount)>> CachedPlans(string table)
+        {
+            await using var connection = new SqlConnection(configuration.ConnectionString);
+            await connection.OpenAsync();
+
+            await using var command = connection.CreateCommand();
+            command.CommandText = """
+                SELECT st.text, cp.usecounts
+                FROM sys.dm_exec_cached_plans AS cp
+                CROSS APPLY sys.dm_exec_sql_text(cp.plan_handle) AS st
+                WHERE CHARINDEX(@statement, st.text) > 0
+                """;
+            command.Parameters.AddWithValue("@statement", $"INSERT INTO [{configuration.Schema}].[{table}]");
+
+            var plans = new List<(string Text, int UseCount)>();
+            await using var reader = await command.ExecuteReaderAsync();
+            while (await reader.ReadAsync())
+            {
+                plans.Add((reader.GetString(0), reader.GetInt32(1)));
+            }
+
+            return plans;
+        }
+    }
+}

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

@@ -15,7 +15,31 @@ abstract class SqlServerDialect
     /// </summary>
     protected static string Table<TEntity>(ServiceControlDbContext dbContext) => SchemaQualifiedTableName.For<TEntity>(dbContext);
 
-    protected static async Task Execute(ServiceControlDbContext dbContext, string sql, IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
+    /// <summary>
+    /// How to declare each column of a statement's rows, taken from the properties the columns are
+    /// mapped to. Pass the property names in the order the values sit in each row. A row may be
+    /// shorter than this list, but not longer. A text column gets its length as its size, or -1 for nvarchar(max).
+    /// </summary>
+    protected static ColumnType[] ColumnTypes<TEntity>(ServiceControlDbContext dbContext, params string[] properties)
+    {
+        var entityType = dbContext.Model.FindEntityType(typeof(TEntity))
+            ?? throw new InvalidOperationException($"{typeof(TEntity).Name} is not part of the model.");
+
+        return [.. properties.Select(name =>
+        {
+            var property = entityType.FindProperty(name)
+                ?? throw new InvalidOperationException($"{typeof(TEntity).Name} has no property {name}.");
+
+            var dbType = property.GetRelationalTypeMapping().DbType
+                ?? throw new InvalidOperationException($"{typeof(TEntity).Name}.{name} has no DbType in its type mapping.");
+
+            return new ColumnType(dbType, dbType == DbType.String ? property.GetMaxLength() ?? -1 : null);
+        })];
+    }
+
+    protected readonly record struct ColumnType(DbType DbType, int? Size);
+
+    protected static async Task Execute(ServiceControlDbContext dbContext, string sql, ColumnType[] columnTypes, IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
     {
         await using var command = dbContext.Database.GetDbConnection().CreateCommand();
         command.Transaction = (dbContext.Database.CurrentTransaction
@@ -25,18 +49,36 @@ protected static async Task Execute(ServiceControlDbContext dbContext, string sq
         var index = 0;
         foreach (var row in rows)
         {
-            foreach (var value in row)
+            for (var column = 0; column < row.Length; column++)
             {
+                var value = row[column];
                 var parameter = command.CreateParameter();
                 parameter.ParameterName = $"@p{index++}";
                 parameter.Value = value ?? DBNull.Value;
 
-                // Attempt de-duplication compares LastAttemptedAt for equality, so datetime
-                // parameters must keep datetime2 precision instead of the datetime default.
+                // Attempt de-duplication compares LastAttemptedAt for equality, so a datetime keeps
+                // datetime2 precision wherever it sits, including the lone "now" that ResolveRetriedMessages sends first.
                 if (value is DateTime)
                 {
                     parameter.DbType = DbType.DateTime2;
                 }
+                else
+                {
+                    // SQL Server reuses a plan only when the declarations match, so every value, null or not,
+                    // is declared by its column. A string longer than its column keeps its own length,
+                    // because SqlClient would silently cut it and SQL Server must refuse it instead.
+                    var (dbType, size) = columnTypes[column];
+                    if (value is string && dbType != DbType.String)
+                    {
+                        throw new InvalidOperationException($"Parameter {parameter.ParameterName} holds text, but its column is declared as {dbType}.");
+                    }
+
+                    parameter.DbType = dbType;
+                    if (size is { } columnSize)
+                    {
+                        parameter.Size = value is string text && columnSize != -1 && text.Length > columnSize ? text.Length : columnSize;
+                    }
+                }
 
                 command.Parameters.Add(parameter);
             }

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

@@ -1,9 +1,5 @@
 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;
@@ -29,6 +25,12 @@ await Execute(
                  WHEN NOT MATCHED THEN INSERT (NormalizedName, ThroughputSource, DateUtc, MessageCount)
                  VALUES (s.NormalizedEndpointName, s.Source, s.Day, s.Count);
                  """,
+                ColumnTypes<LicensingEndpointThroughputEntity>(
+                    dbContext,
+                    nameof(LicensingEndpointThroughputEntity.NormalizedName),
+                    nameof(LicensingEndpointThroughputEntity.ThroughputSource),
+                    nameof(LicensingEndpointThroughputEntity.DateUtc),
+                    nameof(LicensingEndpointThroughputEntity.MessageCount)),
                 [[normalizedName, (int)throughputSource, date, messageCount]],
                 cancellationToken);
         }

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

@@ -29,6 +29,7 @@ await Execute(
                  WHEN NOT MATCHED THEN INSERT ({FailedMessageColumnList})
                  VALUES ({FailedMessageSourceColumnList});
                  """,
+                ColumnTypes<FailedMessageEntity>(dbContext, FailedMessageColumns),
                 chunk.Select(FailedMessageValues),
                 cancellationToken);
         }
@@ -50,6 +51,12 @@ await Execute(
                  WHEN NOT MATCHED THEN INSERT ([FailedMessageUniqueId], [GroupId], [Title], [Type])
                  VALUES (s.[FailedMessageUniqueId], s.[GroupId], s.[Title], s.[Type]);
                  """,
+                ColumnTypes<FailedMessageGroupEntity>(
+                    dbContext,
+                    nameof(FailedMessageGroupEntity.FailedMessageUniqueId),
+                    nameof(FailedMessageGroupEntity.GroupId),
+                    nameof(FailedMessageGroupEntity.Title),
+                    nameof(FailedMessageGroupEntity.Type)),
                 chunk.Select(group => new object?[] { group.FailedMessageUniqueId, group.GroupId, group.Title, group.Type }),
                 cancellationToken);
         }
@@ -71,6 +78,13 @@ await Execute(
                  WHEN NOT MATCHED THEN INSERT ([Id], [Name], [HostId], [Host], [Monitored])
                  VALUES (s.[Id], s.[Name], s.[HostId], s.[Host], s.[Monitored]);
                  """,
+                ColumnTypes<KnownEndpointEntity>(
+                    dbContext,
+                    nameof(KnownEndpointEntity.Id),
+                    nameof(KnownEndpointEntity.Name),
+                    nameof(KnownEndpointEntity.HostId),
+                    nameof(KnownEndpointEntity.Host),
+                    nameof(KnownEndpointEntity.Monitored)),
                 chunk.Select(endpoint => new object?[] { endpoint.Id, endpoint.Name, endpoint.HostId, endpoint.Host, endpoint.Monitored }),
                 cancellationToken);
         }
@@ -98,7 +112,11 @@ INNER JOIN (VALUES
                      ON t.[UniqueMessageId] = s.[UniqueMessageId]
                  WHERE t.[LastAttemptedAt] <= s.[SucceededAt];
                  """,
-                [[now], .. chunk.Select(retry => new object?[] { retry.UniqueMessageId, retry.SucceededAt })],
+                ColumnTypes<FailedMessageEntity>(
+                    dbContext,
+                    nameof(FailedMessageEntity.UniqueMessageId),
+                    nameof(FailedMessageEntity.LastAttemptedAt)),
+                    [[now], .. chunk.Select(retry => new object?[] { retry.UniqueMessageId, retry.SucceededAt })],
                 cancellationToken);
         }
     }

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

@@ -22,6 +22,11 @@ await Execute(
                  WHEN NOT MATCHED THEN INSERT ([UniqueMessageId], [RetryBatchId], [StageAttempts])
                  VALUES (s.[UniqueMessageId], s.[RetryBatchId], s.[StageAttempts]);
                  """,
+                ColumnTypes<FailedMessageRetryEntity>(
+                    dbContext,
+                    nameof(FailedMessageRetryEntity.UniqueMessageId),
+                    nameof(FailedMessageRetryEntity.RetryBatchId),
+                    nameof(FailedMessageRetryEntity.StageAttempts)),
                 chunk.Select(retry => new object?[] { retry.UniqueMessageId, retry.RetryBatchId, retry.StageAttempts }),
                 cancellationToken);
         }