Spruce up async handling in OnNavigateAsync callback in Blazor router (#23835)

* Spruce up async handling in OnNavigateAsync

* Apply suggestions from code review

Co-authored-by: Günther Foidl <gue@korporal.at>

* Ensure previous task awaited before starting next one

* Apply suggestions from code review

Co-authored-by: Pranav K <prkrishn@hotmail.com>

* Validate no exceptions throw on multiple invocations

* Address feedback from peer review

Co-authored-by: Günther Foidl <gue@korporal.at>
Co-authored-by: Pranav K <prkrishn@hotmail.com>

src/Components/Components/src/Properties/AssemblyInfo.cs

@@ -8,3 +8,5 @@ using System.Runtime.CompilerServices;
 [assembly: InternalsVisibleTo("Microsoft.AspNetCore.Components.Tests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f33a29044fa9d740c9b3213a93e57c84b472c84e0b8a0e1ae48e67a9f8f6de9d5f7f3d52ac23e48ac51801f1dc950abe901da34d2a9e3baadb141a17c77ef3c565dd5ee5054b91cf63bb3c6ab83f72ab3aafe93d0fc3c2348b764fafb0b1c0733de51459aeab46580384bf9d74c4e28164b7cde247f891ba07891c9d872ad2bb")]
 [assembly: InternalsVisibleTo("Microsoft.AspNetCore.Components.Web.Tests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f33a29044fa9d740c9b3213a93e57c84b472c84e0b8a0e1ae48e67a9f8f6de9d5f7f3d52ac23e48ac51801f1dc950abe901da34d2a9e3baadb141a17c77ef3c565dd5ee5054b91cf63bb3c6ab83f72ab3aafe93d0fc3c2348b764fafb0b1c0733de51459aeab46580384bf9d74c4e28164b7cde247f891ba07891c9d872ad2bb")]
 [assembly: InternalsVisibleTo("Microsoft.AspNetCore.Components.Web.Extensions.Tests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f33a29044fa9d740c9b3213a93e57c84b472c84e0b8a0e1ae48e67a9f8f6de9d5f7f3d52ac23e48ac51801f1dc950abe901da34d2a9e3baadb141a17c77ef3c565dd5ee5054b91cf63bb3c6ab83f72ab3aafe93d0fc3c2348b764fafb0b1c0733de51459aeab46580384bf9d74c4e28164b7cde247f891ba07891c9d872ad2bb")]
+
+[assembly: InternalsVisibleTo("DynamicProxyGenAssembly2, PublicKey=0024000004800000940000000602000000240000525341310004000001000100c547cac37abd99c8db225ef2f6c8a3602f3b3606cc9891605d02baa56104f4cfc0734aa39b93bf7852f7d9266654753cc297e7d2edfe0bac1cdcf9f717241550e0a7b191195b7667bb4f64bcb8e2121380fd1d9d46ad2d92d2d15605093924cceaf74c4861eff62abf69b9291ed0a340e113be11e6a7d3113e92484cf7045cc7")]

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

@@ -31,6 +31,8 @@ namespace Microsoft.AspNetCore.Components.Routing
 
         private CancellationTokenSource _onNavigateCts;
 
+        private Task _previousOnNavigateTask = Task.CompletedTask;
+
         private readonly HashSet<Assembly> _assemblies = new HashSet<Assembly>();
 
         private bool _onNavigateCalled = false;
@@ -112,7 +114,8 @@ namespace Microsoft.AspNetCore.Components.Routing
             if (!_onNavigateCalled)
             {
                 _onNavigateCalled = true;
-                await RunOnNavigateAsync(NavigationManager.ToBaseRelativePath(_locationAbsolute));
+                await RunOnNavigateWithRefreshAsync(NavigationManager.ToBaseRelativePath(_locationAbsolute), isNavigationIntercepted: false);
+                return;
             }
 
             Refresh(isNavigationIntercepted: false);
@@ -122,7 +125,6 @@ namespace Microsoft.AspNetCore.Components.Routing
         public void Dispose()
         {
             NavigationManager.LocationChanged -= OnLocationChanged;
-            _onNavigateCts?.Dispose();
         }
 
         private static string StringUntilAny(string str, char[] chars)
@@ -147,7 +149,7 @@ namespace Microsoft.AspNetCore.Components.Routing
 
         }
 
-        private void Refresh(bool isNavigationIntercepted)
+        internal virtual void Refresh(bool isNavigationIntercepted)
         {
             RefreshRouteTable();
 
@@ -190,18 +192,21 @@ namespace Microsoft.AspNetCore.Components.Routing
             }
         }
 
-        private async Task RunOnNavigateAsync(string path)
+        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)
             {
-                return;
+                return true;
             }
 
             // If we've already invoked a task and stored its CTS, then
-            // cancel the existing task.
+            // cancel that existing CTS.
-            _onNavigateCts?.Dispose();
+            _onNavigateCts?.Cancel();
+            // Then make sure that the task has been completed cancelled or
+            // completed before continuing with the execution of this current task.
+            await previousOnNavigate;
 
             // Create a new cancellation token source for this instance
             _onNavigateCts = new CancellationTokenSource();
@@ -209,9 +214,9 @@ namespace Microsoft.AspNetCore.Components.Routing
 
             // Create a cancellation task based on the cancellation token
             // associated with the current running task.
-            var cancellationTaskSource = new TaskCompletionSource();
+            var cancellationTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
             navigateContext.CancellationToken.Register(state =>
-                ((TaskCompletionSource)state).SetResult(), cancellationTaskSource);
+                ((TaskCompletionSource)state).SetResult(), cancellationTcs);
 
             var task = OnNavigateAsync.InvokeAsync(navigateContext);
 
@@ -221,13 +226,34 @@ namespace Microsoft.AspNetCore.Components.Routing
                 _renderHandle.Render(Navigating);
             }
 
-            await Task.WhenAny(task, cancellationTaskSource.Task);
+            var completedTask = await Task.WhenAny(task, cancellationTcs.Task);
+            return task == completedTask;
         }
 
-        private async Task RunOnNavigateWithRefreshAsync(string path, bool isNavigationIntercepted)
+        internal async Task RunOnNavigateWithRefreshAsync(string path, bool isNavigationIntercepted)
         {
-            await RunOnNavigateAsync(path);
+            // We cache the Task representing the previously invoked RunOnNavigateWithRefreshAsync
-            Refresh(isNavigationIntercepted);
+            // that is stored
+            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
+            {
+                // And pass an indicator for the previous task to the currently running one.
+                var shouldRefresh = await RunOnNavigateAsync(path, previousTask);
+                if (shouldRefresh)
+                {
+                    Refresh(isNavigationIntercepted);
+                }
+            }
+            finally
+            {
+                tcs.SetResult();
+            }
         }
 
         private void OnLocationChanged(object sender, LocationChangedEventArgs args)

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

@@ -0,0 +1,104 @@
+// 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.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 Moq;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Components.Test.Routing
+{
+    public class RouterTest
+    {
+        [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);
+
+            // Act
+            await router.Object.RunOnNavigateWithRefreshAsync("http://example.com/jan", false);
+
+            // Assert
+            Assert.True(called);
+        }
+
+        [Fact]
+        public async Task 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);
+
+            // Act
+            await router.Object.RunOnNavigateWithRefreshAsync("jan", false);
+            await router.Object.RunOnNavigateWithRefreshAsync("feb", false);
+
+            // Assert
+            var expected = "jan";
+            Assert.Equal(cancelled, expected);
+        }
+
+        [Fact]
+        public async Task RefreshesOnceOnCancelledOnNavigateAsync()
+        {
+            // Arrange
+            var router = CreateMockRouter();
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                if (args.Path.EndsWith("jan"))
+                {
+                    await Task.Delay(Timeout.Infinite);
+                }
+            };
+            router.Object.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+
+            // Act
+            var janTask = router.Object.RunOnNavigateWithRefreshAsync("jan", false);
+            var febTask = router.Object.RunOnNavigateWithRefreshAsync("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());
+        }
+
+        private Mock<Router> CreateMockRouter()
+        {
+            var router = new Mock<Router>() { CallBase = true };
+            router.Setup(x => x.Refresh(It.IsAny<bool>())).Verifiable();
+            return router;
+        }
+
+        [Route("jan")]
+        private class JanComponent : ComponentBase { }
+
+        [Route("feb")]
+        private class FebComponent : ComponentBase { }
+    }
+}