[Blazor] Adds support for persisting and restoring disconnected circuits from storage

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

@@ -23,6 +23,8 @@
     <Compile Include="$(SharedSourceRoot)Components\ServerComponentSerializationSettings.cs" LinkBase="DependencyInjection" />
     <Compile Include="$(SharedSourceRoot)Components\ServerComponent.cs" LinkBase="DependencyInjection" />
     <Compile Include="$(SharedSourceRoot)Components\ResourceCollectionResolver.cs" LinkBase="Assets" />
+    <Compile Include="$(SharedSourceRoot)Components\ServerComponentSerializer.cs" LinkBase="DependencyInjection" />
+    <Compile Include="$(SharedSourceRoot)Components\ServerComponentInvocationSequence.cs" LinkBase="DependencyInjection" />
     <Compile Include="$(RepoRoot)src\Shared\Components\ComponentParameter.cs" LinkBase="DependencyInjection" />
     <Compile Include="$(RepoRoot)src\Shared\Components\PrerenderComponentApplicationStore.cs" LinkBase="DependencyInjection" />
     <Compile Include="$(RepoRoot)src\Shared\Components\ProtectedPrerenderComponentApplicationStore.cs" LinkBase="DependencyInjection" />

src/Components/Server/src/CircuitOptions.cs

@@ -44,6 +44,22 @@ public sealed class CircuitOptions
     /// </value>
     public TimeSpan DisconnectedCircuitRetentionPeriod { get; set; } = TimeSpan.FromMinutes(3);
 
+    /// <summary>
+    /// Gets or sets a value that determines the maximum number of persisted circuits state that
+    /// are retained in memory by the server when no distributed cache is configured.
+    /// </summary>
+    /// <remarks>
+    /// When using a distributed cache like <see cref="Extensions.Caching.Hybrid.HybridCache"/> this value is ignored
+    /// and the configuration from <see cref="Extensions.DependencyInjection.MemoryCacheServiceCollectionExtensions.AddMemoryCache(Extensions.DependencyInjection.IServiceCollection)"/>
+    /// is used instead.
+    /// </remarks>
+    public int PersistedCircuitInMemoryMaxRetained { get; set; } = 1000;
+
+    /// <summary>
+    /// Gets or sets the duration for which a persisted circuit is retained in memory.
+    /// </summary>
+    public TimeSpan PersistedCircuitInMemoryRetentionPeriod { get; set; } = TimeSpan.FromHours(2);
+
     /// <summary>
     /// Gets or sets a value that determines whether or not to send detailed exception messages to JavaScript when an unhandled exception
     /// happens on the circuit or when a .NET method invocation through JS interop results in an exception.

src/Components/Server/src/Circuits/CircuitHost.cs

@@ -32,6 +32,7 @@ internal partial class CircuitHost : IAsyncDisposable
     private bool _isFirstUpdate = true;
     private bool _disposed;
     private long _startTime;
+    private PersistedCircuitState _persistedCircuitState;
 
     // This event is fired when there's an unrecoverable exception coming from the circuit, and
     // it need so be torn down. The registry listens to this even so that the circuit can
@@ -106,6 +107,8 @@ public CircuitHost(
 
     public IServiceProvider Services { get; }
 
+    internal bool HasPendingPersistedCircuitState => _persistedCircuitState != null;
+
     // InitializeAsync is used in a fire-and-forget context, so it's responsible for its own
     // error handling.
     public Task InitializeAsync(ProtectedPrerenderComponentApplicationStore store, ActivityContext httpContext, CancellationToken cancellationToken)
@@ -873,6 +876,23 @@ await HandleInboundActivityAsync(() =>
         }
     }
 
+    internal void AttachPersistedState(PersistedCircuitState persistedCircuitState)
+    {
+        if (_persistedCircuitState != null)
+        {
+            throw new InvalidOperationException("Persisted state has already been attached to this circuit.");
+        }
+
+        _persistedCircuitState = persistedCircuitState;
+    }
+
+    internal PersistedCircuitState TakePersistedCircuitState()
+    {
+        var result = _persistedCircuitState;
+        _persistedCircuitState = null;
+        return result;
+    }
+
     private static partial class Log
     {
         // 100s used for lifecycle stuff

src/Components/Server/src/Circuits/CircuitPersistenceManager.cs

@@ -0,0 +1,172 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Microsoft.AspNetCore.Components.Endpoints;
+using Microsoft.AspNetCore.Components.Infrastructure;
+using Microsoft.AspNetCore.Components.Web;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
+
+namespace Microsoft.AspNetCore.Components.Server.Circuits;
+
+internal partial class CircuitPersistenceManager(
+    IOptions<CircuitOptions> circuitOptions,
+    ServerComponentSerializer serverComponentSerializer,
+    ICircuitPersistenceProvider circuitPersistenceProvider)
+{
+    public async Task PauseCircuitAsync(CircuitHost circuit, CancellationToken cancellation = default)
+    {
+        var renderer = circuit.Renderer;
+        var persistenceManager = circuit.Services.GetRequiredService<ComponentStatePersistenceManager>();
+        var collector = new CircuitPersistenceManagerCollector(circuitOptions, serverComponentSerializer, circuit.Renderer);
+        using var subscription = persistenceManager.State.RegisterOnPersisting(
+            collector.PersistRootComponents,
+            RenderMode.InteractiveServer);
+
+        await persistenceManager.PersistStateAsync(collector, renderer);
+
+        await circuitPersistenceProvider.PersistCircuitAsync(
+            circuit.CircuitId,
+            collector.PersistedCircuitState,
+            cancellation);
+    }
+
+    public async Task<PersistedCircuitState> ResumeCircuitAsync(CircuitId circuitId, CancellationToken cancellation = default)
+    {
+        return await circuitPersistenceProvider.RestoreCircuitAsync(circuitId, cancellation);
+    }
+
+    // We are going to construct a RootComponentOperationBatch but we are going to replace the descriptors from the client with the
+    // descriptors that we have persisted when pausing the circuit.
+    // The way pausing and resuming works is that when the client starts the resume process, it 'simulates' that an SSR has happened and
+    // queues an 'Add' operation for each server-side component that is on the document.
+    // That ends up calling UpdateRootComponents with the old descriptors and no application state.
+    // On the server side, we replace the descriptors with the ones that we have persisted. We can't use the original descriptors because
+    // those have a lifetime of ~ 5 minutes, after which we are not able to unprotect them anymore.
+    internal static RootComponentOperationBatch ToRootComponentOperationBatch(
+        IServerComponentDeserializer serverComponentDeserializer,
+        byte[] rootComponents,
+        string serializedComponentOperations)
+    {
+        // Deserialize the existing batch the client has sent but ignore the markers
+        if (!serverComponentDeserializer.TryDeserializeRootComponentOperations(
+            serializedComponentOperations,
+            out var batch,
+            deserializeDescriptors: false))
+        {
+            return null;
+        }
+
+        var persistedMarkers = TryDeserializeMarkers(rootComponents);
+
+        if (persistedMarkers == null)
+        {
+            return null;
+        }
+
+        if (batch.Operations.Length != persistedMarkers.Count)
+        {
+            return null;
+        }
+
+        // Ensure that all operations in the batch are `Add` operations.
+        for (var i = 0; i < batch.Operations.Length; i++)
+        {
+            var operation = batch.Operations[i];
+            if (operation.Type != RootComponentOperationType.Add)
+            {
+                return null;
+            }
+
+            // Retrieve the marker from the persisted root components, replace it and deserialize the descriptor
+            if (!persistedMarkers.TryGetValue(operation.SsrComponentId, out var marker))
+            {
+                return null;
+            }
+            operation.Marker = marker;
+
+            if (!serverComponentDeserializer.TryDeserializeWebRootComponentDescriptor(operation.Marker.Value, out var descriptor))
+            {
+                return null;
+            }
+
+            operation.Descriptor = descriptor;
+        }
+
+        return batch;
+
+        static Dictionary<int, ComponentMarker> TryDeserializeMarkers(byte[] rootComponents)
+        {
+            if (rootComponents == null || rootComponents.Length == 0)
+            {
+                return null;
+            }
+
+            try
+            {
+                return JsonSerializer.Deserialize<Dictionary<int, ComponentMarker>>(
+                    rootComponents,
+                    JsonSerializerOptionsProvider.Options);
+            }
+            catch
+            {
+                return null;
+            }
+        }
+    }
+
+    private class CircuitPersistenceManagerCollector(
+        IOptions<CircuitOptions> circuitOptions,
+        ServerComponentSerializer serverComponentSerializer,
+        RemoteRenderer renderer)
+        : IPersistentComponentStateStore
+    {
+        internal PersistedCircuitState PersistedCircuitState { get; private set; }
+
+        public Task PersistRootComponents()
+        {
+            var persistedComponents = new Dictionary<int, ComponentMarker>();
+            var components = renderer.GetOrCreateWebRootComponentManager().GetRootComponents();
+            var invocation = new ServerComponentInvocationSequence();
+            foreach (var (id, componentKey, (componentType, parameters)) in components)
+            {
+                var distributedRetention = circuitOptions.Value.PersistedCircuitInMemoryRetentionPeriod;
+                var localRetention = circuitOptions.Value.PersistedCircuitInMemoryRetentionPeriod;
+                var maxRetention = distributedRetention > localRetention ? distributedRetention : localRetention;
+
+                var marker = ComponentMarker.Create(ComponentMarker.ServerMarkerType, prerendered: false, componentKey);
+                serverComponentSerializer.SerializeInvocation(ref marker, invocation, componentType, parameters, maxRetention);
+                persistedComponents.Add(id, marker);
+            }
+
+            PersistedCircuitState = new PersistedCircuitState
+            {
+                RootComponents = JsonSerializer.SerializeToUtf8Bytes(
+                    persistedComponents,
+                    CircuitPersistenceManagerSerializerContext.Default.DictionaryInt32ComponentMarker)
+            };
+
+            return Task.CompletedTask;
+        }
+
+        // This store only support serializing the state
+        Task<IDictionary<string, byte[]>> IPersistentComponentStateStore.GetPersistedStateAsync() => throw new NotImplementedException();
+
+        // During the persisting phase the state is captured into a Dictionary<string, byte[]>, our implementation registers
+        // a callback so that it can run at the same time as the other components' state is persisted.
+        // We then are called to save the persisted state, at which point, we extract the component records
+        // and store them separately from the other state.
+        Task IPersistentComponentStateStore.PersistStateAsync(IReadOnlyDictionary<string, byte[]> state)
+        {
+            PersistedCircuitState.ApplicationState = state;
+            return Task.CompletedTask;
+        }
+    }
+
+    [JsonSerializable(typeof(Dictionary<int, ComponentMarker>))]
+    internal partial class CircuitPersistenceManagerSerializerContext : JsonSerializerContext
+    {
+    }
+}

src/Components/Server/src/Circuits/CircuitRegistry.cs

@@ -41,16 +41,19 @@ internal partial class CircuitRegistry
     private readonly CircuitOptions _options;
     private readonly ILogger _logger;
     private readonly CircuitIdFactory _circuitIdFactory;
+    private readonly CircuitPersistenceManager _circuitPersistenceManager;
     private readonly PostEvictionCallbackRegistration _postEvictionCallback;
 
     public CircuitRegistry(
         IOptions<CircuitOptions> options,
         ILogger<CircuitRegistry> logger,
-        CircuitIdFactory CircuitHostFactory)
+        CircuitIdFactory CircuitHostFactory,
+        CircuitPersistenceManager circuitPersistenceManager)
     {
         _options = options.Value;
         _logger = logger;
         _circuitIdFactory = CircuitHostFactory;
+        _circuitPersistenceManager = circuitPersistenceManager;
         ConnectedCircuits = new ConcurrentDictionary<CircuitId, CircuitHost>();
 
         DisconnectedCircuits = new MemoryCache(new MemoryCacheOptions
@@ -265,7 +268,7 @@ protected virtual void OnEntryEvicted(object key, object value, EvictionReason r
                 // Kick off the dispose in the background.
                 var disconnectedEntry = (DisconnectedCircuitEntry)value;
                 Log.CircuitEvicted(_logger, disconnectedEntry.CircuitHost.CircuitId, reason);
-                _ = DisposeCircuitEntry(disconnectedEntry);
+                _ = PauseAndDisposeCircuitEntry(disconnectedEntry);
                 break;
 
             case EvictionReason.Removed:
@@ -278,12 +281,23 @@ protected virtual void OnEntryEvicted(object key, object value, EvictionReason r
         }
     }
 
-    private async Task DisposeCircuitEntry(DisconnectedCircuitEntry entry)
+    private async Task PauseAndDisposeCircuitEntry(DisconnectedCircuitEntry entry)
     {
         DisposeTokenSource(entry);
 
         try
         {
+            if (!entry.CircuitHost.HasPendingPersistedCircuitState)
+            {
+                // Only pause and persist the circuit state if it has been active at some point,
+                // meaning that the client called UpdateRootComponents on it.
+                await _circuitPersistenceManager.PauseCircuitAsync(entry.CircuitHost);
+            }
+            else
+            {
+                Log.PersistedCircuitStateDiscarded(_logger, entry.CircuitHost.CircuitId);
+            }
+
             entry.CircuitHost.UnhandledException -= CircuitHost_UnhandledException;
             await entry.CircuitHost.DisposeAsync();
         }
@@ -413,5 +427,8 @@ public static void ExceptionDisposingTokenSource(ILogger logger, Exception excep
 
         [LoggerMessage(115, LogLevel.Debug, "Reconnect to circuit with id {CircuitId} succeeded.", EventName = "ReconnectionSucceeded")]
         public static partial void ReconnectionSucceeded(ILogger logger, CircuitId circuitId);
+
+        [LoggerMessage(116, LogLevel.Debug, "Circuit {CircuitId} was not resumed. Persisted circuit state for {CircuitId} discarded.", EventName = "PersistedCircuitStateDiscarded")]
+        public static partial void PersistedCircuitStateDiscarded(ILogger logger, CircuitId circuitId);
     }
 }