Spruce up async handling in OnNavigateAsync

Author: captainsafia

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

@@ -23,7 +23,7 @@ public class Router : IComponent, IHandleAfterRender, IDisposable
         static readonly ReadOnlyDictionary<string, object> _emptyParametersDictionary
             = new ReadOnlyDictionary<string, object>(new Dictionary<string, object>());
 
-        RenderHandle _renderHandle;
+        internal RenderHandle _renderHandle { get; set; }
         string _baseUri;
         string _locationAbsolute;
         bool _navigationInterceptionEnabled;
@@ -35,7 +35,7 @@ static readonly ReadOnlyDictionary<string, object> _emptyParametersDictionary
 
         private bool _onNavigateCalled = false;
 
-        [Inject] private NavigationManager NavigationManager { get; set; }
+        [Inject] internal NavigationManager NavigationManager { get; set; }
 
         [Inject] private INavigationInterception NavigationInterception { get; set; }
 
@@ -72,7 +72,7 @@ static readonly ReadOnlyDictionary<string, object> _emptyParametersDictionary
         /// </summary>
         [Parameter] public EventCallback<NavigationContext> OnNavigateAsync { get; set; }
 
-        private RouteTable Routes { get; set; }
+        internal RouteTable Routes { get; set; }
 
         /// <inheritdoc />
         public void Attach(RenderHandle renderHandle)
@@ -109,20 +109,24 @@ public async Task SetParametersAsync(ParameterView parameters)
                 throw new InvalidOperationException($"The {nameof(Router)} component requires a value for the parameter {nameof(NotFound)}.");
             }
 
+            var shouldRefresh = true;
             if (!_onNavigateCalled)
             {
                 _onNavigateCalled = true;
-                await RunOnNavigateAsync(NavigationManager.ToBaseRelativePath(_locationAbsolute));
+                shouldRefresh = await RunOnNavigateAsync(NavigationManager.ToBaseRelativePath(_locationAbsolute));
+            }
+
+            if (shouldRefresh)
+            {
+                Refresh(isNavigationIntercepted: false);
             }
 
-            Refresh(isNavigationIntercepted: false);
         }
 
         /// <inheritdoc />
         public void Dispose()
         {
             NavigationManager.LocationChanged -= OnLocationChanged;
-            _onNavigateCts?.Dispose();
         }
 
         private static string StringUntilAny(string str, char[] chars)
@@ -190,28 +194,28 @@ private void Refresh(bool isNavigationIntercepted)
             }
         }
 
-        private async Task RunOnNavigateAsync(string path)
+        internal async Task<bool> RunOnNavigateAsync(string path)
         {
             // 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 await Task.FromResult(true);
             }
 
             // If we've already invoked a task and stored its CTS, then
             // cancel the existing task.
-            _onNavigateCts?.Dispose();
+            _onNavigateCts?.Cancel();
 
             // 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 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,16 +225,24 @@ private async Task RunOnNavigateAsync(string path)
                 _renderHandle.Render(Navigating);
             }
 
-            await Task.WhenAny(task, cancellationTaskSource.Task);
+            var completedTask = await Task.WhenAny(task, cancellationTcs.Task);
+            if (completedTask == task)
+            {
+                return await Task.FromResult(true);
+            }
+            return await Task.FromResult(false);
         }
 
         private async Task RunOnNavigateWithRefreshAsync(string path, bool isNavigationIntercepted)
         {
-            await RunOnNavigateAsync(path);
-            Refresh(isNavigationIntercepted);
+            var shouldRefresh = await RunOnNavigateAsync(path);
+            if (shouldRefresh)
+            {
+                Refresh(isNavigationIntercepted);
+            }
         }
 
-        private void OnLocationChanged(object sender, LocationChangedEventArgs args)
+        internal void OnLocationChanged(object sender, LocationChangedEventArgs args)
         {
             _locationAbsolute = args.Location;
             if (_renderHandle.IsInitialized && Routes != null)

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

@@ -0,0 +1,124 @@
+// 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.Threading.Tasks;
+using Microsoft.AspNetCore.Components;
+using Microsoft.AspNetCore.Components.Routing;
+using Microsoft.AspNetCore.Components.Test.Helpers;
+using Microsoft.Extensions.DependencyModel;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Components.Test.Routing
+{
+    public class RouterTest
+    {
+        [Fact]
+        public void CanRunOnNavigateViaLocationChangeAsync()
+        {
+            // Arrange
+            var router = CreateMockRouter();
+            var called = false;
+            async Task OnNavigateAsync(NavigationContext args)
+            {
+                await Task.CompletedTask;
+                called = true;
+            }
+            router.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+
+            // Act
+            router.OnLocationChanged(null, new LocationChangedEventArgs("http://example.com/jan", false));
+
+            // Assert
+            Assert.True(called);
+        }
+
+        [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.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+
+            // Act
+            router.OnLocationChanged(null, new LocationChangedEventArgs("http://example.com/jan", false));
+            router.OnLocationChanged(null, new LocationChangedEventArgs("http://example.com/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(5000);
+                }
+                if (args.Path.EndsWith("feb"))
+                {
+                    await Task.CompletedTask;
+                }
+            };
+            router.OnNavigateAsync = new EventCallbackFactory().Create<NavigationContext>(router, OnNavigateAsync);
+
+            // Act
+            var janTask = router.RunOnNavigateAsync("jan");
+            var febTask = router.RunOnNavigateAsync("feb");
+
+            await janTask;
+            await febTask;
+
+            // Assert
+            Assert.False(janTask.Result);
+            Assert.True(febTask.Result);
+        }
+
+        private Router CreateMockRouter()
+        {
+            var router = new Router();
+            var renderer = new TestRenderer();
+            router._renderHandle = new RenderHandle(renderer, 0);
+            router.Routes = RouteTableFactory.Create(new[] { typeof(JanComponent), typeof(FebComponent) });
+            router.NavigationManager = new TestNavigationManager();
+            return router;
+        }
+
+        [Route("jan")]
+        private class JanComponent : ComponentBase { }
+
+        [Route("feb")]
+        private class FebComponent : ComponentBase { }
+
+        private class TestNavigationManager : NavigationManager
+        {
+            public TestNavigationManager()
+            {
+                Initialize("http://example.com/", "http://example.com/months");
+            }
+
+            public new void Initialize(string baseUri, string uri)
+            {
+                base.Initialize(baseUri, uri);
+            }
+
+            protected override void NavigateToCore(string uri, bool forceLoad)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+    }
+}