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,29 @@ 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, and an active child
+    /// unit of work scope (begun without requiresNew), are left to their owners 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,36 @@ public async override Task InvokeAsync(HttpContext context, RequestDelegate next
 
         using (var uow = _unitOfWorkManager.Reserve(UnitOfWork.UnitOfWorkReservationName))
         {
+            var completionStarted = false;
+
+            if (!context.Response.HasStarted && ShouldCompleteOnResponseStarting(context))
+            {
+                context.Response.OnStarting(async () =>
+                {
+                    // Skip if the completion has already been started at the end of the pipeline;
+                    // the response is then being started from inside that completion (e.g. by an
+                    // event handler writing to the response), so completing again would fail.
+                    // A nested (requiresNew) unit of work that is current and an active child
+                    // unit of work scope are left to their owners; the request unit of work then
+                    // completes at the end of the pipeline as usual.
+                    if (!completionStarted &&
+                        _unitOfWorkManager.Current == uow &&
+                        !uow.HasActiveChildUnitOfWorks())
+                    {
+                        // Set before completing so a post-commit failure isn't masked by the completion below.
+                        completionStarted = true;
+                        await uow.CompleteAsync(_cancellationTokenProvider.Token);
+                    }
+                });
+            }
+
             await next(context);
-            await uow.CompleteAsync(_cancellationTokenProvider.Token);
+
+            if (!completionStarted)
+            {
+                completionStarted = true;
+                await uow.CompleteAsync(_cancellationTokenProvider.Token);
+            }
         }
     }
 
@@ -48,6 +76,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/src/Volo.Abp.Uow/Volo/Abp/Uow/ChildUnitOfWork.cs

@@ -30,12 +30,14 @@ internal class ChildUnitOfWork : IUnitOfWork
     public Dictionary<string, object> Items => _parent.Items;
 
     private readonly IUnitOfWork _parent;
+    private bool _isDisposed;
 
     public ChildUnitOfWork([NotNull] IUnitOfWork parent)
     {
         Check.NotNull(parent, nameof(parent));
 
         _parent = parent;
+        _parent.IncrementActiveChildUnitOfWorkCount();
 
         _parent.Failed += (sender, args) => { Failed.InvokeSafely(sender!, args); };
         _parent.Disposed += (sender, args) => { Disposed.InvokeSafely(sender!, args); };
@@ -122,7 +124,13 @@ public ITransactionApi GetOrAddTransactionApi(string key, Func<ITransactionApi>
 
     public void Dispose()
     {
+        if (_isDisposed)
+        {
+            return;
+        }
 
+        _isDisposed = true;
+        _parent.DecrementActiveChildUnitOfWorkCount();
     }
 
     public override string ToString()

framework/src/Volo.Abp.Uow/Volo/Abp/Uow/UnitOfWorkExtensions.cs

@@ -7,13 +7,40 @@ namespace Volo.Abp.Uow;
 
 public static class UnitOfWorkExtensions
 {
+    private const string ActiveChildUnitOfWorkCountItemKey = "_AbpActiveChildUnitOfWorkCount";
+
     public static bool IsReservedFor([NotNull] this IUnitOfWork unitOfWork, string reservationName)
     {
         Check.NotNull(unitOfWork, nameof(unitOfWork));
 
         return unitOfWork.IsReserved && unitOfWork.ReservationName == reservationName;
     }
 
+    /// <summary>
+    /// Checks if there is an active (not yet disposed) child unit of work scope over the given
+    /// unit of work, i.e. a scope created by <see cref="IUnitOfWorkManager.Begin"/> without
+    /// requiresNew while this unit of work was current. Such a scope shares this unit of work,
+    /// so it should not be completed while the scope is still active.
+    /// </summary>
+    public static bool HasActiveChildUnitOfWorks([NotNull] this IUnitOfWork unitOfWork)
+    {
+        Check.NotNull(unitOfWork, nameof(unitOfWork));
+
+        return unitOfWork.Items.GetOrDefault(ActiveChildUnitOfWorkCountItemKey) is int count && count > 0;
+    }
+
+    internal static void IncrementActiveChildUnitOfWorkCount(this IUnitOfWork unitOfWork)
+    {
+        var count = unitOfWork.Items.GetOrDefault(ActiveChildUnitOfWorkCountItemKey) as int? ?? 0;
+        unitOfWork.Items[ActiveChildUnitOfWorkCountItemKey] = count + 1;
+    }
+
+    internal static void DecrementActiveChildUnitOfWorkCount(this IUnitOfWork unitOfWork)
+    {
+        var count = unitOfWork.Items.GetOrDefault(ActiveChildUnitOfWorkCountItemKey) as int? ?? 0;
+        unitOfWork.Items[ActiveChildUnitOfWorkCountItemKey] = Math.Max(0, count - 1);
+    }
+
     public static void AddItem<TValue>([NotNull] this IUnitOfWork unitOfWork, string key, TValue value)
         where TValue : class
     {

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

@@ -0,0 +1,34 @@
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Volo.Abp.DependencyInjection;
+using Volo.Abp.EventBus;
+
+namespace Volo.Abp.AspNetCore.Mvc.Uow;
+
+public class ResponseWritingTestEvent
+{
+}
+
+/// <summary>
+/// Writes to the HTTP response from a local event handler. When the event is published inside the
+/// request unit of work, this runs during the unit of work completion at the end of the pipeline
+/// and starts the response from inside that completion.
+/// </summary>
+public class ResponseWritingTestEventHandler : ILocalEventHandler<ResponseWritingTestEvent>, ITransientDependency
+{
+    private readonly IHttpContextAccessor _httpContextAccessor;
+
+    public ResponseWritingTestEventHandler(IHttpContextAccessor httpContextAccessor)
+    {
+        _httpContextAccessor = httpContextAccessor;
+    }
+
+    public async Task HandleEventAsync(ResponseWritingTestEvent eventData)
+    {
+        var response = _httpContextAccessor.HttpContext?.Response;
+        if (response != null)
+        {
+            await response.WriteAsync("event-written");
+        }
+    }
+}

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,114 @@ 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");
+    }
+
+    [Fact]
+    public async Task Response_Flush_Inside_A_Child_Uow_Scope_Should_Not_Complete_The_Request_Uow()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        // A child scope (Begin without requiresNew) shares the request unit of work, so completing
+        // it on response start would commit under the still-active scope; it is left to the end of
+        // the pipeline instead, like a nested (requiresNew) unit of work.
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/ChildUowDuringResponseFlush");
+        body.ShouldBe("first:request-not-completed");
+    }
+
+    [Fact]
+    public async Task An_Event_Handler_Starting_The_Response_During_The_End_Of_Pipeline_Completion_Should_Not_Fail()
+    {
+        Options.CompleteUnitOfWorkOnResponseStarting = true;
+
+        // The response does not start during the pipeline here, so the middleware completes the
+        // unit of work at its end; the event handler then starts the response from inside that
+        // completion. The OnStarting callback must not attempt a second completion (which would
+        // throw "Completion has already been requested for this unit of work").
+        var body = await GetResponseAsStringAsync("/api/unitofwork-test/PublishEventThatWritesResponseOnCompletion");
+        body.ShouldBe("event-written");
+    }
 }