Ensure ComponentBase async lifecycle methods return non-null task (imported from Blazor PR 1620) (#4790)

src/Components/src/Microsoft.AspNetCore.Components/ComponentBase.cs

@@ -158,34 +158,51 @@ namespace Microsoft.AspNetCore.Components
                 _hasCalledInit = true;
                 OnInit();
 
-                // If you override OnInitAsync and return a nonnull task, then by default
+                // If you override OnInitAsync and return a noncompleted task, then by default
                 // we automatically re-render once that task completes.
                 var initTask = OnInitAsync();
-                if (initTask != null && initTask.Status != TaskStatus.RanToCompletion)
+                ContinueAfterLifecycleTask(initTask);
-                {
-                    initTask.ContinueWith(ContinueAfterLifecycleTask);
-                }
             }
 
             OnParametersSet();
             var parametersTask = OnParametersSetAsync();
-            if (parametersTask != null && parametersTask.Status != TaskStatus.RanToCompletion)
+            ContinueAfterLifecycleTask(parametersTask);
-            {
-                parametersTask.ContinueWith(ContinueAfterLifecycleTask);
-            }
 
             StateHasChanged();
         }
 
-        private void ContinueAfterLifecycleTask(Task task)
+        private async void ContinueAfterLifecycleTask(Task task)
         {
-            if (task.Exception == null)
+            switch (task == null ? TaskStatus.RanToCompletion : task.Status)
             {
-                StateHasChanged();
+                // If it's already completed synchronously, no need to await and no
-            }
+                // need to issue a further render (we already rerender synchronously).
-            else
+                // Just need to make sure we propagate any errors.
-            {
+                case TaskStatus.RanToCompletion:
-                HandleException(task.Exception);
+                case TaskStatus.Canceled:
+                    break;
+                case TaskStatus.Faulted:
+                    HandleException(task.Exception);
+                    break;
+
+                // For incomplete tasks, automatically re-render on successful completion
+                default:
+                    try
+                    {
+                        await task;
+                        StateHasChanged();
+                    }
+                    catch (Exception ex)
+                    {
+                        // Either the task failed, or it was cancelled, or StateHasChanged threw.
+                        // We want to report task failure or StateHasChanged exceptions only.
+                        if (!task.IsCanceled)
+                        {
+                            HandleException(ex);
+                        }
+                    }
+
+                    break;
             }
         }
 
@@ -203,18 +220,12 @@ namespace Microsoft.AspNetCore.Components
         void IHandleEvent.HandleEvent(EventHandlerInvoker binding, UIEventArgs args)
         {
             var task = binding.Invoke(args);
+            ContinueAfterLifecycleTask(task);
 
             // After each event, we synchronously re-render (unless !ShouldRender())
             // This just saves the developer the trouble of putting "StateHasChanged();"
             // at the end of every event callback.
             StateHasChanged();
-
-            if (task.Status == TaskStatus.RanToCompletion)
-            {
-                return;
-            }
-
-            task.ContinueWith(ContinueAfterLifecycleTask);
         }
 
         void IHandleAfterRender.OnAfterRender()

src/Components/test/Microsoft.AspNetCore.Components.Test/ComponentBaseTest.cs

@@ -0,0 +1,332 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Components.RenderTree;
+using Microsoft.AspNetCore.Components.Test.Helpers;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Components.Test
+{
+    public class ComponentBaseTest
+    {
+        [Fact]
+        public void RunsOnInitWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onInitRuns = 0;
+            component.OnInitLogic = c => onInitRuns++;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onInitRuns);
+        }
+
+        [Fact]
+        public void RunsOnInitAsyncWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onInitAsyncRuns = 0;
+            component.RunsBaseOnInitAsync = false;
+            component.OnInitAsyncLogic = c =>
+            {
+                onInitAsyncRuns++;
+                return Task.CompletedTask;
+            };
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onInitAsyncRuns);
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void RunsOnInitAsyncAlsoOnBaseClassWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onInitAsyncRuns = 0;
+            component.RunsBaseOnInitAsync = true;
+            component.OnInitAsyncLogic = c =>
+            {
+                onInitAsyncRuns++;
+                return Task.CompletedTask;
+            };
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onInitAsyncRuns);
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void RunsOnParametersSetWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onParametersSetRuns = 0;
+            component.OnParametersSetLogic = c => onParametersSetRuns++;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onParametersSetRuns);
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void RunsOnParametersSetAsyncWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onParametersSetAsyncRuns = 0;
+            component.RunsBaseOnParametersSetAsync = false;
+            component.OnParametersSetAsyncLogic = c =>
+            {
+                onParametersSetAsyncRuns++;
+                return Task.CompletedTask;
+            };
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onParametersSetAsyncRuns);
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void RunsOnParametersSetAsyncAlsoOnBaseClassWhenRendered()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            int onParametersSetAsyncRuns = 0;
+            component.RunsBaseOnParametersSetAsync = true;
+            component.OnParametersSetAsyncLogic = c =>
+            {
+                onParametersSetAsyncRuns++;
+                return Task.CompletedTask;
+            };
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Equal(1, onParametersSetAsyncRuns);
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void RendersAfterParametersSetAsyncTaskIsCompleted()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            component.Counter = 1;
+            var parametersSetTask = new TaskCompletionSource<bool>();
+            component.RunsBaseOnParametersSetAsync = false;
+            component.OnParametersSetAsyncLogic = c => parametersSetTask.Task;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Single(renderer.Batches);
+
+            // Completes task started by OnParametersSetAsync
+            component.Counter = 2;
+            parametersSetTask.SetResult(true);
+
+            // Component should be rendered again
+            Assert.Equal(2, renderer.Batches.Count);
+        }
+
+        [Fact]
+        public void RendersAfterParametersSetAndInitAsyncTasksAreCompleted()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent();
+
+            component.Counter = 1;
+            var initTask = new TaskCompletionSource<bool>();
+            var parametersSetTask = new TaskCompletionSource<bool>();
+            component.RunsBaseOnInitAsync = true;
+            component.RunsBaseOnParametersSetAsync = true;
+            component.OnInitAsyncLogic = c => initTask.Task;
+            component.OnParametersSetAsyncLogic = c => parametersSetTask.Task;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Single(renderer.Batches);
+
+            // Completes task started by OnParametersSetAsync
+            component.Counter = 2;
+            parametersSetTask.SetResult(false);
+
+            // Component should be rendered again
+            Assert.Equal(2, renderer.Batches.Count);
+
+            // Completes task started by OnInitAsync
+            // NOTE: We will probably change this behavior. It would make more sense for the base class
+            // to wait until InitAsync is completed before proceeding with SetParametersAsync, rather
+            // that running the two lifecycle methods in parallel. This will come up as a requirement
+            // when implementing async server-side prerendering.
+            component.Counter = 3;
+            initTask.SetResult(true);
+
+            // Component should be rendered again
+            Assert.Equal(3, renderer.Batches.Count);
+        }
+
+        [Fact]
+        public void DoesNotRenderAfterOnInitAsyncTaskIsCancelled()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent() { Counter = 1 };
+            var initTask = new TaskCompletionSource<object>();
+            component.OnInitAsyncLogic = _ => initTask.Task;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Single(renderer.Batches);
+
+            // Cancel task started by OnInitAsync
+            component.Counter = 2;
+            initTask.SetCanceled();
+
+            // Component should not be rendered again
+            Assert.Single(renderer.Batches);
+        }
+
+        [Fact]
+        public void DoesNotRenderAfterOnParametersSetAsyncTaskIsCancelled()
+        {
+            // Arrange
+            var renderer = new TestRenderer();
+            var component = new TestComponent() { Counter = 1 };
+            var onParametersSetTask = new TaskCompletionSource<object>();
+            component.OnParametersSetAsyncLogic = _ => onParametersSetTask.Task;
+
+            // Act
+            var componentId = renderer.AssignRootComponentId(component);
+            renderer.RenderRootComponent(componentId);
+
+            // Assert
+            Assert.Single(renderer.Batches);
+
+            // Cancel task started by OnParametersSet
+            component.Counter = 2;
+            onParametersSetTask.SetCanceled();
+
+            // Component should not be rendered again
+            Assert.Single(renderer.Batches);
+        }
+
+        private class TestComponent : ComponentBase
+        {
+            public bool RunsBaseOnInit { get; set; } = true;
+
+            public bool RunsBaseOnInitAsync { get; set; } = true;
+
+            public bool RunsBaseOnParametersSet { get; set; } = true;
+
+            public bool RunsBaseOnParametersSetAsync { get; set; } = true;
+
+            public Action<TestComponent> OnInitLogic { get; set; }
+
+            public Func<TestComponent, Task> OnInitAsyncLogic { get; set; }
+
+            public Action<TestComponent> OnParametersSetLogic { get; set; }
+
+            public Func<TestComponent, Task> OnParametersSetAsyncLogic { get; set; }
+
+            public int Counter { get; set; }
+
+            protected override void BuildRenderTree(RenderTreeBuilder builder)
+            {
+                base.BuildRenderTree(builder);
+                builder.OpenElement(0, "p");
+                builder.AddContent(1, Counter);
+                builder.CloseElement();
+            }
+
+            protected override void OnInit()
+            {
+                if (RunsBaseOnInit)
+                {
+                    base.OnInit();
+                }
+
+                OnInitLogic?.Invoke(this);
+            }
+
+            protected override async Task OnInitAsync()
+            {
+                if (RunsBaseOnInitAsync)
+                {
+                    await base.OnInitAsync();
+                }
+
+                await OnInitAsyncLogic?.Invoke(this);
+            }
+
+            protected override void OnParametersSet()
+            {
+                if (RunsBaseOnParametersSet)
+                {
+                    base.OnParametersSet();
+                }
+
+                OnParametersSetLogic?.Invoke(this);
+            }
+
+            protected override async Task OnParametersSetAsync()
+            {
+                if (RunsBaseOnParametersSetAsync)
+                {
+                    await base.OnParametersSetAsync();
+                }
+
+                await OnParametersSetAsyncLogic?.Invoke(this);
+            }
+        }
+    }
+}

src/Components/test/shared/TestRenderer.cs

@@ -28,6 +28,9 @@ namespace Microsoft.AspNetCore.Components.Test.Helpers
         public new int AssignRootComponentId(IComponent component)
             => base.AssignRootComponentId(component);
 
+        public new void RenderRootComponent(int componentId)
+            => base.RenderRootComponent(componentId);
+
         public new void DispatchEvent(int componentId, int eventHandlerId, UIEventArgs args)
             => base.DispatchEvent(componentId, eventHandlerId, args);