Support completing request units of work before responses start

framework/Volo.Abp.slnx

@@ -194,6 +194,7 @@
     <Project Path="test/Volo.Abp.AspNetCore.Serilog.Tests/Volo.Abp.AspNetCore.Serilog.Tests.csproj" />
     <Project Path="test/Volo.Abp.AspNetCore.SignalR.Tests/Volo.Abp.AspNetCore.SignalR.Tests.csproj" />
     <Project Path="test/Volo.Abp.AspNetCore.Tests/Volo.Abp.AspNetCore.Tests.csproj" />
+    <Project Path="test/Volo.Abp.AspNetCore.Uow.Tests/Volo.Abp.AspNetCore.Uow.Tests.csproj" />
     <Project Path="test/Volo.Abp.Auditing.Tests/Volo.Abp.Auditing.Tests.csproj" />
     <Project Path="test/Volo.Abp.Authorization.Tests/Volo.Abp.Authorization.Tests.csproj" />
     <Project Path="test/Volo.Abp.Autofac.Tests/Volo.Abp.Autofac.Tests.csproj" />

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

@@ -1,4 +1,4 @@
-using System.Collections.Generic;
+using System.Collections.Generic;
 
 namespace Volo.Abp.AspNetCore.Uow;
 
@@ -11,4 +11,28 @@ public class AbpAspNetCoreUnitOfWorkOptions
     /// starting with an ignored URL.  
     /// </summary>
     public List<string> IgnoredUrls { get; } = new List<string>();
+
+    /// <summary>
+    /// Completes the request unit of work just before the response starts (on
+    /// <c>HttpResponse.OnStarting</c>) instead of at the end of the pipeline, so data written during
+    /// the request is committed before the response is flushed. Disabled by default; enable it here
+    /// globally or opt-in per endpoint via <see cref="CompleteUnitOfWorkOnResponseStartingUrls"/>.
+    /// <para>
+    /// Trade-offs when it applies: an exception after the response starts can no longer roll back the
+    /// committed data (commit and network response are not atomic); database access after the response
+    /// starts is outside the request unit of work (unsuitable for streaming responses); unit of work
+    /// events and completed handlers run before the first response byte (adding to its latency); a
+    /// nested (requiresNew) unit of work that is current when the response starts is left to its owner
+    /// and the request unit of work then completes at the end of the pipeline as usual.
+    /// </para>
+    /// </summary>
+    public bool CompleteUnitOfWorkOnResponseStarting { get; set; } = false;
+
+    /// <summary>
+    /// Request path prefixes that opt-in to <see cref="CompleteUnitOfWorkOnResponseStarting"/> even when
+    /// it is globally disabled. A request whose path starts with one of these values (for example
+    /// "/connect") is included, matched like <see cref="IgnoredUrls"/>. Use
+    /// <see cref="CompleteUnitOfWorkOnResponseStarting"/> to enable it for every request handled by the middleware.
+    /// </summary>
+    public List<string> CompleteUnitOfWorkOnResponseStartingUrls { get; } = new List<string>();
 }

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

@@ -37,8 +37,28 @@ public async override Task InvokeAsync(HttpContext context, RequestDelegate next
 
         using (var uow = _unitOfWorkManager.Reserve(UnitOfWork.UnitOfWorkReservationName))
         {
+            var completionAttemptedOnResponseStarting = false;
+
+            if (!context.Response.HasStarted && ShouldCompleteOnResponseStarting(context))
+            {
+                context.Response.OnStarting(async () =>
+                {
+                    // A nested (requiresNew) unit of work that is current is left to its owner.
+                    if (_unitOfWorkManager.Current == uow)
+                    {
+                        // Set before completing so a post-commit failure isn't masked by the completion below.
+                        completionAttemptedOnResponseStarting = true;
+                        await uow.CompleteAsync(_cancellationTokenProvider.Token);
+                    }
+                });
+            }
+
             await next(context);
-            await uow.CompleteAsync(_cancellationTokenProvider.Token);
+
+            if (!completionAttemptedOnResponseStarting)
+            {
+                await uow.CompleteAsync(_cancellationTokenProvider.Token);
+            }
         }
     }
 
@@ -48,6 +68,13 @@ private bool IsIgnoredUrl(HttpContext context)
                _options.IgnoredUrls.Any(x => context.Request.Path.Value.StartsWith(x, StringComparison.OrdinalIgnoreCase));
     }
 
+    private bool ShouldCompleteOnResponseStarting(HttpContext context)
+    {
+        return _options.CompleteUnitOfWorkOnResponseStarting ||
+               (context.Request.Path.Value != null &&
+                _options.CompleteUnitOfWorkOnResponseStartingUrls.Any(x => context.Request.Path.Value.StartsWith(x, StringComparison.OrdinalIgnoreCase)));
+    }
+
     protected async override Task<bool> ShouldSkipAsync(HttpContext context, RequestDelegate next)
     {
         // Blazor components will render concurrently, so we need to skip the middleware for them.

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,12 +1,19 @@
-using System.Net.Http;
+using System.Net;
+using System.Net.Http;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
 using Shouldly;
+using Volo.Abp.AspNetCore.Uow;
 using Xunit;
 
 namespace Volo.Abp.AspNetCore.Mvc.Uow;
 
 public class UnitOfWorkMiddleware_Tests : AspNetCoreMvcTestBase
 {
+    private AbpAspNetCoreUnitOfWorkOptions Options =>
+        ServiceProvider.GetRequiredService<IOptions<AbpAspNetCoreUnitOfWorkOptions>>().Value;
+
     [Fact]
     public async Task Get_Actions_Should_Not_Be_Transactional()
     {
@@ -27,4 +34,89 @@ public async Task Query_Actions_Should_Not_Be_Transactional()
         var result = await Client.SendAsync(requestMessage);
         result.IsSuccessStatusCode.ShouldBeTrue();
     }
+
+    [Fact]
+    public async Task Ambient_Uow_Should_Be_Completed_Before_Response_Is_Flushed_When_Enabled()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        var result = await GetResponseAsStringAsync("/api/unitofwork-test/CommitBeforeResponseFlush");
+        result.ShouldBe("first:completed");
+    }
+
+    [Fact]
+    public async Task Ambient_Uow_Is_Not_Completed_On_Response_Start_By_Default()
+    {
+        var result = await GetResponseAsStringAsync("/api/unitofwork-test/CommitBeforeResponseFlush");
+        result.ShouldBe("first:not-completed");
+    }
+
+    [Fact]
+    public async Task Ambient_Uow_Is_Already_Completed_When_An_Exception_Is_Raised_After_The_Response_Started()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        // Once the response has started, an exception can't turn it into an error response (the
+        // connection is reset). Database-level rollback/commit is covered by the relational tests.
+        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()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        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()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/RawDatabaseProviderAfterResponseFlush");
+        body.ShouldBe("first:threw-AbpException");
+    }
+
+    [Fact]
+    public async Task Response_Flush_Inside_Nested_Uow_Should_Not_Complete_The_Nested_Uow()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/NestedUowDuringResponseFlush");
+        body.ShouldBe("first:outer-not-completed:nested-completed-by-owner");
+    }
+
+    [Fact]
+    public async Task Completing_The_Uow_In_The_Action_Still_Fails_At_End_Of_Pipeline_By_Default()
+    {
+        var response = await Client.GetAsync("/api/unitofwork-test/CompleteCurrentUow");
+        response.StatusCode.ShouldBe(HttpStatusCode.InternalServerError);
+    }
+
+    [Fact]
+    public async Task Opt_In_Url_Enables_The_Feature_For_A_Matching_Path()
+    {
+        Options.CompleteUnitOfWorkOnResponseStartingUrls.Add("/api/unitofwork-test/CommitBeforeResponseFlush");
+
+        var result = await GetResponseAsStringAsync("/api/unitofwork-test/CommitBeforeResponseFlush");
+        result.ShouldBe("first:completed");
+    }
+
+    [Fact]
+    public async Task Opt_In_Url_With_A_Trailing_Slash_Still_Matches()
+    {
+        Options.CompleteUnitOfWorkOnResponseStartingUrls.Add("/api/unitofwork-test/");
+
+        var result = await GetResponseAsStringAsync("/api/unitofwork-test/CommitBeforeResponseFlush");
+        result.ShouldBe("first:completed");
+    }
 }

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

@@ -1,5 +1,13 @@
-using Microsoft.AspNetCore.Mvc;
+using System;
+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 +72,118 @@ public void ExceptionOnComplete()
 
         _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();
+
+        _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);
+    }
+
+    [HttpGet]
+    [Route("NestedUowDuringResponseFlush")]
+    public async Task NestedUowDuringResponseFlush()
+    {
+        using (var nested = UnitOfWorkManager.Begin(requiresNew: true, isTransactional: false))
+        {
+            await Response.WriteAsync("first");
+            await Response.Body.FlushAsync();
+
+            // The outer request unit of work (nested.Outer) must not have been completed on response
+            // start while a nested unit of work is current.
+            await Response.WriteAsync(nested.Outer!.IsCompleted ? ":outer-completed" : ":outer-not-completed");
+
+            string outcome;
+            try
+            {
+                await nested.CompleteAsync();
+                outcome = ":nested-completed-by-owner";
+            }
+            catch (AbpException)
+            {
+                outcome = ":nested-already-completed";
+            }
+
+            await Response.WriteAsync(outcome);
+        }
+    }
+
+    [HttpGet]
+    [Route("CompleteCurrentUow")]
+    public async Task CompleteCurrentUow()
+    {
+        // Complete the request unit of work inside the action, without writing the response yet.
+        // The middleware must still try to complete it at the end of the pipeline (original behavior).
+        await CurrentUnitOfWork.CompleteAsync();
+    }
 }

framework/test/Volo.Abp.AspNetCore.Uow.Tests/Volo.Abp.AspNetCore.Uow.Tests.csproj

@@ -0,0 +1,26 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <Import Project="..\..\..\common.test.props" />
+
+  <PropertyGroup>
+    <TargetFramework>net10.0</TargetFramework>
+    <AssemblyName>Volo.Abp.AspNetCore.Uow.Tests</AssemblyName>
+    <PackageId>Volo.Abp.AspNetCore.Uow.Tests</PackageId>
+    <GenerateRuntimeConfigurationFiles>true</GenerateRuntimeConfigurationFiles>
+    <GenerateAssemblyConfigurationAttribute>false</GenerateAssemblyConfigurationAttribute>
+    <GenerateAssemblyCompanyAttribute>false</GenerateAssemblyCompanyAttribute>
+    <GenerateAssemblyProductAttribute>false</GenerateAssemblyProductAttribute>
+    <PreserveCompilationReferences>true</PreserveCompilationReferences>
+    <RootNamespace />
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\Volo.Abp.AspNetCore.Mvc\Volo.Abp.AspNetCore.Mvc.csproj" />
+    <ProjectReference Include="..\..\src\Volo.Abp.AspNetCore.TestBase\Volo.Abp.AspNetCore.TestBase.csproj" />
+    <ProjectReference Include="..\..\src\Volo.Abp.EntityFrameworkCore.Sqlite\Volo.Abp.EntityFrameworkCore.Sqlite.csproj" />
+    <ProjectReference Include="..\..\src\Volo.Abp.Autofac\Volo.Abp.Autofac.csproj" />
+    <ProjectReference Include="..\AbpTestBase\AbpTestBase.csproj" />
+    <PackageReference Include="Microsoft.NET.Test.Sdk" />
+  </ItemGroup>
+
+</Project>