Guard UoW completion during response start

This change hardens `AbpUnitOfWorkMiddleware` to avoid invalid second completion attempts when the response starts during end-of-pipeline UoW completion, and to skip early completion while a shared child UoW scope is still active. It adds active child-scope tracking in `UnitOfWorkExtensions`/`ChildUnitOfWork`, updates option docs to clarify behavior, and adds MVC tests plus controller/event-handler scenarios covering child-scope response flush and event-driven response writes during completion.

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

@@ -22,8 +22,9 @@ public class AbpAspNetCoreUnitOfWorkOptions
     /// 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
+    /// nested (requiresNew) unit of work that is current when the response starts, and an active child
-    /// and the request unit of work then completes at the end of the pipeline as usual.
+    /// 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;

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

@@ -37,17 +37,24 @@ public class AbpUnitOfWorkMiddleware : AbpMiddlewareBase, ITransientDependency
 
         using (var uow = _unitOfWorkManager.Reserve(UnitOfWork.UnitOfWorkReservationName))
         {
-            var completionAttemptedOnResponseStarting = false;
+            var completionStarted = 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.
+                    // Skip if the completion has already been started at the end of the pipeline;
-                    if (_unitOfWorkManager.Current == uow)
+                    // 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.
-                        completionAttemptedOnResponseStarting = true;
+                        completionStarted = true;
                         await uow.CompleteAsync(_cancellationTokenProvider.Token);
                     }
                 });
@@ -55,8 +62,9 @@ public class AbpUnitOfWorkMiddleware : AbpMiddlewareBase, ITransientDependency
 
             await next(context);
 
-            if (!completionAttemptedOnResponseStarting)
+            if (!completionStarted)
             {
+                completionStarted = true;
                 await uow.CompleteAsync(_cancellationTokenProvider.Token);
             }
         }

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 @@ internal class ChildUnitOfWork : IUnitOfWork
 
     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,6 +7,8 @@ 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));
@@ -14,6 +16,31 @@ public static class UnitOfWorkExtensions
         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/UnitOfWorkMiddleware_Tests.cs

@@ -119,4 +119,29 @@ public class UnitOfWorkMiddleware_Tests : AspNetCoreMvcTestBase
         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");
+    }
 }

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

@@ -5,6 +5,7 @@ using Microsoft.AspNetCore.Mvc;
 using Shouldly;
 using Volo.Abp;
 using Volo.Abp.Domain.Repositories;
+using Volo.Abp.EventBus.Local;
 using Volo.Abp.MemoryDb;
 using Volo.Abp.TestApp.MemoryDb;
 using Volo.Abp.TestApp.Domain;
@@ -186,4 +187,35 @@ public class UnitOfWorkTestController : AbpController
         // The middleware must still try to complete it at the end of the pipeline (original behavior).
         await CurrentUnitOfWork.CompleteAsync();
     }
+
+    [HttpGet]
+    [Route("ChildUowDuringResponseFlush")]
+    public async Task ChildUowDuringResponseFlush()
+    {
+        var requestUow = CurrentUnitOfWork!;
+
+        using (UnitOfWorkManager.Begin())
+        {
+            await Response.WriteAsync("first");
+            await Response.Body.FlushAsync();
+
+            // The request unit of work must not be completed on response start while a child
+            // unit of work scope (begun without requiresNew) is still active over it.
+            await Response.WriteAsync(requestUow.IsCompleted ? ":request-completed" : ":request-not-completed");
+        }
+    }
+
+    [HttpGet]
+    [Route("PublishEventThatWritesResponseOnCompletion")]
+    public async Task<ActionResult> PublishEventThatWritesResponseOnCompletion()
+    {
+        // Published inside the request unit of work, so the handler runs while the middleware is
+        // completing it at the end of the pipeline. The handler writes to the response, which starts
+        // it mid-completion; the middleware's OnStarting callback must not try to complete again.
+        await LazyServiceProvider.LazyGetRequiredService<ILocalEventBus>()
+            .PublishAsync(new ResponseWritingTestEvent());
+
+        // Ok() sets 200 without writing the body, so the response does not start inside the pipeline.
+        return Ok();
+    }
 }