Removed unnecessary batching boundaries from topological sort (#27713)

Closes #20664

Author: roji

src/EFCore.Relational/Update/Internal/CommandBatchPreparer.cs

@@ -16,7 +16,7 @@ public class CommandBatchPreparer : ICommandBatchPreparer
 {
     private readonly int _minBatchSize;
     private readonly bool _sensitiveLoggingEnabled;
-    private readonly Multigraph<IReadOnlyModificationCommand, IAnnotatable> _modificationCommandGraph = new();
+    private readonly Multigraph<IReadOnlyModificationCommand, IAnnotatable> _modificationCommandGraph;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -29,6 +29,8 @@ public CommandBatchPreparer(CommandBatchPreparerDependencies dependencies)
         _minBatchSize =
             dependencies.Options.Extensions.OfType<RelationalOptionsExtension>().FirstOrDefault()?.MinBatchSize
             ?? 1;
+
+        _modificationCommandGraph = new(dependencies.ModificationCommandComparer);
         Dependencies = dependencies;
 
         if (dependencies.LoggingOptions.IsSensitiveDataLoggingEnabled)
@@ -57,16 +59,14 @@ public CommandBatchPreparer(CommandBatchPreparerDependencies dependencies)
     {
         var parameterNameGenerator = Dependencies.ParameterNameGeneratorFactory.Create();
         var commands = CreateModificationCommands(entries, updateAdapter, parameterNameGenerator.GenerateNext);
-        var sortedCommandSets = TopologicalSort(commands);
+        var commandSets = TopologicalSort(commands);
 
-        for (var commandSetIndex = 0; commandSetIndex < sortedCommandSets.Count; commandSetIndex++)
+        for (var commandSetIndex = 0; commandSetIndex < commandSets.Count; commandSetIndex++)
         {
-            var independentCommandSet = sortedCommandSets[commandSetIndex];
-
-            independentCommandSet.Sort(Dependencies.ModificationCommandComparer);
+            var commandSet = commandSets[commandSetIndex];
 
             var batch = Dependencies.ModificationCommandBatchFactory.Create();
-            foreach (var modificationCommand in independentCommandSet)
+            foreach (var modificationCommand in commandSet)
             {
                 (modificationCommand as ModificationCommand)?.AssertColumnsNotInitialized();
                 if (modificationCommand.EntityState == EntityState.Modified
@@ -108,7 +108,7 @@ public CommandBatchPreparer(CommandBatchPreparerDependencies dependencies)
                 }
             }
 
-            var hasMoreCommandSets = commandSetIndex < sortedCommandSets.Count - 1;
+            var hasMoreCommandSets = commandSetIndex < commandSets.Count - 1;
 
             if (batch.ModificationCommands.Count == 1
                 || batch.ModificationCommands.Count >= _minBatchSize)
@@ -295,10 +295,10 @@ private string FormatCycle(
         var builder = new StringBuilder();
         for (var i = 0; i < data.Count; i++)
         {
-            var (command1, command2, annotatables) = data[i];
+            var (command1, command2, edges) = data[i];
             Format(command1, builder);
 
-            switch (annotatables.First())
+            switch (edges.First())
             {
                 case IForeignKeyConstraint foreignKey:
                     Format(foreignKey, command1, command2, builder);
@@ -536,10 +536,40 @@ private void AddForeignKeyEdges(
 
                     var dependentKeyValue = ((ForeignKeyConstraint)foreignKey).GetRowForeignKeyValueFactory()
                         .CreateDependentValueIndex(command);
-                    if (dependentKeyValue != null)
+
+                    if (dependentKeyValue is null || !predecessorsMap.TryGetValue(dependentKeyValue, out var predecessorCommands))
                     {
-                        AddMatchingPredecessorEdge(
-                            predecessorsMap, dependentKeyValue, commandGraph, command, foreignKey);
+                        continue;
+                    }
+
+                    foreach (var predecessor in predecessorCommands)
+                    {
+                        if (predecessor != command)
+                        {
+                            // If we're adding/inserting a dependent where the principal key is being database-generated, then
+                            // the dependency edge represents a batching boundary: fetch the principal database-generated
+                            // property from the database in separate batch, in order to populate the dependent foreign key
+                            // property in the next.
+                            var requiresBatchingBoundary = false;
+
+                            for (var i = 0; i < foreignKey.PrincipalColumns.Count; i++)
+                            {
+                                for (var j = 0; j < predecessor.Entries.Count; j++)
+                                {
+                                    var entry = predecessor.Entries[j];
+
+                                    if (foreignKey.PrincipalColumns[i].FindColumnMapping(entry.EntityType) is IColumnMapping columnMapping
+                                        && entry.IsStoreGenerated(columnMapping.Property))
+                                    {
+                                        requiresBatchingBoundary = true;
+                                        goto AfterLoop;
+                                    }
+                                }
+                            }
+                            AfterLoop:
+                            
+                            commandGraph.AddEdge(predecessor, command, foreignKey, requiresBatchingBoundary);
+                        }
                     }
                 }
             }
@@ -623,7 +653,7 @@ private static void AddMatchingPredecessorEdge<T>(
         T keyValue,
         Multigraph<IReadOnlyModificationCommand, IAnnotatable> commandGraph,
         IReadOnlyModificationCommand command,
-        IAnnotatable edge)
+        IAnnotatable edgeAnnotatable)
         where T : notnull
     {
         if (predecessorsMap.TryGetValue(keyValue, out var predecessorCommands))
@@ -632,7 +662,7 @@ private static void AddMatchingPredecessorEdge<T>(
             {
                 if (predecessor != command)
                 {
-                    commandGraph.AddEdge(predecessor, command, edge);
+                    commandGraph.AddEdge(predecessor, command, edgeAnnotatable);
                 }
             }
         }

src/EFCore.SqlServer/Update/Internal/SqlServerModificationCommandBatch.cs

@@ -18,7 +18,6 @@ public class SqlServerModificationCommandBatch : AffectedCountModificationComman
     private const int DefaultNetworkPacketSizeBytes = 4096;
     private const int MaxScriptLength = 65536 * DefaultNetworkPacketSizeBytes / 2;
     private const int MaxParameterCount = 2100;
-    private readonly int _maxBatchSize;
     private readonly List<IReadOnlyModificationCommand> _pendingBulkInsertCommands = new();
 
     /// <summary>
@@ -31,7 +30,7 @@ public SqlServerModificationCommandBatch(
         ModificationCommandBatchFactoryDependencies dependencies,
         int maxBatchSize)
         : base(dependencies)
-        => _maxBatchSize = maxBatchSize;
+        => MaxBatchSize = maxBatchSize;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -48,8 +47,7 @@ public SqlServerModificationCommandBatch(
     /// <remarks>
     ///     For SQL Server, this is 42 by default, and cannot exceed 1000.
     /// </remarks>
-    protected override int MaxBatchSize
-        => _maxBatchSize;
+    protected override int MaxBatchSize { get; }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -165,8 +163,8 @@ protected override void AddCommand(IReadOnlyModificationCommand modificationComm
     private static bool CanBeInsertedInSameStatement(
         IReadOnlyModificationCommand firstCommand,
         IReadOnlyModificationCommand secondCommand)
-        => string.Equals(firstCommand.TableName, secondCommand.TableName, StringComparison.Ordinal)
-            && string.Equals(firstCommand.Schema, secondCommand.Schema, StringComparison.Ordinal)
+        => firstCommand.TableName == secondCommand.TableName
+            && firstCommand.Schema == secondCommand.Schema
             && firstCommand.ColumnModifications.Where(o => o.IsWrite).Select(o => o.ColumnName).SequenceEqual(
                 secondCommand.ColumnModifications.Where(o => o.IsWrite).Select(o => o.ColumnName))
             && firstCommand.ColumnModifications.Where(o => o.IsRead).Select(o => o.ColumnName).SequenceEqual(

src/EFCore/Metadata/IReadOnlyEntityType.cs

@@ -125,7 +125,7 @@ IEnumerable<IReadOnlyEntityType> GetDerivedTypesInclusive()
         => new[] { this }.Concat(GetDerivedTypes());
 
     /// <summary>
-    ///     Gets all types in the model that directly derive from a given entity type.
+    ///     Gets all types in the model that directly derive from a given entity type, in a deterministic top-to-bottom ordering.
     /// </summary>
     /// <returns>The derived types.</returns>
     IEnumerable<IReadOnlyEntityType> GetDirectlyDerivedTypes();

src/EFCore/Metadata/Internal/EntityType.cs

@@ -427,9 +427,8 @@ private void UpdateBaseTypeConfigurationSource(ConfigurationSource configuration
     ///     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>
-    // Note this is ISet because there is no suitable readonly interface in the profiles we are using
     [DebuggerStepThrough]
-    public virtual ISet<EntityType> GetDirectlyDerivedTypes()
+    public virtual IReadOnlySet<EntityType> GetDirectlyDerivedTypes()
         => _directlyDerivedTypes;
 
     /// <summary>

src/EFCore/Metadata/Internal/ModelExtensions.cs

@@ -1,6 +1,8 @@
 // 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;
+
 namespace Microsoft.EntityFrameworkCore.Metadata.Internal;
 
 /// <summary>
@@ -36,18 +38,25 @@ public static IEnumerable<IEntityType> GetRootEntityTypes(this IModel model)
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public static IEnumerable<IEntityType> GetEntityTypesInHierarchicalOrder(this IModel model)
-        => Sort(model.GetEntityTypes());
-
-    private static IEnumerable<IEntityType> Sort(IEnumerable<IEntityType> entityTypes)
     {
-        var entityTypeGraph = new Multigraph<IEntityType, int>();
-        entityTypeGraph.AddVertices(entityTypes);
-        foreach (var entityType in entityTypes.Where(et => et.BaseType != null))
+        var entityTypes = new Queue<IEntityType>();
+
+        foreach (var root in model.GetRootEntityTypes())
         {
-            entityTypeGraph.AddEdge(entityType.BaseType!, entityType, 0);
+            entityTypes.Enqueue(root);
         }
 
-        return entityTypeGraph.BatchingTopologicalSort().SelectMany(b => b.OrderBy(et => et.Name));
+        while (entityTypes.Count > 0)
+        {
+            var current = entityTypes.Dequeue();
+
+            yield return current;
+
+            foreach (var descendant in current.GetDirectlyDerivedTypes())
+            {
+                entityTypes.Enqueue(descendant);
+            }
+        }
     }
 
     /// <summary>