Blazor 404 code re-execution is compatible with OnNavigateAsync

src/Components/Components/src/Routing/RouteTable.cs

@@ -78,6 +78,8 @@ public void Route(RouteContext routeContext)
 
     private static void ProcessParameters(InboundRouteEntry entry, RouteValueDictionary routeValues)
     {
+        routeValues.Remove(Router.AllowRenderDuringPendingNavigationKey);
+
         // Add null values for unused route parameters.
         if (entry.UnusedRouteParameterNames != null)
         {

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

@@ -42,6 +42,8 @@ static readonly IReadOnlyDictionary<string, object> _emptyParametersDictionary
 
     private bool _onNavigateCalled;
 
+    internal const string AllowRenderDuringPendingNavigationKey = "__BlazorAllowRenderDuringPendingNavigation";
+
     [Inject] private NavigationManager NavigationManager { get; set; }
 
     [Inject] private INavigationInterception NavigationInterception { get; set; }
@@ -220,11 +222,14 @@ private void ClearRouteCaches()
 
     internal virtual void Refresh(bool isNavigationIntercepted)
     {
+        var providerRouteData = RoutingStateProvider?.RouteData;
+        var allowRenderDuringPendingNavigation = TryConsumeAllowRenderDuringPendingNavigation(providerRouteData);
+
         // If an `OnNavigateAsync` task is currently in progress, then wait
         // for it to complete before rendering. Note: because _previousOnNavigateTask
         // is initialized to a CompletedTask on initialization, this will still
         // allow first-render to complete successfully.
-        if (_previousOnNavigateTask.Status != TaskStatus.RanToCompletion)
+        if (_previousOnNavigateTask.Status != TaskStatus.RanToCompletion && !allowRenderDuringPendingNavigation)
         {
             if (Navigating != null)
             {
@@ -239,7 +244,7 @@ internal virtual void Refresh(bool isNavigationIntercepted)
         ComponentsActivityHandle activityHandle;
 
         // In order to avoid routing twice we check for RouteData
-        if (RoutingStateProvider?.RouteData is { } endpointRouteData)
+        if (providerRouteData is { } endpointRouteData)
         {
             activityHandle = RecordDiagnostics(endpointRouteData.PageType.FullName, endpointRouteData.Template);
 
@@ -312,6 +317,17 @@ internal virtual void Refresh(bool isNavigationIntercepted)
         _renderHandle.ComponentActivitySource?.StopNavigateActivity(activityHandle, null);
     }
 
+    private static bool TryConsumeAllowRenderDuringPendingNavigation(RouteData? routeData)
+    {
+        if (routeData?.RouteValues.TryGetValue(AllowRenderDuringPendingNavigationKey, out var value) == true && value is true)
+        {
+            (routeData.RouteValues as IDictionary<string, object?>)?.Remove(AllowRenderDuringPendingNavigationKey);
+            return true;
+        }
+
+        return false;
+    }
+
     private ComponentsActivityHandle RecordDiagnostics(string componentType, string template)
     {
         ComponentsActivityHandle activityHandle = default;

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -87,7 +87,8 @@ await _renderer.InitializeStandardComponentServicesAsync(
             context,
             componentType: pageComponent,
             handler: result.HandlerName,
-            form: result.HandlerName != null && context.Request.HasFormContentType ? await context.Request.ReadFormAsync() : null);
+            form: result.HandlerName != null && context.Request.HasFormContentType ? await context.Request.ReadFormAsync() : null,
+            allowRenderingDuringPendingNavigation: isReExecuted);
 
         // Matches MVC's MemoryPoolHttpResponseStreamWriterFactory.DefaultBufferSize
         var defaultBufferSize = 16 * 1024;

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.cs

@@ -54,6 +54,8 @@ internal partial class EndpointHtmlRenderer : StaticHtmlRenderer, IComponentPrer
 
     private string _notFoundUrl = string.Empty;
 
+    private const string AllowRenderDuringPendingNavigationKey = "__BlazorAllowRenderDuringPendingNavigation";
+
     public EndpointHtmlRenderer(IServiceProvider serviceProvider, ILoggerFactory loggerFactory)
         : base(serviceProvider, loggerFactory)
     {
@@ -81,12 +83,13 @@ internal async Task InitializeStandardComponentServicesAsync(
         HttpContext httpContext,
         [DynamicallyAccessedMembers(Component)] Type? componentType = null,
         string? handler = null,
-        IFormCollection? form = null)
+        IFormCollection? form = null,
+        bool allowRenderingDuringPendingNavigation = false)
     {
         var navigationManager = httpContext.RequestServices.GetRequiredService<NavigationManager>();
         ((IHostEnvironmentNavigationManager)navigationManager)?.Initialize(
-            GetContextBaseUri(httpContext.Request), 
-            GetFullUri(httpContext.Request), 
+            GetContextBaseUri(httpContext.Request),
+            GetFullUri(httpContext.Request),
             uri => GetErrorHandledTask(OnNavigateTo(uri)));
 
         navigationManager?.OnNotFound += (sender, args) => NotFoundEventArgs = args;
@@ -132,7 +135,13 @@ internal async Task InitializeStandardComponentServicesAsync(
         {
             // Saving RouteData to avoid routing twice in Router component
             var routingStateProvider = httpContext.RequestServices.GetRequiredService<EndpointRoutingStateProvider>();
-            routingStateProvider.RouteData = new RouteData(componentType, httpContext.GetRouteData().Values);
+            var routeValues = new RouteValueDictionary(httpContext.GetRouteData().Values);
+            if (allowRenderingDuringPendingNavigation)
+            {
+                routeValues[AllowRenderDuringPendingNavigationKey] = true;
+            }
+
+            routingStateProvider.RouteData = new RouteData(componentType, routeValues);
             if (httpContext.GetEndpoint() is RouteEndpoint routeEndpoint)
             {
                 routingStateProvider.RouteData.Template = routeEndpoint.RoutePattern.RawText;

src/Components/test/E2ETest/ServerRenderingTests/NoInteractivityTest.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net.Http;
+using System;
 using Components.TestServer.RazorComponents;
 using Microsoft.AspNetCore.Components.E2ETest;
 using Microsoft.AspNetCore.Components.E2ETest.Infrastructure;
@@ -125,6 +126,42 @@ public void BrowserNavigationToNotExistingPath_ReExecutesTo404(bool streaming)
         AssertReExecutionPageRendered();
     }
 
+    [Fact]
+    public void BrowserNavigationToNotExistingPath_WithOnNavigateAsync_ReExecutesTo404()
+    {
+        AppContext.SetSwitch("Microsoft.AspNetCore.Components.Endpoints.NavigationManager.DisableThrowNavigationException", isEnabled: true);
+
+        // using query for controlling router parameters does not work in re-execution scenario, we have to rely on other communication channel
+        const string useOnNavigateAsyncSwitch = "Components.TestServer.RazorComponents.UseOnNavigateAsync";
+        AppContext.SetSwitch(useOnNavigateAsyncSwitch, true);
+        try
+        {
+            Navigate($"{ServerPathBase}/reexecution/not-existing-page");
+            AssertReExecutionPageRendered();
+        }
+        finally
+        {
+            AppContext.SetSwitch(useOnNavigateAsyncSwitch, false);
+        }
+    }
+
+    [Fact]
+    public void BrowserNavigationToLazyLoadedRoute_WaitsForOnNavigateAsyncGuard()
+    {
+        const string navigationGuardSwitch = "Components.TestServer.RazorComponents.UseNavigationCompletionGuard";
+        AppContext.SetSwitch(navigationGuardSwitch, true);
+
+        try
+        {
+            Navigate($"{ServerPathBase}/routing/with-lazy-assembly");
+            Browser.Equal("Lazy route rendered", () => Browser.Exists(By.Id("lazy-route-status")).Text);
+        }
+        finally
+        {
+            AppContext.SetSwitch(navigationGuardSwitch, false);
+        }
+    }
+
     private void AssertReExecutionPageRendered() =>
         Browser.Equal("Welcome On Page Re-executed After Not Found Event", () => Browser.Exists(By.Id("test-info")).Text);
 
@@ -434,3 +471,6 @@ public void StatusCodePagesWithReExecution()
         Browser.Equal("Re-executed page", () => Browser.Title);
     }
 }
+
+#pragma warning restore RS0037 // PublicAPI files must include '#nullable enable'
+#pragma warning restore RS0016 // Add public types and members to the declared API

src/Components/test/testassets/Components.TestServer/RazorComponents/App.razor

@@ -3,6 +3,10 @@
 @using Components.WasmMinimal.Pages.NotFound
 @using TestContentPackage.NotFound
 @using Components.TestServer.RazorComponents
+@using Microsoft.AspNetCore.Components
+@using Microsoft.AspNetCore.Components.Routing
+@using Microsoft.AspNetCore.Components.Web
+@using System.Threading.Tasks
 
 @code {
     [Parameter]
@@ -17,8 +21,12 @@
     [SupplyParameterFromQuery(Name = "appSetsEventArgsPath")]
     public bool AppSetsEventArgsPath { get; set; }
 
+    private const string UseOnNavigateAsyncSwitchName = "Components.TestServer.RazorComponents.UseOnNavigateAsync";
+
     private Type? NotFoundPageType { get; set; }
     private NavigationManager _navigationManager = default!;
+    private bool ShouldDelayOnNavigateAsync =>
+        AppContext.TryGetSwitch(UseOnNavigateAsyncSwitchName, out var switchEnabled) && switchEnabled;
 
     [Inject]
     private NavigationManager NavigationManager
@@ -70,6 +78,26 @@
             _navigationManager.OnNotFound -= OnNotFoundEvent;
         }
     }
+
+    private Task HandleOnNavigateAsync(NavigationContext args)
+    {
+        if (NavigationCompletionTracker.TryGetGuardTask(args.Path, out var guardTask))
+        {
+            return guardTask;
+        }
+
+        if (!ShouldDelayOnNavigateAsync)
+        {
+            return Task.CompletedTask;
+        }
+
+        return PerformOnNavigateAsyncWork();
+    }
+
+    private async Task PerformOnNavigateAsyncWork()
+    {
+        await Task.Yield();
+    }
 }
 
 <!DOCTYPE html>
@@ -93,7 +121,7 @@
     {
         @if (NotFoundPageType is not null)
         {
-            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" NotFoundPage="NotFoundPageType">
+            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" NotFoundPage="NotFoundPageType" OnNavigateAsync="HandleOnNavigateAsync">
                 <Found Context="routeData">
                     <RouteView RouteData="@routeData" />
                     <FocusOnNavigate RouteData="@routeData" Selector="[data-focus-on-navigate]" />
@@ -102,7 +130,7 @@
         }
         else
         {
-            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }">
+            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" OnNavigateAsync="HandleOnNavigateAsync">
                 <Found Context="routeData">
                     <RouteView RouteData="@routeData" />
                     <FocusOnNavigate RouteData="@routeData" Selector="[data-focus-on-navigate]" />

src/Components/test/testassets/Components.TestServer/RazorComponents/NavigationCompletionTracker.cs

@@ -0,0 +1,65 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Components.TestServer.RazorComponents;
+
+internal static class NavigationCompletionTracker
+{
+    internal const string GuardSwitchName = "Components.TestServer.RazorComponents.UseNavigationCompletionGuard";
+
+    private const string TrackedPathSuffix = "with-lazy-assembly";
+    private static int _isNavigationTracked;
+    private static int _isNavigationCompleted;
+
+    public static bool TryGetGuardTask(string? path, out Task guardTask)
+    {
+        if (!IsGuardEnabledForPath(path))
+        {
+            guardTask = Task.CompletedTask;
+            return false;
+        }
+
+        guardTask = TrackNavigationAsync();
+        return true;
+    }
+
+    public static void AssertNavigationCompleted()
+    {
+        if (Volatile.Read(ref _isNavigationTracked) == 1 && Volatile.Read(ref _isNavigationCompleted) == 0)
+        {
+            throw new InvalidOperationException("Navigation finished before OnNavigateAsync work completed.");
+        }
+
+        Volatile.Write(ref _isNavigationTracked, 0);
+    }
+
+    private static bool IsGuardEnabledForPath(string? path)
+    {
+        if (!AppContext.TryGetSwitch(GuardSwitchName, out var isEnabled) || !isEnabled)
+        {
+            return false;
+        }
+
+        return path is not null && path.EndsWith(TrackedPathSuffix, StringComparison.OrdinalIgnoreCase);
+    }
+
+    private static async Task TrackNavigationAsync()
+    {
+        Volatile.Write(ref _isNavigationTracked, 1);
+        Volatile.Write(ref _isNavigationCompleted, 0);
+
+        try
+        {
+            await Task.Yield();
+            await Task.Delay(TimeSpan.FromMilliseconds(50)).ConfigureAwait(false);
+        }
+        finally
+        {
+            Volatile.Write(ref _isNavigationCompleted, 1);
+        }
+    }
+}

src/Components/test/testassets/Components.TestServer/RazorComponents/Pages/Routing/WithLazyAssembly.razor

@@ -0,0 +1,12 @@
+@page "/routing/with-lazy-assembly"
+@using Components.TestServer.RazorComponents;
+
+<h1 id="lazy-route-status">Lazy route rendered</h1>
+
+@code
+{
+    protected override void OnInitialized()
+    {
+        NavigationCompletionTracker.AssertNavigationCompleted();
+    }
+}