Commit ambient unit of work before the response starts

framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Uow/AbpUnitOfWorkMiddleware.cs

@@ -37,8 +37,23 @@ public class AbpUnitOfWorkMiddleware : AbpMiddlewareBase, ITransientDependency
 
         using (var uow = _unitOfWorkManager.Reserve(UnitOfWork.UnitOfWorkReservationName))
         {
+            // Commit the ambient unit of work before the response starts, so data written
+            // during the request is committed before the response is flushed to the client.
+            context.Response.OnStarting(async () =>
+            {
+                var currentUow = _unitOfWorkManager.Current;
+                if (currentUow != null && !currentUow.IsCompleted)
+                {
+                    await currentUow.CompleteAsync(_cancellationTokenProvider.Token);
+                }
+            });
+
             await next(context);
-            await uow.CompleteAsync(_cancellationTokenProvider.Token);
+
+            if (!uow.IsCompleted)
+            {
+                await uow.CompleteAsync(_cancellationTokenProvider.Token);
+            }
         }
     }
 

framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Uow/TestUnitOfWorkConfig.cs

@@ -7,4 +7,6 @@ public class TestUnitOfWorkConfig : ISingletonDependency
     public const string ExceptionOnCompleteMessage = "TestUnitOfWork configured for exception";
 
     public bool ThrowExceptionOnComplete { get; set; }
+
+    public bool? UowCompletedAfterResponseFlush { get; set; }
 }

framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Uow/UnitOfWorkMiddleware_Tests.cs

@@ -1,5 +1,7 @@
-using System.Net.Http;
+using System.Net;
+using System.Net.Http;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Shouldly;
 using Xunit;
 
@@ -27,4 +29,43 @@ public class UnitOfWorkMiddleware_Tests : AspNetCoreMvcTestBase
         var result = await Client.SendAsync(requestMessage);
         result.IsSuccessStatusCode.ShouldBeTrue();
     }
+
+    [Fact]
+    public async Task Ambient_Uow_Should_Be_Completed_Before_Response_Is_Flushed()
+    {
+        var result = await GetResponseAsStringAsync("/api/unitofwork-test/CommitBeforeResponseFlush");
+        result.ShouldBe("first:completed");
+    }
+
+    [Fact]
+    public async Task Exception_After_Response_Flush_Should_Not_Undo_Committed_Work()
+    {
+        // Once the response has started, an exception can't turn it into an error response
+        // (the connection is reset). What matters: the uow was committed before the throw.
+        await Should.ThrowAsync<HttpRequestException>(async () =>
+        {
+            var response = await Client.GetAsync("/api/unitofwork-test/CommitThenThrowAfterResponseFlush");
+            await response.Content.ReadAsStringAsync();
+        });
+
+        ServiceProvider.GetRequiredService<TestUnitOfWorkConfig>()
+            .UowCompletedAfterResponseFlush.ShouldBe(true);
+    }
+
+    [Fact]
+    public async Task Repository_Access_After_Response_Flush_Runs_Outside_The_Request_Uow()
+    {
+        // After the response starts the request uow is gone; a repository still works via its
+        // own implicit uow (ambient=null), so it no longer joins the request transaction.
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/ReadRepositoryAfterResponseFlush");
+        body.ShouldBe("before=ok(1);after=ok(1,ambient=null)");
+    }
+
+    [Fact]
+    public async Task Raw_Database_Provider_After_Response_Flush_Throws()
+    {
+        // Unlike repositories, raw provider access after the response started has no uow and throws.
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/RawDatabaseProviderAfterResponseFlush");
+        body.ShouldBe("first:threw-AbpException");
+    }
 }

framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Uow/UnitOfWorkTestController.cs

@@ -1,5 +1,14 @@
-using Microsoft.AspNetCore.Mvc;
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Mvc;
 using Shouldly;
+using Volo.Abp;
+using Volo.Abp.Domain.Repositories;
+using Volo.Abp.MemoryDb;
+using Volo.Abp.TestApp.MemoryDb;
+using Volo.Abp.TestApp.Domain;
 using Volo.Abp.Uow;
 
 namespace Volo.Abp.AspNetCore.Mvc.Uow;
@@ -64,4 +73,82 @@ public class UnitOfWorkTestController : AbpController
 
         _testUnitOfWorkConfig.ThrowExceptionOnComplete = true;
     }
+
+    [HttpGet]
+    [Route("CommitBeforeResponseFlush")]
+    public async Task CommitBeforeResponseFlush()
+    {
+        var uow = CurrentUnitOfWork;
+        uow.ShouldNotBeNull();
+
+        // Start the response from inside the pipeline, before the middleware would commit.
+        await Response.WriteAsync("first");
+        await Response.Body.FlushAsync();
+
+        await Response.WriteAsync(uow.IsCompleted ? ":completed" : ":not-completed");
+    }
+
+    [HttpGet]
+    [Route("CommitThenThrowAfterResponseFlush")]
+    public async Task CommitThenThrowAfterResponseFlush()
+    {
+        var uow = CurrentUnitOfWork;
+
+        await Response.WriteAsync("first");
+        await Response.Body.FlushAsync();
+
+        // Record the commit state so the test can assert the throw below doesn't undo it.
+        _testUnitOfWorkConfig.UowCompletedAfterResponseFlush = uow.IsCompleted;
+
+        throw new UserFriendlyException("boom after the response was already flushed");
+    }
+
+    [HttpGet]
+    [Route("ReadRepositoryAfterResponseFlush")]
+    public async Task ReadRepositoryAfterResponseFlush()
+    {
+        var repository = LazyServiceProvider.LazyGetRequiredService<IRepository<Person, Guid>>();
+
+        var before = (await repository.GetListAsync()).Count;
+        await Response.WriteAsync($"before=ok({before})");
+        await Response.Body.FlushAsync();
+
+        string after;
+        try
+        {
+            var count = (await repository.GetListAsync()).Count;
+            after = $";after=ok({count},ambient={(UnitOfWorkManager.Current == null ? "null" : "present")})";
+        }
+        catch (Exception ex)
+        {
+            after = $";after=threw:{ex.GetType().Name}";
+        }
+
+        await Response.WriteAsync(after);
+    }
+
+    [HttpGet]
+    [Route("RawDatabaseProviderAfterResponseFlush")]
+    public async Task RawDatabaseProviderAfterResponseFlush()
+    {
+        var databaseProvider = LazyServiceProvider
+            .LazyGetRequiredService<IMemoryDatabaseProvider<TestAppMemoryDbContext>>();
+
+        await Response.WriteAsync("first");
+        await Response.Body.FlushAsync();
+
+        string outcome;
+        try
+        {
+            await databaseProvider.GetDatabaseAsync();
+            outcome = ":ok";
+        }
+        catch (AbpException)
+        {
+            // Raw provider access has no ambient uow once the response started, so it throws.
+            outcome = ":threw-AbpException";
+        }
+
+        await Response.WriteAsync(outcome);
+    }
 }