[Blazor] Add IPersistentComponentStateSerializer<T> interface for custom serialization extensibility

src/Components/Components/src/IPersistentComponentStateSerializer.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.Buffers;
+
+namespace Microsoft.AspNetCore.Components;
+
+/// <summary>
+/// Provides custom serialization logic for persistent component state values.
+/// </summary>
+internal interface IPersistentComponentStateSerializer
+{
+    /// <summary>
+    /// Serializes the provided <paramref name="value"/> and writes it to the <paramref name="writer"/>.
+    /// </summary>
+    /// <param name="type">The type of the value to serialize.</param>
+    /// <param name="value">The value to serialize.</param>
+    /// <param name="writer">The buffer writer to write the serialized data to.</param>
+    /// <returns>A task that represents the asynchronous serialization operation.</returns>
+    Task PersistAsync(Type type, object value, IBufferWriter<byte> writer);
+
+    /// <summary>
+    /// Deserializes a value from the provided <paramref name="data"/>.
+    /// This method must be synchronous to avoid UI tearing during component state restoration.
+    /// </summary>
+    /// <param name="type">The type of the value to deserialize.</param>
+    /// <param name="data">The serialized data to deserialize.</param>
+    /// <returns>The deserialized value.</returns>
+    object Restore(Type type, ReadOnlySequence<byte> data);
+}
+
+/// <summary>
+/// Provides custom serialization logic for persistent component state values of type <typeparamref name="T"/>.
+/// </summary>
+/// <typeparam name="T">The type of the value to serialize.</typeparam>
+internal interface IPersistentComponentStateSerializer<T> : IPersistentComponentStateSerializer
+{
+    /// <summary>
+    /// Serializes the provided <paramref name="value"/> and writes it to the <paramref name="writer"/>.
+    /// </summary>
+    /// <param name="value">The value to serialize.</param>
+    /// <param name="writer">The buffer writer to write the serialized data to.</param>
+    /// <returns>A task that represents the asynchronous serialization operation.</returns>
+    Task PersistAsync(T value, IBufferWriter<byte> writer);
+
+    /// <summary>
+    /// Deserializes a value of type <typeparamref name="T"/> from the provided <paramref name="data"/>.
+    /// This method must be synchronous to avoid UI tearing during component state restoration.
+    /// </summary>
+    /// <param name="data">The serialized data to deserialize.</param>
+    /// <returns>The deserialized value.</returns>
+    T Restore(ReadOnlySequence<byte> data);
+
+    /// <summary>
+    /// Default implementation of the non-generic PersistAsync method.
+    /// </summary>
+    Task IPersistentComponentStateSerializer.PersistAsync(Type type, object value, IBufferWriter<byte> writer)
+        => PersistAsync((T)value, writer);
+
+    /// <summary>
+    /// Default implementation of the non-generic Restore method.
+    /// </summary>
+    object IPersistentComponentStateSerializer.Restore(Type type, ReadOnlySequence<byte> data)
+        => Restore(data)!;
+}

src/Components/Components/src/Microsoft.AspNetCore.Components.csproj

@@ -19,6 +19,7 @@
     <Compile Include="$(ComponentsSharedSourceRoot)src\HotReloadManager.cs" LinkBase="HotReload" />
     <Compile Include="$(ComponentsSharedSourceRoot)src\RootTypeCache.cs" LinkBase="Shared" />
     <Compile Include="$(SharedSourceRoot)LinkerFlags.cs" LinkBase="Shared" />
+    <Compile Include="$(SharedSourceRoot)PooledArrayBufferWriter.cs" LinkBase="Shared" />
     <Compile Include="$(SharedSourceRoot)QueryStringEnumerable.cs" LinkBase="Shared" />
     <Compile Include="$(SharedSourceRoot)Debugger\DictionaryItemDebugView.cs" LinkBase="Shared" />
     <Compile Include="$(SharedSourceRoot)Debugger\DictionaryDebugView.cs" LinkBase="Shared" />

src/Components/Components/src/PersistentComponentState.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Buffers;
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 using static Microsoft.AspNetCore.Internal.LinkerFlags;
@@ -110,6 +111,55 @@ internal void PersistAsJson(string key, object instance, [DynamicallyAccessedMem
         _currentState.Add(key, JsonSerializer.SerializeToUtf8Bytes(instance, type, JsonSerializerOptionsProvider.Options));
     }
 
+    /// <summary>
+    /// Persists the provided byte array under the given key.
+    /// </summary>
+    /// <param name="key">The key to use to persist the state.</param>
+    /// <param name="data">The byte array to persist.</param>
+    internal void PersistAsBytes(string key, byte[] data)
+    {
+        ArgumentNullException.ThrowIfNull(key);
+
+        if (!PersistingState)
+        {
+            throw new InvalidOperationException("Persisting state is only allowed during an OnPersisting callback.");
+        }
+
+        if (_currentState.ContainsKey(key))
+        {
+            throw new ArgumentException($"There is already a persisted object under the same key '{key}'");
+        }
+
+        _currentState.Add(key, data);
+    }
+
+    /// <summary>
+    /// Serializes <paramref name="instance"/> using the provided <paramref name="serializer"/> and persists it under the given <paramref name="key"/>.
+    /// </summary>
+    /// <typeparam name="TValue">The <paramref name="instance"/> type.</typeparam>
+    /// <param name="key">The key to use to persist the state.</param>
+    /// <param name="instance">The instance to persist.</param>
+    /// <param name="serializer">The custom serializer to use for serialization.</param>
+    internal async Task PersistAsync<TValue>(string key, TValue instance, IPersistentComponentStateSerializer<TValue> serializer)
+    {
+        ArgumentNullException.ThrowIfNull(key);
+        ArgumentNullException.ThrowIfNull(serializer);
+
+        if (!PersistingState)
+        {
+            throw new InvalidOperationException("Persisting state is only allowed during an OnPersisting callback.");
+        }
+
+        if (_currentState.ContainsKey(key))
+        {
+            throw new ArgumentException($"There is already a persisted object under the same key '{key}'");
+        }
+
+        using var writer = new PooledArrayBufferWriter<byte>();
+        await serializer.PersistAsync(instance, writer);
+        _currentState.Add(key, writer.WrittenMemory.ToArray());
+    }
+
     /// <summary>
     /// Tries to retrieve the persisted state as JSON with the given <paramref name="key"/> and deserializes it into an
     /// instance of type <typeparamref name="TValue"/>.
@@ -155,6 +205,47 @@ internal bool TryTakeFromJson(string key, [DynamicallyAccessedMembers(JsonSerial
         }
     }
 
+    /// <summary>
+    /// Tries to retrieve the persisted state with the given <paramref name="key"/> and deserializes it using the provided <paramref name="serializer"/> into an
+    /// instance of type <typeparamref name="TValue"/>.
+    /// When the key is present, the state is successfully returned via <paramref name="instance"/>
+    /// and removed from the <see cref="PersistentComponentState"/>.
+    /// </summary>
+    /// <param name="key">The key used to persist the instance.</param>
+    /// <param name="serializer">The custom serializer to use for deserialization.</param>
+    /// <param name="instance">The persisted instance.</param>
+    /// <returns><c>true</c> if the state was found; <c>false</c> otherwise.</returns>
+    internal bool TryTake<TValue>(string key, IPersistentComponentStateSerializer<TValue> serializer, [MaybeNullWhen(false)] out TValue instance)
+    {
+        ArgumentNullException.ThrowIfNull(key);
+        ArgumentNullException.ThrowIfNull(serializer);
+
+        if (TryTake(key, out var data))
+        {
+            var sequence = new ReadOnlySequence<byte>(data!);
+            instance = serializer.Restore(sequence);
+            return true;
+        }
+        else
+        {
+            instance = default;
+            return false;
+        }
+    }
+
+    /// <summary>
+    /// Tries to retrieve the persisted state as raw bytes with the given <paramref name="key"/>.
+    /// When the key is present, the raw bytes are successfully returned via <paramref name="data"/>
+    /// and removed from the <see cref="PersistentComponentState"/>.
+    /// </summary>
+    /// <param name="key">The key used to persist the data.</param>
+    /// <param name="data">The persisted raw bytes.</param>
+    /// <returns><c>true</c> if the state was found; <c>false</c> otherwise.</returns>
+    internal bool TryTakeBytes(string key, [MaybeNullWhen(false)] out byte[]? data)
+    {
+        return TryTake(key, out data);
+    }
+
     private bool TryTake(string key, out byte[]? value)
     {
         ArgumentNullException.ThrowIfNull(key);

src/Components/Components/src/PersistentStateValueProvider.cs

@@ -15,10 +15,11 @@
 
 namespace Microsoft.AspNetCore.Components.Infrastructure;
 
-internal sealed class PersistentStateValueProvider(PersistentComponentState state) : ICascadingValueSupplier
+internal sealed class PersistentStateValueProvider(PersistentComponentState state, IServiceProvider serviceProvider) : ICascadingValueSupplier
 {
     private static readonly ConcurrentDictionary<(string, string, string), byte[]> _keyCache = new();
     private static readonly ConcurrentDictionary<(Type, string), PropertyGetter> _propertyGetterCache = new();
+    private readonly ConcurrentDictionary<Type, IPersistentComponentStateSerializer?> _serializerCache = new();
 
     private readonly Dictionary<ComponentState, PersistingComponentStateSubscription> _subscriptions = [];
 
@@ -42,6 +43,20 @@ public bool CanSupplyValue(in CascadingParameterInfo parameterInfo)
         var componentState = (ComponentState)key!;
         var storageKey = ComputeKey(componentState, parameterInfo.PropertyName);
 
+        // Try to get a custom serializer for this type first
+        var customSerializer = ResolveSerializer(parameterInfo.PropertyType);
+
+        if (customSerializer != null)
+        {
+            if (state.TryTakeBytes(storageKey, out var data))
+            {
+                var sequence = new ReadOnlySequence<byte>(data!);
+                return customSerializer.Restore(parameterInfo.PropertyType, sequence);
+            }
+            return null;
+        }
+
+        // Fallback to JSON serialization
         return state.TryTakeFromJson(storageKey, parameterInfo.PropertyType, out var value) ? value : null;
     }
 
@@ -52,17 +67,30 @@ public void Subscribe(ComponentState subscriber, in CascadingParameterInfo param
     {
         var propertyName = parameterInfo.PropertyName;
         var propertyType = parameterInfo.PropertyType;
-        _subscriptions[subscriber] = state.RegisterOnPersisting(() =>
+
+        // Resolve serializer outside the lambda
+        var customSerializer = ResolveSerializer(propertyType);
+
+        _subscriptions[subscriber] = state.RegisterOnPersisting(async () =>
             {
                 var storageKey = ComputeKey(subscriber, propertyName);
                 var propertyGetter = ResolvePropertyGetter(subscriber.Component.GetType(), propertyName);
                 var property = propertyGetter.GetValue(subscriber.Component);
                 if (property == null)
                 {
-                    return Task.CompletedTask;
+                    return;
+                }
+
+                if (customSerializer != null)
+                {
+                    using var writer = new PooledArrayBufferWriter<byte>();
+                    await customSerializer.PersistAsync(propertyType, property, writer);
+                    state.PersistAsBytes(storageKey, writer.WrittenMemory.ToArray());
+                    return;
                 }
+
+                // Fallback to JSON serialization
                 state.PersistAsJson(storageKey, property, propertyType);
-                return Task.CompletedTask;
             }, subscriber.Renderer.GetComponentRenderMode(subscriber.Component));
     }
 
@@ -71,6 +99,18 @@ private static PropertyGetter ResolvePropertyGetter(Type type, string propertyNa
         return _propertyGetterCache.GetOrAdd((type, propertyName), PropertyGetterFactory);
     }
 
+    private IPersistentComponentStateSerializer? ResolveSerializer(Type type)
+    {
+        return _serializerCache.GetOrAdd(type, SerializerFactory);
+    }
+
+    private IPersistentComponentStateSerializer? SerializerFactory(Type type)
+    {
+        var serializerType = typeof(IPersistentComponentStateSerializer<>).MakeGenericType(type);
+        var serializer = serviceProvider.GetService(serializerType);
+        return serializer as IPersistentComponentStateSerializer;
+    }
+
     [UnconditionalSuppressMessage(
     "Trimming",
     "IL2077:Target parameter argument does not satisfy 'DynamicallyAccessedMembersAttribute' in call to target method. The source field does not have matching annotations.",

src/Components/Components/test/IPersistentComponentStateSerializerTests.cs

@@ -0,0 +1,75 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Buffers;
+using System.Text;
+using System.Text.Json;
+using Microsoft.AspNetCore.Components.Infrastructure;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.AspNetCore.Components;
+
+public class IPersistentComponentStateSerializerTests
+{
+    [Fact]
+    public async Task PersistAsync_CanUseCustomSerializer()
+    {
+        // Arrange
+        var currentState = new Dictionary<string, byte[]>();
+        var state = new PersistentComponentState(currentState, []);
+        var customSerializer = new TestStringSerializer();
+        var testValue = "Hello, World!";
+
+        state.PersistingState = true;
+
+        // Act
+        await state.PersistAsync("test-key", testValue, customSerializer);
+
+        // Assert
+        state.PersistingState = false;
+
+        // Simulate the state transfer that happens between persist and restore phases
+        var newState = new PersistentComponentState(new Dictionary<string, byte[]>(), []);
+        newState.InitializeExistingState(currentState);
+
+        Assert.True(newState.TryTake("test-key", customSerializer, out var retrievedValue));
+        Assert.Equal(testValue, retrievedValue);
+    }
+
+    [Fact]
+    public void TryTake_CanUseCustomSerializer()
+    {
+        // Arrange
+        var customData = "Custom Data";
+        var customBytes = Encoding.UTF8.GetBytes(customData);
+        var existingState = new Dictionary<string, byte[]> { { "test-key", customBytes } };
+
+        var state = new PersistentComponentState(new Dictionary<string, byte[]>(), []);
+        state.InitializeExistingState(existingState);
+
+        var customSerializer = new TestStringSerializer();
+
+        // Act
+        var success = state.TryTake("test-key", customSerializer, out var retrievedValue);
+
+        // Assert
+        Assert.True(success);
+        Assert.Equal(customData, retrievedValue);
+    }
+
+    private class TestStringSerializer : IPersistentComponentStateSerializer<string>
+    {
+        public Task PersistAsync(string value, IBufferWriter<byte> writer)
+        {
+            var bytes = Encoding.UTF8.GetBytes(value);
+            writer.Write(bytes);
+            return Task.CompletedTask;
+        }
+
+        public string Restore(ReadOnlySequence<byte> data)
+        {
+            var bytes = data.ToArray();
+            return Encoding.UTF8.GetString(bytes);
+        }
+    }
+}