Cosmos: flow transactional batch failures through the execution strategy so they can be retried

src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs

@@ -713,17 +713,12 @@ public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string
     ///     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 Task<CosmosTransactionalBatchResult> ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default)
-        => _executionStrategy.ExecuteAsync((batch, sessionTokenStorage, this), ExecuteTransactionalBatchOnceAsync, null, cancellationToken);
-
-    private static async Task<CosmosTransactionalBatchResult> ExecuteTransactionalBatchOnceAsync(DbContext _,
-        (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters,
+    public virtual async Task ExecuteTransactionalBatchAsync(
+        ICosmosTransactionalBatchWrapper batch,
+        ISessionTokenStorage sessionTokenStorage,
         CancellationToken cancellationToken = default)
     {
-        var batch = parameters.Batch;
         var transactionalBatch = batch.GetTransactionalBatch();
-        var wrapper = parameters.Wrapper;
-        var sessionTokenStorage = parameters.SessionTokenStorage;
 
         var options = new TransactionalBatchRequestOptions
         {
@@ -732,15 +727,15 @@ private static async Task<CosmosTransactionalBatchResult> ExecuteTransactionalBa
 
         using var response = await transactionalBatch.ExecuteAsync(options, cancellationToken).ConfigureAwait(false);
 
-        wrapper._commandLogger.ExecutedTransactionalBatch(
+        _commandLogger.ExecutedTransactionalBatch(
             response.Diagnostics.GetClientElapsedTime(),
             response.Headers.RequestCharge,
             response.Headers.ActivityId,
             batch.CollectionId,
             batch.PartitionKeyValue,
             "[ \"" + string.Join("\", \"", batch.Entries.Select(x => x.Id)) + "\" ]");
 
-        return ProcessBatchResponse(batch.CollectionId, response, batch.Entries, sessionTokenStorage);
+        ProcessBatchResponse(batch.CollectionId, response, batch.Entries, sessionTokenStorage);
     }
 
     private static ItemRequestOptions CreateItemRequestOptions(IUpdateEntry entry, string? sessionToken)
@@ -821,21 +816,14 @@ private static void TryTrackSessionTokenFromFailure(string containerId, HttpStat
         }
     }
 
-    private static CosmosTransactionalBatchResult ProcessBatchResponse(string containerId, TransactionalBatchResponse response, IReadOnlyList<CosmosTransactionalBatchEntry> entries, ISessionTokenStorage sessionTokenStorage)
+    private static void ProcessBatchResponse(string containerId, TransactionalBatchResponse response, IReadOnlyList<CosmosTransactionalBatchEntry> entries, ISessionTokenStorage sessionTokenStorage)
     {
         if (!response.IsSuccessStatusCode)
         {
             TryTrackSessionTokenFromFailure(containerId, response.StatusCode, response.Headers, sessionTokenStorage);
 
-            var errorCode = response.StatusCode;
-            var errorEntries = response
-                .Select((opResult, index) => (opResult, index))
-                .Where(r => r.opResult.StatusCode == errorCode)
-                .Select(r => entries[r.index].Entry)
-                .ToList();
-
-            var exception = new CosmosException(response.ErrorMessage, errorCode, 0, response.ActivityId, response.RequestCharge);
-            return new CosmosTransactionalBatchResult(errorEntries, exception);
+            throw new CosmosException(
+                response.ErrorMessage, response.StatusCode, 0, response.ActivityId, response.RequestCharge);
         }
 
         sessionTokenStorage.TrackSessionToken(containerId, response.Headers.Session);
@@ -847,8 +835,6 @@ private static CosmosTransactionalBatchResult ProcessBatchResponse(string contai
 
             ProcessWriteResponse(entry.Entry, item.ETag, item.ResourceStream);
         }
-
-        return CosmosTransactionalBatchResult.Success;
     }
 
     private static void ProcessWriteResponse(IUpdateEntry entry, string eTag, Stream? content)

src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs

@@ -24,6 +24,7 @@ public class CosmosDatabaseWrapper : Database, IResettableService
     private readonly Dictionary<IEntityType, DocumentSource> _documentCollections = new();
 
     private readonly ICosmosClientWrapper _cosmosClient;
+    private readonly IExecutionStrategy _executionStrategy;
     private readonly bool _sensitiveLoggingEnabled;
     private readonly bool _bulkExecutionEnabled;
 
@@ -39,13 +40,15 @@ public CosmosDatabaseWrapper(
         DatabaseDependencies dependencies,
         ICurrentDbContext currentDbContext,
         ICosmosClientWrapper cosmosClient,
+        IExecutionStrategy executionStrategy,
         ICosmosSingletonOptions cosmosSingletonOptions,
         ISessionTokenStorageFactory sessionTokenStorageFactory,
         ILoggingOptions loggingOptions)
         : base(dependencies)
     {
         _currentDbContext = currentDbContext;
         _cosmosClient = cosmosClient;
+        _executionStrategy = executionStrategy;
         _bulkExecutionEnabled = cosmosSingletonOptions.EnableBulkExecution == true;
         SessionTokenStorage = sessionTokenStorageFactory.Create(currentDbContext.Context);
 
@@ -108,46 +111,89 @@ public override async Task<int> SaveChangesAsync(
             }
         }
 
-        foreach (var batch in groups.BatchableUpdateEntries)
+        var batchableUpdateEntries = groups.BatchableUpdateEntries as IReadOnlyList<(Grouping Key, List<CosmosUpdateEntry> UpdateEntries)>
+            ?? groups.BatchableUpdateEntries.ToList();
+
+        if (batchableUpdateEntries.Count > 0)
+        {
+            // The execution strategy is invoked around the whole set of batches (rather than around each individual
+            // batch) so that a transient failure retries the remaining work, while skipping the operations that were
+            // already committed by a previous attempt.
+            var state = new BatchExecutionState();
+            rowsAffected += await _executionStrategy.ExecuteAsync(
+                (Database: this, Batches: batchableUpdateEntries, State: state),
+                static (_, s, ct) => s.Database.ExecuteBatchesAsync(s.Batches, s.State, ct),
+                verifySucceeded: null,
+                cancellationToken).ConfigureAwait(false);
+        }
+
+        return rowsAffected;
+    }
+
+    private async Task<int> ExecuteBatchesAsync(
+        IReadOnlyList<(Grouping Key, List<CosmosUpdateEntry> UpdateEntries)> batches,
+        BatchExecutionState state,
+        CancellationToken cancellationToken)
+    {
+        var operationIndex = 0;
+        foreach (var batch in batches)
         {
             if (batch.UpdateEntries.Count == 1 && _currentDbContext.Context.Database.AutoTransactionBehavior != AutoTransactionBehavior.Always)
             {
+                // Skip the operations that were already committed by a previous execution strategy attempt.
+                if (operationIndex < state.CommittedOperations)
+                {
+                    operationIndex++;
+                    continue;
+                }
+
                 if (await SaveAsync(batch.UpdateEntries[0], cancellationToken).ConfigureAwait(false))
                 {
-                    rowsAffected++;
+                    state.RowsAffected++;
                 }
 
+                state.CommittedOperations++;
+                operationIndex++;
                 continue;
             }
 
             foreach (var transaction in CreateTransactions(batch))
             {
+                // Skip the operations that were already committed by a previous execution strategy attempt.
+                if (operationIndex < state.CommittedOperations)
+                {
+                    operationIndex++;
+                    continue;
+                }
+
                 try
                 {
-                    var response = await _cosmosClient.ExecuteTransactionalBatchAsync(transaction, SessionTokenStorage, cancellationToken).ConfigureAwait(false);
-                    if (!response.IsSuccess)
+                    await _cosmosClient.ExecuteTransactionalBatchAsync(transaction, SessionTokenStorage, cancellationToken).ConfigureAwait(false);
+                }
+                catch (CosmosException cosmosException)
+                {
+                    var entries = transaction.Entries.Select(x => x.Entry).ToArray();
+                    var exception = WrapUpdateException(cosmosException, entries);
+                    if (exception is not DbUpdateConcurrencyException
+                        || !(await Dependencies.Logger.OptimisticConcurrencyExceptionAsync(
+                                transaction.Entries.First().Entry.Context, entries, (DbUpdateConcurrencyException)exception, null, cancellationToken)
+                            .ConfigureAwait(false)).IsSuppressed)
                     {
-                        var exception = WrapUpdateException(response.Exception, response.ErroredEntries);
-                        if (exception is not DbUpdateConcurrencyException
-                            || !(await Dependencies.Logger.OptimisticConcurrencyExceptionAsync(
-                                    transaction.Entries.First().Entry.Context, transaction.Entries.Select(x => x.Entry).ToArray(), (DbUpdateConcurrencyException)exception, null, cancellationToken)
-                                .ConfigureAwait(false)).IsSuppressed)
-                        {
-                            throw exception;
-                        }
+                        throw exception;
                     }
                 }
                 catch (Exception ex) when (!ex.IsCritical() && ex is not DbUpdateException)
                 {
-                    var exception = WrapUpdateException(ex, transaction.Entries.Select(x => x.Entry).ToArray());
-                    throw exception;
+                    throw WrapUpdateException(ex, transaction.Entries.Select(x => x.Entry).ToArray());
                 }
 
-                rowsAffected += transaction.Entries.Count;
+                state.RowsAffected += transaction.Entries.Count;
+                state.CommittedOperations++;
+                operationIndex++;
             }
         }
 
-        return rowsAffected;
+        return state.RowsAffected;
     }
 
     private SaveGroups CreateSaveGroups(IList<IUpdateEntry> entries)
@@ -587,6 +633,13 @@ private sealed class SaveGroups
         public required IEnumerable<(Grouping Key, List<CosmosUpdateEntry> UpdateEntries)> BatchableUpdateEntries { get; init; }
     }
 
+    private sealed class BatchExecutionState
+    {
+        public int CommittedOperations { get; set; }
+
+        public int RowsAffected { get; set; }
+    }
+
     private sealed class CosmosUpdateEntry
     {
         public required IUpdateEntry Entry { get; init; }

src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs

@@ -146,5 +146,5 @@ IAsyncEnumerable<JToken> ExecuteSqlQueryAsync(
     ///     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>
-    Task<CosmosTransactionalBatchResult> ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default);
+    Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default);
 }

test/EFCore.Cosmos.FunctionalTests/CosmosTransactionalBatchTest.cs

@@ -5,6 +5,7 @@
 using Microsoft.Azure.Cosmos;
 using Microsoft.Azure.Cosmos.Scripts;
 using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 
 namespace Microsoft.EntityFrameworkCore;
 
@@ -39,6 +40,59 @@ public virtual async Task SaveChanges_fails_for_duplicate_key_in_same_partition_
         Assert.Equal(1, customersCount);
     }
 
+    [Fact]
+    public virtual async Task SaveChanges_transactional_batch_failure_flows_through_execution_strategy()
+    {
+        using (var arrangeContext = Fixture.CreateContext())
+        {
+            arrangeContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" });
+            await arrangeContext.SaveChangesAsync();
+        }
+
+        var recordedExceptions = new List<Exception>();
+
+        var options = new DbContextOptionsBuilder<TransactionalBatchContext>()
+            .UseCosmos(
+                ((CosmosTestStore)Fixture.TestStore).ConnectionUri,
+                ((CosmosTestStore)Fixture.TestStore).AuthToken,
+                Fixture.TestStore.Name,
+                cfg =>
+                {
+                    cfg.ApplyConfiguration();
+                    cfg.ExecutionStrategy(d => new RecordingExecutionStrategy(d, recordedExceptions));
+                })
+            .Options;
+
+        using var context = new TransactionalBatchContext(options);
+        context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always;
+
+        context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" });
+        context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" });
+
+        // Before the fix the transactional batch failure never reached the execution strategy, so it could not be
+        // retried and RetryLimitExceededException was never thrown.
+        await Assert.ThrowsAsync<RetryLimitExceededException>(() => context.SaveChangesAsync());
+
+        Assert.NotEmpty(recordedExceptions);
+        Assert.All(
+            recordedExceptions,
+            e => Assert.Equal(HttpStatusCode.Conflict, Assert.IsType<CosmosException>(e).StatusCode));
+
+        using var assertContext = Fixture.CreateContext();
+        var customersCount = await assertContext.Customers.CountAsync();
+        Assert.Equal(1, customersCount);
+    }
+
+    private sealed class RecordingExecutionStrategy(ExecutionStrategyDependencies dependencies, List<Exception> recordedExceptions)
+        : CosmosExecutionStrategy(dependencies, maxRetryCount: 2, maxRetryDelay: TimeSpan.FromMilliseconds(1))
+    {
+        protected override bool ShouldRetryOn(Exception exception)
+        {
+            recordedExceptions.Add(exception);
+            return true;
+        }
+    }
+
     [Fact]
     public virtual async Task SaveChanges_fails_for_duplicate_key_in_same_partition_writes_only_partition_staged_before_error()
     {