Tests, tweaks, and other follow-ups to lazy-loading (#23947)

* Render 'OnNavigateError' fragment on unhandled exception in OnNavigateAsync

* Address first round of feedback from peer review

* Refactor OnNavigateAsync handling and fix tests

* Make OnNavigateAsync cancellation cooperative with user tasks

* Fix aggressive re-rendering and cancellation handling

* Fix up tests based on peer review

src/Components/Components/src/Routing/Router.cs

@@ -4,6 +4,7 @@
 #nullable disable warnings
 
 using System;
+using System.Runtime.ExceptionServices;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Linq;
@@ -72,7 +73,7 @@ namespace Microsoft.AspNetCore.Components.Routing
         /// <summary>
         /// Gets or sets a handler that should be called before navigating to a new page.
         /// </summary>
-        [Parameter] public EventCallback<NavigationContext> OnNavigateAsync { get; set; }
+        [Parameter] public Func<NavigationContext, Task> OnNavigateAsync { get; set; }
 
         private RouteTable Routes { get; set; }
 
@@ -194,51 +195,58 @@ namespace Microsoft.AspNetCore.Components.Routing
 
         private async ValueTask<bool> RunOnNavigateAsync(string path, Task previousOnNavigate)
         {
-            // If this router instance does not provide an OnNavigateAsync parameter
-            // then we render the component associated with the route as per usual.
-            if (!OnNavigateAsync.HasDelegate)
+            if (OnNavigateAsync == null)
             {
                 return true;
             }
 
-            // If we've already invoked a task and stored its CTS, then
-            // cancel that existing CTS.
+            // Cancel the CTS instead of disposing it, since disposing does not
+            // actually cancel and can cause unintended Object Disposed Exceptions.
+            // This effectivelly cancels the previously running task and completes it.
             _onNavigateCts?.Cancel();
-            // Then make sure that the task has been completed cancelled or
-            // completed before continuing with the execution of this current task.
+            // Then make sure that the task has been completely cancelled or completed
+            // before starting the next one. This avoid race conditions where the cancellation
+            // for the previous task was set but not fully completed by the time we get to this
+            // invocation.
             await previousOnNavigate;
 
-            // Create a new cancellation token source for this instance
             _onNavigateCts = new CancellationTokenSource();
             var navigateContext = new NavigationContext(path, _onNavigateCts.Token);
 
-            // Create a cancellation task based on the cancellation token
-            // associated with the current running task.
-            var cancellationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            navigateContext.CancellationToken.Register(state =>
-                ((TaskCompletionSource)state).SetResult(), cancellationTcs);
-
-            var task = OnNavigateAsync.InvokeAsync(navigateContext);
-
-            // If the user provided a Navigating render fragment, then show it.
-            if (Navigating != null && task.Status != TaskStatus.RanToCompletion)
+            try
+            {
+                if (Navigating != null)
+                {
+                    _renderHandle.Render(Navigating);
+                }
+                await OnNavigateAsync(navigateContext);
+                return true;
+            }
+            catch (OperationCanceledException e)
+            {
+                if (e.CancellationToken != navigateContext.CancellationToken)
+                {
+                    var rethrownException = new InvalidOperationException("OnNavigateAsync can only be cancelled via NavigateContext.CancellationToken.", e);
+                    _renderHandle.Render(builder => ExceptionDispatchInfo.Capture(rethrownException).Throw());
+                    return false;
+                }
+            }
+            catch (Exception e)
             {
-                _renderHandle.Render(Navigating);
+                _renderHandle.Render(builder => ExceptionDispatchInfo.Capture(e).Throw());
+                return false;
             }
 
-            var completedTask = await Task.WhenAny(task, cancellationTcs.Task);
-            return task == completedTask;
+            return false;
         }
 
         internal async Task RunOnNavigateWithRefreshAsync(string path, bool isNavigationIntercepted)
         {
             // We cache the Task representing the previously invoked RunOnNavigateWithRefreshAsync
-            // that is stored
+            // that is stored. Then we create a new one that represents our current invocation and store it
+            // globally for the next invocation. This allows us to check inside `RunOnNavigateAsync` if the
+            // previous OnNavigateAsync task has fully completed before starting the next one.
             var previousTask = _previousOnNavigateTask;
-            // Then we create a new one that represents our current invocation and store it
-            // globally for the next invocation. Note to the developer, if the WASM runtime
-            // support multi-threading then we'll need to implement the appropriate locks
-            // here to ensure that the cached previous task is overwritten incorrectly.
             var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
             _previousOnNavigateTask = tcs.Task;
             try

src/Components/Components/test/Routing/RouterTest.cs

@@ -2,103 +2,253 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
-using System.Linq;
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Routing;
 using Microsoft.AspNetCore.Components.Test.Helpers;
-using Microsoft.Extensions.DependencyModel;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
 using Moq;
 using Xunit;
+using Microsoft.AspNetCore.Components;
 
 namespace Microsoft.AspNetCore.Components.Test.Routing
 {
     public class RouterTest
     {
+        private readonly Router _router;
+        private readonly TestRenderer _renderer;
+
+        public RouterTest()
+        {
+            var services = new ServiceCollection();
+            services.AddSingleton<ILoggerFactory>(NullLoggerFactory.Instance);
+            services.AddSingleton<NavigationManager, TestNavigationManager>();
+            services.AddSingleton<INavigationInterception, TestNavigationInterception>();
+            var serviceProvider = services.BuildServiceProvider();
+
+            _renderer = new TestRenderer(serviceProvider);
+            _renderer.ShouldHandleExceptions = true;
+            _router = (Router)_renderer.InstantiateComponent<Router>();
+            _router.AppAssembly = Assembly.GetExecutingAssembly();
+            _router.Found = routeData => (builder) => builder.AddContent(0, "Rendering route...");
+            _renderer.AssignRootComponentId(_router);
+        }
+
         [Fact]
         public async Task CanRunOnNavigateAsync()
         {
             // Arrange
-            var router = CreateMockRouter();
             var called = false;
             async Task OnNavigateAsync(NavigationContext args)
             {
                 await Task.CompletedTask;
                 called = true;
             }
-            router.Object.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+            _router.OnNavigateAsync = OnNavigateAsync;
 
             // Act
-            await router.Object.RunOnNavigateWithRefreshAsync("http://example.com/jan", false);
+            await _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
 
             // Assert
             Assert.True(called);
         }
 
         [Fact]
-        public async Task CanCancelPreviousOnNavigateAsync()
+        public async Task CanHandleSingleFailedOnNavigateAsync()
+        {
+            // Arrange
+            var called = false;
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                called = true;
+                await Task.CompletedTask;
+                throw new Exception("This is an uncaught exception.");
+            }
+            _router.OnNavigateAsync = OnNavigateAsync;
+
+            // Act
+            await _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
+
+            // Assert
+            Assert.True(called);
+            Assert.Single(_renderer.HandledExceptions);
+            var unhandledException = _renderer.HandledExceptions[0];
+            Assert.Equal("This is an uncaught exception.", unhandledException.Message);
+        }
+
+        [Fact]
+        public async Task CanceledFailedOnNavigateAsyncDoesNothing()
+        {
+            // Arrange
+            var onNavigateInvoked = 0;
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                onNavigateInvoked += 1;
+                if (args.Path.EndsWith("jan"))
+                {
+                    await Task.Delay(Timeout.Infinite, args.CancellationToken);
+                    throw new Exception("This is an uncaught exception.");
+                }
+            }
+            var refreshCalled = false;
+            _renderer.OnUpdateDisplay = (renderBatch) =>
+            {
+                if (!refreshCalled)
+                {
+                    refreshCalled = true;
+                    return;
+                }
+                Assert.True(false, "OnUpdateDisplay called more than once.");
+            };
+            _router.OnNavigateAsync = OnNavigateAsync;
+
+            // Act
+            var janTask = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
+            var febTask = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/feb", false));
+
+            await janTask;
+            await febTask;
+
+            // Assert that we render the second route component and don't throw an exception
+            Assert.Empty(_renderer.HandledExceptions);
+            Assert.Equal(2, onNavigateInvoked);
+        }
+
+        [Fact]
+        public async Task CanHandleSingleCancelledOnNavigateAsync()
+        {
+            // Arrange
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                var tcs = new TaskCompletionSource<int>();
+                tcs.TrySetCanceled();
+                await tcs.Task;
+            }
+            _renderer.OnUpdateDisplay = (renderBatch) => Assert.True(false, "OnUpdateDisplay called more than once.");
+            _router.OnNavigateAsync = OnNavigateAsync;
+
+            // Act
+            await _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
+
+            // Assert
+            Assert.Single(_renderer.HandledExceptions);
+            var unhandledException = _renderer.HandledExceptions[0];
+            Assert.Equal("OnNavigateAsync can only be cancelled via NavigateContext.CancellationToken.", unhandledException.Message);
+        }
+
+        [Fact]
+        public async Task AlreadyCanceledOnNavigateAsyncDoesNothing()
+        {
+            // Arrange
+            var triggerCancel = new TaskCompletionSource();
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                if (args.Path.EndsWith("jan"))
+                {
+                    var tcs = new TaskCompletionSource();
+                    await triggerCancel.Task;
+                    tcs.TrySetCanceled();
+                    await tcs.Task;
+                }
+            }
+            var refreshCalled = false;
+            _renderer.OnUpdateDisplay = (renderBatch) =>
+            {
+                if (!refreshCalled)
+                {
+                    Assert.True(true);
+                    return;
+                }
+                Assert.True(false, "OnUpdateDisplay called more than once.");
+            };
+            _router.OnNavigateAsync = OnNavigateAsync;
+
+            // Act (start the operations then await them)
+            var jan = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
+            var feb = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/feb", false));
+            triggerCancel.TrySetResult();
+
+            await jan;
+            await feb;
+        }
+
+        [Fact]
+        public void CanCancelPreviousOnNavigateAsync()
         {
             // Arrange
-            var router = CreateMockRouter();
             var cancelled = "";
             async Task OnNavigateAsync(NavigationContext args)
             {
                 await Task.CompletedTask;
                 args.CancellationToken.Register(() => cancelled = args.Path);
             };
-            router.Object.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+            _router.OnNavigateAsync = OnNavigateAsync;
 
             // Act
-            await router.Object.RunOnNavigateWithRefreshAsync("jan", false);
-            await router.Object.RunOnNavigateWithRefreshAsync("feb", false);
+            _ = _router.RunOnNavigateWithRefreshAsync("jan", false);
+            _ = _router.RunOnNavigateWithRefreshAsync("feb", false);
 
             // Assert
             var expected = "jan";
-            Assert.Equal(cancelled, expected);
+            Assert.Equal(expected, cancelled);
         }
 
         [Fact]
         public async Task RefreshesOnceOnCancelledOnNavigateAsync()
         {
             // Arrange
-            var router = CreateMockRouter();
             async Task OnNavigateAsync(NavigationContext args)
             {
                 if (args.Path.EndsWith("jan"))
                 {
-                    await Task.Delay(Timeout.Infinite);
+                    await Task.Delay(Timeout.Infinite, args.CancellationToken);
+                }
+            };
+            var refreshCalled = false;
+            _renderer.OnUpdateDisplay = (renderBatch) =>
+            {
+                if (!refreshCalled)
+                {
+                    Assert.True(true);
+                    return;
                 }
+                Assert.True(false, "OnUpdateDisplay called more than once.");
             };
-            router.Object.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+            _router.OnNavigateAsync = OnNavigateAsync;
 
             // Act
-            var janTask = router.Object.RunOnNavigateWithRefreshAsync("jan", false);
-            var febTask = router.Object.RunOnNavigateWithRefreshAsync("feb", false);
+            var jan = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/jan", false));
+            var feb = _renderer.Dispatcher.InvokeAsync(() => _router.RunOnNavigateWithRefreshAsync("http://example.com/feb", false));
 
-            var janTaskException = await Record.ExceptionAsync(() => janTask);
-            var febTaskException = await Record.ExceptionAsync(() => febTask);
-
-            // Assert neither exceution threw an exception
-            Assert.Null(janTaskException);
-            Assert.Null(febTaskException);
-            // Assert refresh should've only been called once for the second route
-            router.Verify(x => x.Refresh(false), Times.Once());
+            await jan;
+            await feb;
         }
 
-        private Mock<Router> CreateMockRouter()
+        internal class TestNavigationManager : NavigationManager
         {
-            var router = new Mock<Router>() { CallBase = true };
-            router.Setup(x => x.Refresh(It.IsAny<bool>())).Verifiable();
-            return router;
+            public TestNavigationManager() =>
+                Initialize("https://www.example.com/subdir/", "https://www.example.com/subdir/jan");
+
+            protected override void NavigateToCore(string uri, bool forceLoad) => throw new NotImplementedException();
         }
 
-        [Route("jan")]
-        private class JanComponent : ComponentBase { }
+        internal sealed class TestNavigationInterception : INavigationInterception
+        {
+            public static readonly TestNavigationInterception Instance = new TestNavigationInterception();
+
+            public Task EnableNavigationInterceptionAsync()
+            {
+                return Task.CompletedTask;
+            }
+        }
 
         [Route("feb")]
-        private class FebComponent : ComponentBase { }
+        public class FebComponent : ComponentBase { }
+
+        [Route("jan")]
+        public class JanComponent : ComponentBase { }
     }
 }

src/Components/test/E2ETest/Tests/RoutingTest.cs

@@ -554,6 +554,19 @@ namespace Microsoft.AspNetCore.Components.E2ETest.Tests
             AssertDidNotLog("I'm not happening...");
         }
 
+        [Fact]
+        public void OnNavigate_CanRenderUIForExceptions()
+        {
+            var app = Browser.MountTestComponent<TestRouterWithOnNavigate>();
+
+            // Navigating from one page to another should
+            // cancel the previous OnNavigate Task
+            SetUrlViaPushState("/Other");
+
+            var errorUiElem = Browser.Exists(By.Id("blazor-error-ui"), TimeSpan.FromSeconds(10));
+            Assert.NotNull(errorUiElem);
+        }
+
         private long BrowserScrollY
         {
             get => (long)((IJavaScriptExecutor)Browser).ExecuteScript("return window.scrollY");

src/Components/test/testassets/BasicTestApp/RouterTest/TestRouterWithOnNavigate.razor

@@ -20,7 +20,8 @@
     private Dictionary<string, Func<NavigationContext, Task>> preNavigateTasks = new Dictionary<string, Func<NavigationContext, Task>>()
     {
         { "LongPage1", new Func<NavigationContext, Task>(TestLoadingPageShows) },
-        { "LongPage2", new Func<NavigationContext, Task>(TestOnNavCancel) }
+        { "LongPage2", new Func<NavigationContext, Task>(TestOnNavCancel) },
+        { "Other", new Func<NavigationContext, Task>(TestOnNavException) }
     };
 
     private async Task OnNavigateAsync(NavigationContext args)
@@ -43,4 +44,10 @@
         await Task.Delay(2000, args.CancellationToken);
         Console.WriteLine("I'm not happening...");
     }
+
+    public static async Task TestOnNavException(NavigationContext args)
+    {
+        await Task.CompletedTask;
+        throw new Exception("This is an uncaught exception.");
+    }
 }