[Blazor] Prerendered state

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()