Persist 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="PersistedStateSerializationMode"/>.
+    /// </summary>
+    /// <param name="serializationMode">The <see cref="PersistedStateSerializationMode"/> in question.</param>
+    /// <returns><c>true</c> if the serialization mode is supported by the store, otherwise <c>false</c>.</returns>
+    bool SupportsSerializationMode(PersistedStateSerializationMode serializationMode) => true;
 }

src/Components/Components/src/ISerializationModeHandler.cs

@@ -0,0 +1,17 @@
+// 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;
+
+/// <summary>
+/// A service that can infer <see cref="IComponent"/>'s <see cref="PersistedStateSerializationMode"/>.
+/// </summary>
+public interface ISerializationModeHandler
+{
+    /// <summary>
+    /// Infers <see cref="IComponent"/>'s <see cref="PersistedStateSerializationMode"/>.
+    /// </summary>
+    /// <param name="callbackTarget">The callback target</param>
+    /// <returns>The <see cref="PersistedStateSerializationMode"/> for the component.</returns>
+    public PersistedStateSerializationMode GetCallbackTargetSerializationMode(object? callbackTarget);
+}

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

@@ -11,25 +11,33 @@ namespace Microsoft.AspNetCore.Components.Infrastructure;
 /// </summary>
 public class ComponentStatePersistenceManager
 {
-    private bool _stateIsPersisted;
-    private readonly List<Func<Task>> _pauseCallbacks = new();
-    private readonly Dictionary<string, byte[]> _currentState = new(StringComparer.Ordinal);
+    private readonly List<PersistenceCallback> _registeredCallbacks = new();
     private readonly ILogger<ComponentStatePersistenceManager> _logger;
 
     /// <summary>
     /// Initializes a new instance of <see cref="ComponentStatePersistenceManager"/>.
     /// </summary>
-    public ComponentStatePersistenceManager(ILogger<ComponentStatePersistenceManager> logger)
+    public ComponentStatePersistenceManager(ILogger<ComponentStatePersistenceManager> logger, ISerializationModeHandler serializationModeHandler)
     {
-        State = new PersistentComponentState(_currentState, _pauseCallbacks);
         _logger = logger;
+
+        State = new(_registeredCallbacks, serializationModeHandler);
     }
 
     /// <summary>
     /// Gets the <see cref="ComponentStatePersistenceManager"/> associated with the <see cref="ComponentStatePersistenceManager"/>.
     /// </summary>
     public PersistentComponentState State { get; }
 
+    /// <summary>
+    /// 
+    /// </summary>
+    /// <param name="serializationModeHandler"></param>
+    public void SetSerializationModeHandler(ISerializationModeHandler serializationModeHandler)
+    {
+        State.SetSerializationModeHandler(serializationModeHandler);
+    }
+
     /// <summary>
     /// Restores the component application state from the given <see cref="IPersistentComponentStateStore"/>.
     /// </summary>
@@ -44,47 +52,48 @@ public async Task RestoreStateAsync(IPersistentComponentStateStore store)
     /// <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="store">The <see cref="IPersistentComponentStateStore"/> to persist the application state into.</param>
     /// <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"/>.
+    /// Persists the component application state into the given <see cref="IPersistentComponentStateStore"/>
+    /// so that it could be restored on Server.
     /// </summary>
-    /// <param name="store">The <see cref="IPersistentComponentStateStore"/> to restore the application state from.</param>
+    /// <param name="store">The <see cref="IPersistentComponentStateStore"/> to persist the application state into.</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);
 
         async Task PauseAndPersistState()
         {
-            State.PersistingState = true;
-            await PauseAsync();
-            State.PersistingState = false;
+            var currentState = new Dictionary<string, byte[]>();
 
-            await store.PersistStateAsync(_currentState);
+            State.PersistenceContext = new(currentState);
+            await PauseAsync(store);
+            State.PersistenceContext = default;
+
+            await store.PersistStateAsync(currentState);
         }
     }
 
-    internal Task PauseAsync()
+    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 callback = _registeredCallbacks[i];
+            if (!store.SupportsSerializationMode(callback.SerializationMode))
+            {
+                continue;
+            }
+
+            var result = ExecuteCallback(callback.Callback, _logger);
             if (!result.IsCompletedSuccessfully)
             {
                 pendingCallbackTasks ??= new();

src/Components/Endpoints/src/DependencyInjection/PersistedStateSerializationMode.cs → src/Components/Components/src/PersistedStateSerializationMode.cs

@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.Components;
 public enum PersistedStateSerializationMode
 {
     /// <summary>
-    /// Indicates that the serialization mode should be inferred from the current request context.
+    /// Indicates that the serialization mode should be inferred.
     /// </summary>
     Infer = 1,
 
@@ -22,4 +22,9 @@ public enum PersistedStateSerializationMode
     /// Indicates that the state should be persisted so that execution may resume on WebAssembly.
     /// </summary>
     WebAssembly = 3,
+
+    /// <summary>
+    /// Indicates that the state should be persisted so that execution may resume on both Server and WebAssembly.
+    /// </summary>
+    ServerAndWebAssembly = 4,
 }

src/Components/Components/src/PersistenceCallback.cs

@@ -0,0 +1,6 @@
+// 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 record struct PersistenceCallback(Func<Task> Callback, PersistedStateSerializationMode SerializationMode);

src/Components/Components/src/PersistenceContext.cs

@@ -0,0 +1,6 @@
+// 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 record struct PersistenceContext(IDictionary<string, byte[]> State);

src/Components/Components/src/PersistentComponentState.cs

@@ -13,29 +13,29 @@ namespace Microsoft.AspNetCore.Components;
 public class PersistentComponentState
 {
     private IDictionary<string, byte[]>? _existingState;
-    private readonly IDictionary<string, byte[]> _currentState;
 
-    private readonly List<Func<Task>> _registeredCallbacks;
+    private readonly List<PersistenceCallback> _registeredCallbacks;
+    private ISerializationModeHandler _serializationModeHandler;
 
-    internal PersistentComponentState(
-        IDictionary<string, byte[]> currentState,
-        List<Func<Task>> pauseCallbacks)
+    internal PersistenceContext? PersistenceContext { get; set; }
+
+    internal PersistentComponentState(List<PersistenceCallback> registeredCallbacks, ISerializationModeHandler serializationModeHandler)
     {
-        _currentState = currentState;
-        _registeredCallbacks = pauseCallbacks;
+        _registeredCallbacks = registeredCallbacks;
+        _serializationModeHandler = serializationModeHandler;
     }
 
-    internal bool PersistingState { get; set; }
-
     internal void InitializeExistingState(IDictionary<string, byte[]> existingState)
     {
-        if (_existingState != null)
-        {
-            throw new InvalidOperationException("PersistentComponentState already initialized.");
-        }
+        // The existing state is either Server or WebAssembly
         _existingState = existingState ?? throw new ArgumentNullException(nameof(existingState));
     }
 
+    internal void SetSerializationModeHandler(ISerializationModeHandler serializationModeHandler)
+    {
+        _serializationModeHandler = serializationModeHandler;
+    }
+
     /// <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.
@@ -46,9 +46,26 @@ public PersistingComponentStateSubscription RegisterOnPersisting(Func<Task> call
     {
         ArgumentNullException.ThrowIfNull(callback);
 
-        _registeredCallbacks.Add(callback);
+        var serializationMode = _serializationModeHandler.GetCallbackTargetSerializationMode(callback.Target);
 
-        return new PersistingComponentStateSubscription(_registeredCallbacks, callback);
+        return RegisterOnPersisting(callback, serializationMode);
+    }
+
+    /// <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="serializationMode">The <see cref="PersistedStateSerializationMode"/> to register the callback.</param>
+    /// <returns>A subscription that can be used to unregister the callback when disposed.</returns>
+    public PersistingComponentStateSubscription RegisterOnPersisting(Func<Task> callback, PersistedStateSerializationMode serializationMode)
+    {
+        ArgumentNullException.ThrowIfNull(callback);
+
+        var persistenceCallback = new PersistenceCallback(callback, serializationMode);
+        _registeredCallbacks.Add(persistenceCallback);
+
+        return new PersistingComponentStateSubscription(_registeredCallbacks, persistenceCallback);
     }
 
     /// <summary>
@@ -60,19 +77,17 @@ public PersistingComponentStateSubscription RegisterOnPersisting(Func<Task> call
     [RequiresUnreferencedCode("JSON serialization and deserialization might require types that cannot be statically analyzed.")]
     public void PersistAsJson<[DynamicallyAccessedMembers(JsonSerialized)] TValue>(string key, TValue instance)
     {
-        ArgumentNullException.ThrowIfNull(key);
-
-        if (!PersistingState)
+        if (PersistenceContext is not { State: var currentState })
         {
             throw new InvalidOperationException("Persisting state is only allowed during an OnPersisting callback.");
         }
 
-        if (_currentState.ContainsKey(key))
+        if (currentState.ContainsKey(key))
         {
             throw new ArgumentException($"There is already a persisted object under the same key '{key}'");
         }
 
-        _currentState.Add(key, JsonSerializer.SerializeToUtf8Bytes(instance, JsonSerializerOptionsProvider.Options));
+        currentState.Add(key, JsonSerializer.SerializeToUtf8Bytes(instance, JsonSerializerOptionsProvider.Options));
     }
 
     /// <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<PersistenceCallback>? _callbacks;
+    private readonly PersistenceCallback? _callback;
 
-    internal PersistingComponentStateSubscription(List<Func<Task>> callbacks, Func<Task> callback)
+    internal PersistingComponentStateSubscription(List<PersistenceCallback> callbacks, PersistenceCallback 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,9 +16,21 @@ 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
+*REMOVED*Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.ComponentStatePersistenceManager(Microsoft.Extensions.Logging.ILogger<Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager!>! logger) -> void
+Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.ComponentStatePersistenceManager(Microsoft.Extensions.Logging.ILogger<Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager!>! logger, Microsoft.AspNetCore.Components.ISerializationModeHandler! serializationModeHandler) -> void
 Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.PersistStateAsync(Microsoft.AspNetCore.Components.IPersistentComponentStateStore! store, Microsoft.AspNetCore.Components.Dispatcher! dispatcher) -> System.Threading.Tasks.Task!
+Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.SetSerializationModeHandler(Microsoft.AspNetCore.Components.ISerializationModeHandler! serializationModeHandler) -> void
+Microsoft.AspNetCore.Components.IPersistentComponentStateStore.SupportsSerializationMode(Microsoft.AspNetCore.Components.PersistedStateSerializationMode serializationMode) -> bool
+Microsoft.AspNetCore.Components.ISerializationModeHandler
+Microsoft.AspNetCore.Components.ISerializationModeHandler.GetCallbackTargetSerializationMode(object? callbackTarget) -> Microsoft.AspNetCore.Components.PersistedStateSerializationMode
 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.PersistedStateSerializationMode
+Microsoft.AspNetCore.Components.PersistedStateSerializationMode.Infer = 1 -> Microsoft.AspNetCore.Components.PersistedStateSerializationMode
+Microsoft.AspNetCore.Components.PersistedStateSerializationMode.Server = 2 -> Microsoft.AspNetCore.Components.PersistedStateSerializationMode
+Microsoft.AspNetCore.Components.PersistedStateSerializationMode.ServerAndWebAssembly = 4 -> Microsoft.AspNetCore.Components.PersistedStateSerializationMode
+Microsoft.AspNetCore.Components.PersistedStateSerializationMode.WebAssembly = 3 -> Microsoft.AspNetCore.Components.PersistedStateSerializationMode
+Microsoft.AspNetCore.Components.PersistentComponentState.RegisterOnPersisting(System.Func<System.Threading.Tasks.Task!>! callback, Microsoft.AspNetCore.Components.PersistedStateSerializationMode serializationMode) -> 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!
@@ -42,6 +54,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

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

@@ -133,7 +133,12 @@ private static IComponentActivator GetComponentActivatorOrDefault(IServiceProvid
     protected ComponentState GetComponentState(int componentId)
         => GetRequiredComponentState(componentId);
 
-    internal ComponentState GetComponentState(IComponent component)
+    /// <summary>
+    /// Gets the <see cref="ComponentState"/> associated with the specified component.
+    /// </summary>
+    /// <param name="component">The component</param>
+    /// <returns>The corresponding <see cref="ComponentState"/>.</returns>
+    protected internal ComponentState GetComponentState(IComponent component)
         => _componentStateByComponent.GetValueOrDefault(component);
 
     private async void RenderRootComponentsOnHotReload()