Ensure ComponentState is always disposed (#48406)

Author: SteveSandersonMS

src/Components/Components/src/PublicAPI.Unshipped.txt

@@ -34,7 +34,6 @@ Microsoft.AspNetCore.Components.Rendering.ComponentState
 Microsoft.AspNetCore.Components.Rendering.ComponentState.Component.get -> Microsoft.AspNetCore.Components.IComponent!
 Microsoft.AspNetCore.Components.Rendering.ComponentState.ComponentId.get -> int
 Microsoft.AspNetCore.Components.Rendering.ComponentState.ComponentState(Microsoft.AspNetCore.Components.RenderTree.Renderer! renderer, int componentId, Microsoft.AspNetCore.Components.IComponent! component, Microsoft.AspNetCore.Components.Rendering.ComponentState? parentComponentState) -> void
-Microsoft.AspNetCore.Components.Rendering.ComponentState.Dispose() -> void
 Microsoft.AspNetCore.Components.Rendering.ComponentState.ParentComponentState.get -> Microsoft.AspNetCore.Components.Rendering.ComponentState?
 Microsoft.AspNetCore.Components.RenderTree.Renderer.GetComponentState(int componentId) -> Microsoft.AspNetCore.Components.Rendering.ComponentState!
 Microsoft.AspNetCore.Components.Sections.SectionContent
@@ -61,6 +60,7 @@ override Microsoft.AspNetCore.Components.EventCallback.GetHashCode() -> int
 override Microsoft.AspNetCore.Components.EventCallback.Equals(object? obj) -> bool
 override Microsoft.AspNetCore.Components.EventCallback<TValue>.GetHashCode() -> int
 override Microsoft.AspNetCore.Components.EventCallback<TValue>.Equals(object? obj) -> bool
+virtual Microsoft.AspNetCore.Components.Rendering.ComponentState.DisposeAsync() -> System.Threading.Tasks.ValueTask
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.AddPendingTask(Microsoft.AspNetCore.Components.Rendering.ComponentState? componentState, System.Threading.Tasks.Task! task) -> void
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.CreateComponentState(int componentId, Microsoft.AspNetCore.Components.IComponent! component, Microsoft.AspNetCore.Components.Rendering.ComponentState? parentComponentState) -> Microsoft.AspNetCore.Components.Rendering.ComponentState!
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.DispatchEventAsync(ulong eventHandlerId, Microsoft.AspNetCore.Components.RenderTree.EventFieldInfo? fieldInfo, System.EventArgs! eventArgs, bool quiesce) -> System.Threading.Tasks.Task!

src/Components/Components/src/RenderTree/Renderer.cs

@@ -876,28 +876,20 @@ private void ProcessDisposalQueueInExistingBatch()
             var disposeComponentId = _batchBuilder.ComponentDisposalQueue.Dequeue();
             var disposeComponentState = GetRequiredComponentState(disposeComponentId);
             Log.DisposingComponent(_logger, disposeComponentState);
-            if (!(disposeComponentState.Component is IAsyncDisposable))
-            {
-                if (!disposeComponentState.TryDisposeInBatch(_batchBuilder, out var exception))
-                {
-                    exceptions ??= new List<Exception>();
-                    exceptions.Add(exception);
-                }
-            }
-            else
+
+            try
             {
-                var result = disposeComponentState.DisposeInBatchAsync(_batchBuilder);
-                if (result.IsCompleted)
+                var disposalTask = disposeComponentState.DisposeInBatchAsync(_batchBuilder);
+                if (disposalTask.IsCompletedSuccessfully)
                 {
-                    if (!result.IsCompletedSuccessfully)
-                    {
-                        exceptions ??= new List<Exception>();
-                        exceptions.Add(result.Exception);
-                    }
+                    // If it's a IValueTaskSource backed ValueTask,
+                    // inform it its result has been read so it can reset
+                    disposalTask.GetAwaiter().GetResult();
                 }
                 else
                 {
                     // We set owningComponentState to null because we don't want exceptions during disposal to be recoverable
+                    var result = disposalTask.AsTask();
                     AddToPendingTasksWithErrorHandling(GetHandledAsynchronousDisposalErrorsTask(result), owningComponentState: null);
 
                     async Task GetHandledAsynchronousDisposalErrorsTask(Task result)
@@ -913,6 +905,11 @@ async Task GetHandledAsynchronousDisposalErrorsTask(Task result)
                     }
                 }
             }
+            catch (Exception exception)
+            {
+                exceptions ??= new List<Exception>();
+                exceptions.Add(exception);
+            }
 
             _componentStateById.Remove(disposeComponentId);
             _componentStateByComponent.Remove(disposeComponentState.Component);
@@ -1080,39 +1077,25 @@ protected virtual void Dispose(bool disposing)
         {
             Log.DisposingComponent(_logger, componentState);
 
-            // Components shouldn't need to implement IAsyncDisposable and IDisposable simultaneously,
-            // but in case they do, we prefer the async overload since we understand the sync overload
-            // is implemented for more "constrained" scenarios.
-            // Component authors are responsible for their IAsyncDisposable implementations not taking
-            // forever.
-            if (componentState.Component is IAsyncDisposable asyncDisposable)
+            try
             {
-                try
+                var task = componentState.DisposeAsync();
+                if (task.IsCompletedSuccessfully)
                 {
-                    var task = asyncDisposable.DisposeAsync();
-                    if (!task.IsCompletedSuccessfully)
-                    {
-                        asyncDisposables ??= new();
-                        asyncDisposables.Add(task.AsTask());
-                    }
+                    // If it's a IValueTaskSource backed ValueTask,
+                    // inform it its result has been read so it can reset
+                    task.GetAwaiter().GetResult();
                 }
-                catch (Exception exception)
+                else
                 {
-                    exceptions ??= new List<Exception>();
-                    exceptions.Add(exception);
+                    asyncDisposables ??= new();
+                    asyncDisposables.Add(task.AsTask());
                 }
             }
-            else if (componentState.Component is IDisposable disposable)
+            catch (Exception exception)
             {
-                try
-                {
-                    componentState.Dispose();
-                }
-                catch (Exception exception)
-                {
-                    exceptions ??= new List<Exception>();
-                    exceptions.Add(exception);
-                }
+                exceptions ??= new List<Exception>();
+                exceptions.Add(exception);
             }
         }
 

src/Components/Components/src/Rendering/ComponentState.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Diagnostics;
-using System.Diagnostics.CodeAnalysis;
 using Microsoft.AspNetCore.Components.RenderTree;
 
 namespace Microsoft.AspNetCore.Components.Rendering;
@@ -13,7 +12,7 @@ namespace Microsoft.AspNetCore.Components.Rendering;
 /// detail of <see cref="Renderer"/>.
 /// </summary>
 [DebuggerDisplay($"{{{nameof(GetDebuggerDisplay)}(),nq}}")]
-public class ComponentState : IDisposable
+public class ComponentState : IAsyncDisposable
 {
     private readonly Renderer _renderer;
     private readonly IReadOnlyList<CascadingParameterState> _cascadingParameters;
@@ -108,41 +107,6 @@ internal void RenderIntoBatch(RenderBatchBuilder batchBuilder, RenderFragment re
         batchBuilder.InvalidateParameterViews();
     }
 
-    internal bool TryDisposeInBatch(RenderBatchBuilder batchBuilder, [NotNullWhen(false)] out Exception? exception)
-    {
-        _componentWasDisposed = true;
-        exception = null;
-
-        try
-        {
-            if (Component is IDisposable disposable)
-            {
-                disposable.Dispose();
-            }
-        }
-        catch (Exception ex)
-        {
-            exception = ex;
-        }
-
-        CleanupComponentStateResources(batchBuilder);
-
-        return exception == null;
-    }
-
-    private void CleanupComponentStateResources(RenderBatchBuilder batchBuilder)
-    {
-        // We don't expect these things to throw.
-        RenderTreeDiffBuilder.DisposeFrames(batchBuilder, CurrentRenderTree.GetFrames());
-
-        if (_hasAnyCascadingParameterSubscriptions)
-        {
-            RemoveCascadingParameterSubscriptions();
-        }
-
-        DisposeBuffers();
-    }
-
     // Callers expect this method to always return a faulted task.
     internal Task NotifyRenderCompletedAsync()
     {
@@ -254,15 +218,26 @@ private void RemoveCascadingParameterSubscriptions()
     }
 
     /// <summary>
-    /// Disposes this instance.
+    /// Disposes this instance and its associated component.
     /// </summary>
-    public void Dispose()
+    public virtual ValueTask DisposeAsync()
     {
+        _componentWasDisposed = true;
         DisposeBuffers();
 
-        if (Component is IDisposable disposable)
+        // Components shouldn't need to implement IAsyncDisposable and IDisposable simultaneously,
+        // but in case they do, we prefer the async overload since we understand the sync overload
+        // is implemented for more "constrained" scenarios.
+        // Component authors are responsible for their IAsyncDisposable implementations not taking
+        // forever.
+        if (Component is IAsyncDisposable asyncDisposable)
         {
-            disposable.Dispose();
+            return asyncDisposable.DisposeAsync();
+        }
+        else
+        {
+            (Component as IDisposable)?.Dispose();
+            return ValueTask.CompletedTask;
         }
     }
 
@@ -273,33 +248,17 @@ private void DisposeBuffers()
         _latestDirectParametersSnapshot?.Dispose();
     }
 
-    internal Task DisposeInBatchAsync(RenderBatchBuilder batchBuilder)
+    internal ValueTask DisposeInBatchAsync(RenderBatchBuilder batchBuilder)
     {
-        _componentWasDisposed = true;
-
-        CleanupComponentStateResources(batchBuilder);
+        // We don't expect these things to throw.
+        RenderTreeDiffBuilder.DisposeFrames(batchBuilder, CurrentRenderTree.GetFrames());
 
-        try
-        {
-            var result = ((IAsyncDisposable)Component).DisposeAsync();
-            if (result.IsCompletedSuccessfully)
-            {
-                // If it's a IValueTaskSource backed ValueTask,
-                // inform it its result has been read so it can reset
-                result.GetAwaiter().GetResult();
-                return Task.CompletedTask;
-            }
-            else
-            {
-                // We know we are dealing with an exception that happened asynchronously, so return a task
-                // to the caller so that he can unwrap it.
-                return result.AsTask();
-            }
-        }
-        catch (Exception e)
+        if (_hasAnyCascadingParameterSubscriptions)
         {
-            return Task.FromException(e);
+            RemoveCascadingParameterSubscriptions();
         }
+
+        return DisposeAsync();
     }
 
     private string GetDebuggerDisplay()

src/Components/Components/test/RendererTest.cs

@@ -2334,9 +2334,8 @@ public void RenderBatch_HandlesSynchronousExceptionsInAsyncDisposableComponents(
 
         // Outer component is still alive and not disposed.
         Assert.False(component.Disposed);
-        var aex = Assert.IsType<AggregateException>(Assert.Single(renderer.HandledExceptions));
-        var innerException = Assert.Single(aex.Flatten().InnerExceptions);
-        Assert.Same(exception1, innerException);
+        var aex = Assert.Single(renderer.HandledExceptions);
+        Assert.Same(exception1, aex);
     }
 
     [Fact]
@@ -2493,8 +2492,7 @@ public void RenderBatch_ReportsSynchronousCancelationsAsErrors()
 
         // Outer component is still alive and not disposed.
         Assert.False(component.Disposed);
-        var aex = Assert.IsType<AggregateException>(Assert.Single(renderer.HandledExceptions));
-        Assert.IsType<TaskCanceledException>(Assert.Single(aex.Flatten().InnerExceptions));
+        Assert.IsType<TaskCanceledException>(Assert.Single(renderer.HandledExceptions));
     }
 
     [Fact]

src/Components/Shared/test/TestRenderer.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Runtime.ExceptionServices;
+using Microsoft.AspNetCore.Components.Rendering;
 using Microsoft.AspNetCore.Components.RenderTree;
 using Microsoft.Extensions.Logging.Abstractions;
 
@@ -60,6 +61,8 @@ protected internal override bool ShouldTrackNamedEventHandlers()
 
     public Task NextRenderResultTask { get; set; } = Task.CompletedTask;
 
+    private HashSet<TestRendererComponentState> UndisposedComponentStates { get; } = new();
+
     public new int AssignRootComponentId(IComponent component)
         => base.AssignRootComponentId(component);
 
@@ -145,4 +148,35 @@ protected override Task UpdateDisplayAsync(in RenderBatch renderBatch)
 
     public new void ProcessPendingRender()
         => base.ProcessPendingRender();
+
+    protected override ComponentState CreateComponentState(int componentId, IComponent component, ComponentState parentComponentState)
+        => new TestRendererComponentState(this, componentId, component, parentComponentState);
+
+    protected override void Dispose(bool disposing)
+    {
+        base.Dispose(disposing);
+
+        if (UndisposedComponentStates.Count > 0)
+        {
+            throw new InvalidOperationException("Did not dispose all the ComponentState instances. This could lead to ArrayBuffer not returning buffers to its pool.");
+        }
+    }
+
+    class TestRendererComponentState : ComponentState, IAsyncDisposable
+    {
+        private readonly TestRenderer _renderer;
+
+        public TestRendererComponentState(Renderer renderer, int componentId, IComponent component, ComponentState parentComponentState)
+            : base(renderer, componentId, component, parentComponentState)
+        {
+            _renderer = (TestRenderer)renderer;
+            _renderer.UndisposedComponentStates.Add(this);
+        }
+
+        public override ValueTask DisposeAsync()
+        {
+            _renderer.UndisposedComponentStates.Remove(this);
+            return base.DisposeAsync();
+        }
+    }
 }