Add Cosmos bulk execution (#37033)

Adds BulkExecutionEnabled to CosmosDbContextOptionsBuilder
Execute non batchable document writes in parallel when bulk execution enabled
Add a warning exception when batching is enabled when using bulk execution

Fixes #36599

Author: JoasE

src/EFCore.Cosmos/Diagnostics/AutoTransactionBehaviorEventData.cs

@@ -0,0 +1,32 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Cosmos.Diagnostics;
+
+/// <summary>
+///     A <see cref="DiagnosticSource" /> event payload class for Cosmos events related to AutoTransactionBehavior
+/// </summary>
+/// <remarks>
+///     See <see href="https://aka.ms/efcore-docs-diagnostics">Logging, events, and diagnostics</see> for more information and examples.
+/// </remarks>
+public class AutoTransactionBehaviorEventData : EventData
+{
+    /// <summary>
+    ///     Constructs the event payload.
+    /// </summary>
+    /// <param name="eventDefinition">The event definition.</param>
+    /// <param name="messageGenerator">A delegate that generates a log message for this event.</param>
+    /// <param name="autoTransactionBehavior">The AutoTransactionBehavior that was used.</param>
+    public AutoTransactionBehaviorEventData(
+        EventDefinitionBase eventDefinition,
+        Func<EventDefinitionBase, EventData, string> messageGenerator,
+        AutoTransactionBehavior autoTransactionBehavior) : base(eventDefinition, messageGenerator)
+    {
+        AutoTransactionBehavior = autoTransactionBehavior;
+    }
+
+    /// <summary>
+    ///     The AutoTransactionBehavior that was used.
+    /// </summary>
+    public virtual AutoTransactionBehavior AutoTransactionBehavior { get; }
+}

src/EFCore.Cosmos/Diagnostics/CosmosEventId.cs

@@ -40,6 +40,7 @@ private enum Id
 
         // Update events
         PrimaryKeyValueNotSet = CoreEventId.ProviderBaseId + 200,
+        BulkExecutionWithTransactionalBatch,
 
         // Model validation events
         NoPartitionKeyDefined = CoreEventId.ProviderBaseId + 600,
@@ -193,4 +194,19 @@ private static EventId MakeUpdateId(Id id)
     ///     </para>
     /// </remarks>
     public static readonly EventId PrimaryKeyValueNotSet = MakeUpdateId(Id.PrimaryKeyValueNotSet);
+
+    /// <summary>
+    ///     SaveChanges was invoked with both bulk execution and batching being enabled. Transactional batches can not be run in bulk thus they
+    ///     will skip bulk execution. Use AutoTransactionBehavior.Never to leverage bulk execution. If batching was intended, suppress this warning
+    ///     using <c>DbContextOptionsBuilder.ConfigureWarnings(w => w.Ignore(CosmosEventId.BulkExecutionWithTransactionalBatch))</c>
+    /// </summary>
+    /// <remarks>
+    ///     <para>
+    ///         This event is in the <see cref="DbLoggerCategory.Update" /> category.
+    ///     </para>
+    ///     <para>
+    ///         This event uses the <see cref="Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.AutoTransactionBehaviorEventData" /> payload when used with a <see cref="DiagnosticSource" />.
+    ///     </para>
+    /// </remarks>
+    public static readonly EventId BulkExecutionWithTransactionalBatch = MakeUpdateId(Id.BulkExecutionWithTransactionalBatch);
 }

src/EFCore.Cosmos/Diagnostics/CosmosTransactionalBatchExecutedEventData.cs

@@ -1,13 +1,10 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Collections.Generic;
-using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
-
 namespace Microsoft.EntityFrameworkCore.Diagnostics;
 
 /// <summary>
-///     A <see cref="DiagnosticSource" /> event payload class for Cosmos item command executed events.
+///     A <see cref="DiagnosticSource" /> event payload class for Cosmos transactional batch executed events.
 /// </summary>
 /// <remarks>
 ///     See <see href="https://aka.ms/efcore-docs-diagnostics">Logging, events, and diagnostics</see> for more information and examples.

src/EFCore.Cosmos/Diagnostics/Internal/CosmosLoggerExtensions.cs

@@ -547,6 +547,33 @@ public static void PrimaryKeyValueNotSet(
         }
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public static void BulkExecutionWithTransactionalBatch(
+        this IDiagnosticsLogger<DbLoggerCategory.Update> diagnostics,
+        AutoTransactionBehavior autoTransactionBehavior)
+    {
+        var definition = CosmosResources.LogBulkExecutionWithTransactionalBatch(diagnostics);
+
+        if (diagnostics.ShouldLog(definition))
+        {
+            definition.Log(diagnostics);
+        }
+
+        if (diagnostics.NeedsEventData(definition, out var diagnosticSourceEnabled, out var simpleLogEnabled))
+        {
+            var eventData = new AutoTransactionBehaviorEventData(
+                definition,
+                (d, b) => ((EventDefinition)d).GenerateMessage(),
+                autoTransactionBehavior);
+            diagnostics.DispatchEventData(definition, eventData, diagnosticSourceEnabled, simpleLogEnabled);
+        }
+    }
+
     private static string FormatParameters(IReadOnlyList<(string Name, object? Value)> parameters, bool shouldLogParameterValues)
         => FormatParameters(parameters.Select(p => new SqlParameter(p.Name, p.Value)).ToList(), shouldLogParameterValues);
 

src/EFCore.Cosmos/Diagnostics/Internal/CosmosLoggingDefinitions.cs

@@ -90,4 +90,12 @@ public class CosmosLoggingDefinitions : LoggingDefinitions
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public EventDefinitionBase? LogPrimaryKeyValueNotSet;
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public EventDefinitionBase? LogBulkExecutionWithTransactionalBatch;
 }

src/EFCore.Cosmos/Extensions/CosmosDbContextOptionsExtensions.cs

@@ -52,6 +52,8 @@ public static DbContextOptionsBuilder UseCosmos(
         Check.NotNull(optionsBuilder);
         Check.NotNull(cosmosOptionsAction);
 
+        ConfigureWarnings(optionsBuilder);
+
         cosmosOptionsAction.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));
 
         return optionsBuilder;
@@ -118,6 +120,8 @@ public static DbContextOptionsBuilder UseCosmos(
             .WithAccountKey(accountKey)
             .WithDatabaseName(databaseName);
 
+        ConfigureWarnings(optionsBuilder);
+
         ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
 
         cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));
@@ -186,6 +190,8 @@ public static DbContextOptionsBuilder UseCosmos(
             .WithTokenCredential(tokenCredential)
             .WithDatabaseName(databaseName);
 
+        ConfigureWarnings(optionsBuilder);
+
         ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
 
         cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));
@@ -247,6 +253,8 @@ public static DbContextOptionsBuilder UseCosmos(
             .WithConnectionString(connectionString)
             .WithDatabaseName(databaseName);
 
+        ConfigureWarnings(optionsBuilder);
+
         ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
 
         cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));
@@ -260,6 +268,10 @@ var coreOptionsExtension
             = optionsBuilder.Options.FindExtension<CoreOptionsExtension>()
             ?? new CoreOptionsExtension();
 
+        coreOptionsExtension = coreOptionsExtension.WithWarningsConfiguration(
+            coreOptionsExtension.WarningsConfiguration.TryWithExplicit(
+                CosmosEventId.BulkExecutionWithTransactionalBatch, WarningBehavior.Throw));
+
         ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(coreOptionsExtension);
     }
 }

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -229,6 +229,22 @@ public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool
     public virtual CosmosDbContextOptionsBuilder SessionTokenManagementMode(SessionTokenManagementMode mode)
         => WithOption(e => e.WithSessionTokenManagementMode(mode));
 
+    /// <summary>
+    ///     Sets the boolean to enable the <see href="https://learn.microsoft.com/en-us/azure/cosmos-db/bulk-executor-overview">Cosmos DB SDK bulk execution feature</see>.
+    ///     Enabling this feature can improve throughput for small write operations but may increase latency.
+    ///     It is recommended only for high-throughput, non-latency-sensitive scenarios.
+    ///     Because <see href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/transactional-batch">Transactional Batches</see> cannot be executed in bulk mode,
+    ///     any operations batched by EF will not use bulk execution. To ensure operations are executed in bulk, consider disabling batching by setting <see cref="AutoTransactionBehavior.Never"/>.
+    ///     For more information, see <see href="https://learn.microsoft.com/en-us/ef/core/providers/cosmos/saving">Saving Data - Azure Cosmos DB Provider</see>.
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
+    ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <param name="enabled"><see langword="true" /> to enable the Cosmos DB SDK bulk feature.</param>
+    public virtual CosmosDbContextOptionsBuilder BulkExecutionEnabled(bool enabled = true)
+        => WithOption(e => e.BulkExecutionEnabled(enabled));
+
     /// <summary>
     ///     Sets an option by cloning the extension used to store the settings. This ensures the builder
     ///     does not modify options that are already in use elsewhere.

src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs

@@ -37,6 +37,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension
     private DbContextOptionsExtensionInfo? _info;
     private Func<HttpClient>? _httpClientFactory;
     private SessionTokenManagementMode _sessionTokenManagementMode = SessionTokenManagementMode.FullyAutomatic;
+    private bool? _enableBulkExecution;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -75,6 +76,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom)
         _maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection;
         _httpClientFactory = copyFrom._httpClientFactory;
         _sessionTokenManagementMode = copyFrom._sessionTokenManagementMode;
+        _enableBulkExecution = copyFrom._enableBulkExecution;
     }
 
     /// <summary>
@@ -590,6 +592,29 @@ public virtual CosmosOptionsExtension WithSessionTokenManagementMode(SessionToke
         return clone;
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual bool? EnableBulkExecution => _enableBulkExecution;
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual CosmosOptionsExtension BulkExecutionEnabled(bool enabled)
+    {
+        var clone = Clone();
+
+        clone._enableBulkExecution = enabled;
+
+        return clone;
+    }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -659,6 +684,7 @@ public override int GetServiceProviderHashCode()
                 hashCode.Add(Extension._maxRequestsPerTcpConnection);
                 hashCode.Add(Extension._httpClientFactory);
                 hashCode.Add(Extension._sessionTokenManagementMode);
+                hashCode.Add(Extension._enableBulkExecution);
 
                 _serviceProviderHash = hashCode.ToHashCode();
             }
@@ -684,7 +710,8 @@ public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo
                 && Extension._maxTcpConnectionsPerEndpoint == otherInfo.Extension._maxTcpConnectionsPerEndpoint
                 && Extension._maxRequestsPerTcpConnection == otherInfo.Extension._maxRequestsPerTcpConnection
                 && Extension._httpClientFactory == otherInfo.Extension._httpClientFactory
-                && Extension._sessionTokenManagementMode == otherInfo.Extension._sessionTokenManagementMode;
+                && Extension._sessionTokenManagementMode == otherInfo.Extension._sessionTokenManagementMode
+                && Extension._enableBulkExecution == otherInfo.Extension._enableBulkExecution;
 
         public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
         {

src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs

@@ -157,7 +157,7 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public virtual SessionTokenManagementMode SessionTokenManagementMode { get; private set; }
+    public virtual bool? EnableBulkExecution { get; private set; }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -186,6 +186,7 @@ public virtual void Initialize(IDbContextOptions options)
             MaxTcpConnectionsPerEndpoint = cosmosOptions.MaxTcpConnectionsPerEndpoint;
             MaxRequestsPerTcpConnection = cosmosOptions.MaxRequestsPerTcpConnection;
             HttpClientFactory = cosmosOptions.HttpClientFactory;
+            EnableBulkExecution = cosmosOptions.EnableBulkExecution;
         }
     }
 
@@ -216,6 +217,7 @@ public virtual void Validate(IDbContextOptions options)
                 || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint
                 || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection
                 || HttpClientFactory != cosmosOptions.HttpClientFactory
+                || EnableBulkExecution != cosmosOptions.EnableBulkExecution
             ))
         {
             throw new InvalidOperationException(

src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs

@@ -155,4 +155,12 @@ public interface ICosmosSingletonOptions : ISingletonOptions
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     Func<HttpClient>? HttpClientFactory { get; }
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    bool? EnableBulkExecution { get; }
 }