[Blazor] Prerendered state (#50742)

[Blazor] Adds support for persting prerendered state on Blazor Web applications.
* Persists state both for server and webassembly as necessary.
* Initializes the state when a given interactive runtime is initialized and renders the first set of components.
  * On WebAssembly, this is the first time the app starts.
  * On Server this happens every time a circuit starts.
* The state is available during the first render, until the components reach quiescence.

The approach we follow is different for server and webassembly:
* On Server, we support initializing the circuit with an empty set of descriptors and in that case, we delay initialization until the first `UpdateRootComponents` call is issued.
  * This is because it's hard to deal with the security constraints imposed by starting a new circuit multiple times, and its easier to handle them within UpdateRootComponents. We might switch this approach in the future to go through `StartCircuit` too.
* On WebAssembly, we query for the initial set of webassembly components when we are starting the runtime in a Blazor Web Scenario.
  * We do this because Blazor WebAssembly offers a programatic API to render root components at a given location defined by their selectors, so we need to make sure that those components can receive state at the same time the initial set of WebAssembly components added to the page.

There are a set of tests validating different behaviors with regards to enhanced navigation and streaming rendering, as well as making sure that auto mode can access the state on Server and WebAssembly, and that Server gets new state every time a circuit is opened.

Author: javiercn

src/Components/Components/src/IPersistentComponentStateStore.cs

@@ -20,4 +20,11 @@ public interface IPersistentComponentStateStore
     /// <param name="state">The serialized state to persist.</param>
     /// <returns>A <see cref="Task" /> that completes when the state is persisted to disk.</returns>
     Task PersistStateAsync(IReadOnlyDictionary<string, byte[]> state);
+
+    /// <summary>
+    /// Returns a value that indicates whether the store supports the given <see cref="IComponentRenderMode"/>.
+    /// </summary>
+    /// <param name="renderMode">The <see cref="IComponentRenderMode"/> in question.</param>
+    /// <returns><c>true</c> if the render mode is supported by the store, otherwise <c>false</c>.</returns>
+    bool SupportsRenderMode(IComponentRenderMode renderMode) => true;
 }

src/Components/Components/src/Infrastructure/ComponentStatePersistenceManager.cs

@@ -11,17 +11,18 @@ namespace Microsoft.AspNetCore.Components.Infrastructure;
 /// </summary>
 public class ComponentStatePersistenceManager
 {
+    private readonly List<PersistComponentStateRegistration> _registeredCallbacks = new();
+    private readonly ILogger<ComponentStatePersistenceManager> _logger;
+
     private bool _stateIsPersisted;
-    private readonly List<Func<Task>> _pauseCallbacks = new();
     private readonly Dictionary<string, byte[]> _currentState = new(StringComparer.Ordinal);
-    private readonly ILogger<ComponentStatePersistenceManager> _logger;
 
     /// <summary>
     /// Initializes a new instance of <see cref="ComponentStatePersistenceManager"/>.
     /// </summary>
     public ComponentStatePersistenceManager(ILogger<ComponentStatePersistenceManager> logger)
     {
-        State = new PersistentComponentState(_currentState, _pauseCallbacks);
+        State = new PersistentComponentState(_currentState, _registeredCallbacks);
         _logger = logger;
     }
 
@@ -48,43 +49,100 @@ public async Task RestoreStateAsync(IPersistentComponentStateStore store)
     /// <param name="renderer">The <see cref="Renderer"/> that components are being rendered.</param>
     /// <returns>A <see cref="Task"/> that will complete when the state has been restored.</returns>
     public Task PersistStateAsync(IPersistentComponentStateStore store, Renderer renderer)
-        => PersistStateAsync(store, renderer.Dispatcher);
-
-    /// <summary>
-    /// Persists the component application state into the given <see cref="IPersistentComponentStateStore"/>.
-    /// </summary>
-    /// <param name="store">The <see cref="IPersistentComponentStateStore"/> to restore the application state from.</param>
-    /// <param name="dispatcher">The <see cref="Dispatcher"/> corresponding to the components' renderer.</param>
-    /// <returns>A <see cref="Task"/> that will complete when the state has been restored.</returns>
-    public Task PersistStateAsync(IPersistentComponentStateStore store, Dispatcher dispatcher)
     {
         if (_stateIsPersisted)
         {
             throw new InvalidOperationException("State already persisted.");
         }
 
-        _stateIsPersisted = true;
-
-        return dispatcher.InvokeAsync(PauseAndPersistState);
+        return renderer.Dispatcher.InvokeAsync(PauseAndPersistState);
 
         async Task PauseAndPersistState()
         {
             State.PersistingState = true;
-            await PauseAsync();
+
+            if (store is IEnumerable<IPersistentComponentStateStore> compositeStore)
+            {
+                // We only need to do inference when there is more than one store. This is determined by
+                // the set of rendered components.
+                InferRenderModes(renderer);
+
+                // Iterate over each store and give it a chance to run against the existing declared
+                // render modes. After we've run through a store, we clear the current state so that
+                // the next store can start with a clean slate.
+                foreach (var store in compositeStore)
+                {
+                    await PersistState(store);
+                    _currentState.Clear();
+                }
+            }
+            else
+            {
+                await PersistState(store);
+            }
+
             State.PersistingState = false;
+            _stateIsPersisted = true;
+        }
 
+        async Task PersistState(IPersistentComponentStateStore store)
+        {
+            await PauseAsync(store);
             await store.PersistStateAsync(_currentState);
         }
     }
 
-    internal Task PauseAsync()
+    private void InferRenderModes(Renderer renderer)
+    {
+        for (var i = 0; i < _registeredCallbacks.Count; i++)
+        {
+            var registration = _registeredCallbacks[i];
+            if (registration.RenderMode != null)
+            {
+                // Explicitly set render mode, so nothing to do.
+                continue;
+            }
+
+            if (registration.Callback.Target is IComponent component)
+            {
+                var componentRenderMode = renderer.GetComponentRenderMode(component);
+                if (componentRenderMode != null)
+                {
+                    _registeredCallbacks[i] = new PersistComponentStateRegistration(registration.Callback, componentRenderMode);
+                }
+                else
+                {
+                    // If we can't find a render mode, it's an SSR only component and we don't need to
+                    // persist its state at all.
+                    _registeredCallbacks[i] = default;
+                }
+                continue;
+            }
+
+            throw new InvalidOperationException(
+                $"The registered callback {registration.Callback.Method.Name} must be associated with a component or define" +
+                $" an explicit render mode type during registration.");
+        }
+    }
+
+    internal Task PauseAsync(IPersistentComponentStateStore store)
     {
         List<Task>? pendingCallbackTasks = null;
 
-        for (var i = 0; i < _pauseCallbacks.Count; i++)
+        for (var i = 0; i < _registeredCallbacks.Count; i++)
         {
-            var callback = _pauseCallbacks[i];
-            var result = ExecuteCallback(callback, _logger);
+            var registration = _registeredCallbacks[i];
+
+            if (!store.SupportsRenderMode(registration.RenderMode!))
+            {
+                // The callback does not have an associated render mode and we are in a multi-store scenario.
+                // Otherwise, in a single store scenario, we just run the callback.
+                // If the registration callback is null, it's because it was associated with a component and we couldn't infer
+                // its render mode, which means is an SSR only component and we don't need to persist its state at all.
+                continue;
+            }
+
+            var result = ExecuteCallback(registration.Callback, _logger);
             if (!result.IsCompletedSuccessfully)
             {
                 pendingCallbackTasks ??= new();

src/Components/Components/src/PersistComponentStateRegistration.cs

@@ -0,0 +1,13 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+internal readonly struct PersistComponentStateRegistration(
+    Func<Task> callback,
+    IComponentRenderMode? renderMode)
+{
+    public Func<Task> Callback { get; } = callback;
+
+    public IComponentRenderMode? RenderMode { get; } = renderMode;
+}

src/Components/Components/src/PersistentComponentState.cs

@@ -15,11 +15,11 @@ public class PersistentComponentState
     private IDictionary<string, byte[]>? _existingState;
     private readonly IDictionary<string, byte[]> _currentState;
 
-    private readonly List<Func<Task>> _registeredCallbacks;
+    private readonly List<PersistComponentStateRegistration> _registeredCallbacks;
 
     internal PersistentComponentState(
-        IDictionary<string, byte[]> currentState,
-        List<Func<Task>> pauseCallbacks)
+        IDictionary<string , byte[]> currentState,
+        List<PersistComponentStateRegistration> pauseCallbacks)
     {
         _currentState = currentState;
         _registeredCallbacks = pauseCallbacks;
@@ -43,12 +43,24 @@ internal void InitializeExistingState(IDictionary<string, byte[]> existingState)
     /// <param name="callback">The callback to invoke when the application is being paused.</param>
     /// <returns>A subscription that can be used to unregister the callback when disposed.</returns>
     public PersistingComponentStateSubscription RegisterOnPersisting(Func<Task> callback)
+        => RegisterOnPersisting(callback, null);
+
+    /// <summary>
+    /// Register a callback to persist the component state when the application is about to be paused.
+    /// Registered callbacks can use this opportunity to persist their state so that it can be retrieved when the application resumes.
+    /// </summary>
+    /// <param name="callback">The callback to invoke when the application is being paused.</param>
+    /// <param name="renderMode"></param>
+    /// <returns>A subscription that can be used to unregister the callback when disposed.</returns>
+    public PersistingComponentStateSubscription RegisterOnPersisting(Func<Task> callback, IComponentRenderMode? renderMode)
     {
         ArgumentNullException.ThrowIfNull(callback);
 
-        _registeredCallbacks.Add(callback);
+        var persistenceCallback = new PersistComponentStateRegistration(callback, renderMode);
+
+        _registeredCallbacks.Add(persistenceCallback);
 
-        return new PersistingComponentStateSubscription(_registeredCallbacks, callback);
+        return new PersistingComponentStateSubscription(_registeredCallbacks, persistenceCallback);
     }
 
     /// <summary>

src/Components/Components/src/PersistingComponentStateSubscription.cs

@@ -11,10 +11,10 @@ namespace Microsoft.AspNetCore.Components;
 /// </summary>
 public readonly struct PersistingComponentStateSubscription : IDisposable
 {
-    private readonly List<Func<Task>>? _callbacks;
-    private readonly Func<Task>? _callback;
+    private readonly List<PersistComponentStateRegistration>? _callbacks;
+    private readonly PersistComponentStateRegistration? _callback;
 
-    internal PersistingComponentStateSubscription(List<Func<Task>> callbacks, Func<Task> callback)
+    internal PersistingComponentStateSubscription(List<PersistComponentStateRegistration> callbacks, PersistComponentStateRegistration callback)
     {
         _callbacks = callbacks;
         _callback = callback;
@@ -23,9 +23,9 @@ internal PersistingComponentStateSubscription(List<Func<Task>> callbacks, Func<T
     /// <inheritdoc />
     public void Dispose()
     {
-        if (_callback != null)
+        if (_callback.HasValue)
         {
-            _callbacks?.Remove(_callback);
+            _callbacks?.Remove(_callback.Value);
         }
     }
 }

src/Components/Components/src/PublicAPI.Unshipped.txt

@@ -16,11 +16,12 @@ Microsoft.AspNetCore.Components.CascadingValueSource<TValue>.NotifyChangedAsync(
 Microsoft.AspNetCore.Components.CascadingValueSource<TValue>.NotifyChangedAsync(TValue newValue) -> System.Threading.Tasks.Task!
 Microsoft.AspNetCore.Components.ComponentBase.DispatchExceptionAsync(System.Exception! exception) -> System.Threading.Tasks.Task!
 Microsoft.AspNetCore.Components.IComponentRenderMode
-Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.PersistStateAsync(Microsoft.AspNetCore.Components.IPersistentComponentStateStore! store, Microsoft.AspNetCore.Components.Dispatcher! dispatcher) -> System.Threading.Tasks.Task!
 Microsoft.AspNetCore.Components.InjectAttribute.Key.get -> object?
 Microsoft.AspNetCore.Components.InjectAttribute.Key.init -> void
+Microsoft.AspNetCore.Components.IPersistentComponentStateStore.SupportsRenderMode(Microsoft.AspNetCore.Components.IComponentRenderMode! renderMode) -> bool
 Microsoft.AspNetCore.Components.ParameterView.ToDictionary() -> System.Collections.Generic.IReadOnlyDictionary<string!, object?>!
 *REMOVED*Microsoft.AspNetCore.Components.ParameterView.ToDictionary() -> System.Collections.Generic.IReadOnlyDictionary<string!, object!>!
+Microsoft.AspNetCore.Components.PersistentComponentState.RegisterOnPersisting(System.Func<System.Threading.Tasks.Task!>! callback, Microsoft.AspNetCore.Components.IComponentRenderMode? renderMode) -> Microsoft.AspNetCore.Components.PersistingComponentStateSubscription
 Microsoft.AspNetCore.Components.RenderHandle.DispatchExceptionAsync(System.Exception! exception) -> System.Threading.Tasks.Task!
 *REMOVED*Microsoft.AspNetCore.Components.NavigationManager.ToAbsoluteUri(string! relativeUri) -> System.Uri!
 Microsoft.AspNetCore.Components.NavigationManager.ToAbsoluteUri(string? relativeUri) -> System.Uri!
@@ -44,6 +45,7 @@ Microsoft.AspNetCore.Components.RenderTree.NamedEventChangeType
 Microsoft.AspNetCore.Components.RenderTree.NamedEventChangeType.Added = 0 -> Microsoft.AspNetCore.Components.RenderTree.NamedEventChangeType
 Microsoft.AspNetCore.Components.RenderTree.NamedEventChangeType.Removed = 1 -> Microsoft.AspNetCore.Components.RenderTree.NamedEventChangeType
 Microsoft.AspNetCore.Components.RenderTree.RenderBatch.NamedEventChanges.get -> Microsoft.AspNetCore.Components.RenderTree.ArrayRange<Microsoft.AspNetCore.Components.RenderTree.NamedEventChange>?
+Microsoft.AspNetCore.Components.RenderTree.Renderer.GetComponentState(Microsoft.AspNetCore.Components.IComponent! component) -> Microsoft.AspNetCore.Components.Rendering.ComponentState!
 Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrame.ComponentFrameFlags.get -> Microsoft.AspNetCore.Components.RenderTree.ComponentFrameFlags
 Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrameType.ComponentRenderMode = 9 -> Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrameType
 Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrameType.NamedEvent = 10 -> Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrameType
@@ -101,6 +103,7 @@ virtual Microsoft.AspNetCore.Components.Rendering.ComponentState.DisposeAsync()
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.AddPendingTask(Microsoft.AspNetCore.Components.Rendering.ComponentState? componentState, System.Threading.Tasks.Task! task) -> void
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.CreateComponentState(int componentId, Microsoft.AspNetCore.Components.IComponent! component, Microsoft.AspNetCore.Components.Rendering.ComponentState? parentComponentState) -> Microsoft.AspNetCore.Components.Rendering.ComponentState!
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.DispatchEventAsync(ulong eventHandlerId, Microsoft.AspNetCore.Components.RenderTree.EventFieldInfo? fieldInfo, System.EventArgs! eventArgs, bool waitForQuiescence) -> System.Threading.Tasks.Task!
+virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.GetComponentRenderMode(Microsoft.AspNetCore.Components.IComponent! component) -> Microsoft.AspNetCore.Components.IComponentRenderMode?
 virtual Microsoft.AspNetCore.Components.RenderTree.Renderer.ResolveComponentForRenderMode(System.Type! componentType, int? parentComponentId, Microsoft.AspNetCore.Components.IComponentActivator! componentActivator, Microsoft.AspNetCore.Components.IComponentRenderMode! renderMode) -> Microsoft.AspNetCore.Components.IComponent!
 ~Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrame.ComponentRenderMode.get -> Microsoft.AspNetCore.Components.IComponentRenderMode
 ~Microsoft.AspNetCore.Components.RenderTree.RenderTreeFrame.NamedEventAssignedName.get -> string

src/Components/Components/src/RenderTree/Renderer.cs

@@ -133,7 +133,20 @@ private static IComponentActivator GetComponentActivatorOrDefault(IServiceProvid
     protected ComponentState GetComponentState(int componentId)
         => GetRequiredComponentState(componentId);
 
-    internal ComponentState GetComponentState(IComponent component)
+    /// <summary>
+    /// Gets the <see cref="IComponentRenderMode"/> for a given component if available.
+    /// </summary>
+    /// <param name="component">The component type</param>
+    /// <returns></returns>
+    protected internal virtual IComponentRenderMode? GetComponentRenderMode(IComponent component)
+        => null;
+
+    /// <summary>
+    /// Resolves the component state for a given <see cref="IComponent"/> instance.
+    /// </summary>
+    /// <param name="component">The <see cref="IComponent"/> instance</param>
+    /// <returns></returns>
+    protected internal ComponentState GetComponentState(IComponent component)
         => _componentStateByComponent.GetValueOrDefault(component);
 
     private async void RenderRootComponentsOnHotReload()