Address feedback

Author: javiercn

src/Components/Server/src/Circuits/CircuitHost.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.Collections.ObjectModel;
 using System.Diagnostics;
 using System.Globalization;
 using System.Linq;
@@ -161,7 +162,7 @@ public Task InitializeAsync(ProtectedPrerenderComponentApplicationStore store, A
                 // web service, preventing UI updates.
                 if (Descriptors.Count > 0)
                 {
-                    store.ExistingState.Clear();
+                    store.ExistingState = ReadOnlyDictionary<string, byte[]>.Empty;
                 }
 
                 // This variable is used to track that this is the first time we are updating components.
@@ -821,7 +822,7 @@ internal Task UpdateRootComponents(
                     // At this point all components have successfully produced an initial render and we can clear the contents of the component
                     // application state store. This ensures the memory that was not used during the initial render of these components gets
                     // reclaimed since no-one else is holding on to it any longer.
-                    store.ExistingState.Clear();
+                    store.ExistingState = ReadOnlyDictionary<string, byte[]>.Empty;
                 }
             }
         });

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

@@ -16,45 +16,23 @@ internal partial class CircuitPersistenceManager(
     ServerComponentSerializer serverComponentSerializer,
     ICircuitPersistenceProvider circuitPersistenceProvider)
 {
-    private const string CircuitPersistenceManagerKey = $"Microsoft.AspNetCore.Components.Server.Circuits.{nameof(CircuitPersistenceManager)}";
-
     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(
-            () => PersistRootComponents(renderer, persistenceManager.State),
+            collector.PersistRootComponents,
             RenderMode.InteractiveServer);
-        var store = new CircuitPersistenceManagerStore();
-        await persistenceManager.PersistStateAsync(store, renderer);
+
+        await persistenceManager.PersistStateAsync(collector, renderer);
 
         await circuitPersistenceProvider.PersistCircuitAsync(
             circuit.CircuitId,
-            store.PersistedCircuitState,
+            collector.PersistedCircuitState,
             cancellation);
     }
 
-    private Task PersistRootComponents(RemoteRenderer renderer, PersistentComponentState state)
-    {
-        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, false, componentKey);
-            serverComponentSerializer.SerializeInvocation(ref marker, invocation, componentType, parameters, maxRetention);
-            persistedComponents.Add(id, marker);
-        }
-
-        state.PersistAsJson(CircuitPersistenceManagerKey, persistedComponents);
-
-        return Task.CompletedTask;
-    }
-
     public async Task<PersistedCircuitState> ResumeCircuitAsync(CircuitId circuitId, CancellationToken cancellation = default)
     {
         return await circuitPersistenceProvider.RestoreCircuitAsync(circuitId, cancellation);
@@ -63,7 +41,7 @@ public async Task<PersistedCircuitState> ResumeCircuitAsync(CircuitId circuitId,
     // 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 and 'Add' operation for each server-side component that is on the document.
+    // 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.
@@ -139,10 +117,40 @@ static Dictionary<int, ComponentMarker> TryDeserializeMarkers(byte[] rootCompone
         }
     }
 
-    private class CircuitPersistenceManagerStore : IPersistentComponentStateStore
+    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();
 
@@ -152,26 +160,7 @@ private class CircuitPersistenceManagerStore : IPersistentComponentStateStore
         // and store them separately from the other state.
         Task IPersistentComponentStateStore.PersistStateAsync(IReadOnlyDictionary<string, byte[]> state)
         {
-            var dictionary = new Dictionary<string, byte[]>(state.Count - 1);
-            byte[] rootComponentMarkers = null;
-            foreach (var (key, value) in state)
-            {
-                if (key == CircuitPersistenceManagerKey)
-                {
-                    rootComponentMarkers = value;
-                }
-                else
-                {
-                    dictionary[key] = value;
-                }
-            }
-
-            PersistedCircuitState = new PersistedCircuitState
-            {
-                ApplicationState = dictionary,
-                RootComponents = rootComponentMarkers
-            };
-
+            PersistedCircuitState.ApplicationState = state;
             return Task.CompletedTask;
         }
     }

src/Components/Server/src/Circuits/DefaultPersistedCircuitCache.cs renamed to src/Components/Server/src/Circuits/DefaultInMemoryCircuitPersistenceProvider.cs

@@ -16,7 +16,7 @@ internal sealed partial class DefaultInMemoryCircuitPersistenceProvider : ICircu
     private readonly Lock _lock = new();
     private readonly CircuitOptions _options;
     private readonly MemoryCache _persistedCircuits;
-    private readonly Task<PersistedCircuitState> _noMatch = Task.FromResult<PersistedCircuitState>(null);
+    private static readonly Task<PersistedCircuitState> _noMatch = Task.FromResult<PersistedCircuitState>(null);
     private readonly ILogger<ICircuitPersistenceProvider> _logger;
 
     public PostEvictionCallbackRegistration PostEvictionCallback { get; internal set; }
@@ -66,7 +66,8 @@ private void PersistCore(CircuitId circuitId, PersistedCircuitState persistedCir
         var persistedCircuitEntry = new PersistedCircuitEntry
         {
             State = persistedCircuitState,
-            TokenSource = cancellationTokenSource
+            TokenSource = cancellationTokenSource,
+            CircuitId = circuitId
         };
 
         _persistedCircuits.Set(circuitId.Secret, persistedCircuitEntry, options);
@@ -81,13 +82,13 @@ private void OnEntryEvicted(object key, object value, EvictionReason reason, obj
             // Happens after the circuit state times out, this is triggered by the CancellationTokenSource we register
             // with the entry, which is what controls the expiration
             case EvictionReason.Capacity:
-                // Happens when the cache is full
+            // Happens when the cache is full
                 var persistedCircuitEntry = (PersistedCircuitEntry)value;
                 Log.CircuitStateEvicted(_logger, persistedCircuitEntry.CircuitId, reason);
                 break;
 
             case EvictionReason.Removed:
-                // Happens when the entry is explicitly removed as part of resuming a circuit.
+            // Happens when the entry is explicitly removed as part of resuming a circuit.
                 return;
             default:
                 Debug.Fail($"Unexpected {nameof(EvictionReason)} {reason}");

src/Components/Server/src/Circuits/IServerComponentDeserializer.cs

@@ -12,5 +12,5 @@ bool TryDeserializeComponentDescriptorCollection(
         out List<ComponentDescriptor> descriptors);
     bool TryDeserializeRootComponentOperations(string serializedComponentOperations, [NotNullWhen(true)] out RootComponentOperationBatch? operationBatch, bool deserializeDescriptors = true);
 
-    public bool TryDeserializeWebRootComponentDescriptor(ComponentMarker record, [NotNullWhen(true)] out WebRootComponentDescriptor? result);
+    bool TryDeserializeWebRootComponentDescriptor(ComponentMarker record, [NotNullWhen(true)] out WebRootComponentDescriptor? result);
 }

src/Components/Server/src/Circuits/PersistedCircuitState.cs

@@ -8,7 +8,7 @@ namespace Microsoft.AspNetCore.Components.Server.Circuits;
 [DebuggerDisplay($"{{{nameof(GetDebuggerDisplay)}(),nq}}")]
 internal class PersistedCircuitState
 {
-    public Dictionary<string, byte[]> ApplicationState { get; internal set; }
+    public IReadOnlyDictionary<string, byte[]> ApplicationState { get; internal set; }
 
     public byte[] RootComponents { get; internal set; }
 

src/Components/Server/src/ComponentHub.cs

@@ -312,7 +312,7 @@ public async ValueTask<string> ResumeCircuit(
         }
 
         PersistedCircuitState? persistedCircuitState;
-        if (rootComponents == "[]" && string.IsNullOrEmpty(applicationState))
+        if (RootComponentIsEmpty(rootComponents) && string.IsNullOrEmpty(applicationState))
         {
             persistedCircuitState = await _circuitPersistenceManager.ResumeCircuitAsync(circuitId, Context.ConnectionAborted);
             if (persistedCircuitState == null)

src/Components/Server/test/Circuits/ComponentHubTest.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.Collections.ObjectModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Security.Claims;
 using System.Text.RegularExpressions;
@@ -164,7 +165,7 @@ public async Task CanCallUpdateRootComponentsOnResumedCircuit()
             .ReturnsAsync(new PersistedCircuitState
             {
                 RootComponents = [.. """{}"""u8],
-                ApplicationState = []
+                ApplicationState = ReadOnlyDictionary<string, byte[]>.Empty
             });
 
         var (mockClientProxy, hub) = InitializeComponentHub(deserializer, handleRegistryMock.Object, providerMock.Object);
@@ -264,7 +265,7 @@ public async Task CanResumeAppWhenPersistedComponentStateIsAvailable()
             .ReturnsAsync(new PersistedCircuitState
             {
                 RootComponents = [],
-                ApplicationState = []
+                ApplicationState = ReadOnlyDictionary<string, byte[]>.Empty,
             });
 
         var (mockClientProxy, hub) = InitializeComponentHub(null, handleRegistryMock.Object, providerMock.Object);

src/Components/Server/test/Circuits/DefaultPersistedCircuitCacheTest.cs renamed to src/Components/Server/test/Circuits/DefaultInMemoryCircuitPersistenceProviderTest.cs

@@ -9,7 +9,7 @@
 
 namespace Microsoft.AspNetCore.Components.Server.Tests.Circuits;
 
-public class DefaultPersistedCircuitCacheTest
+public class DefaultInMemoryCircuitPersistenceProviderTest
 {
     [Fact]
     public async Task PersistCircuitAsync_StoresCircuitState()

src/Components/WebAssembly/WebAssembly/src/Hosting/WebAssemblyHost.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.Collections.ObjectModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection.Metadata;
 using Microsoft.AspNetCore.Components.Infrastructure;
@@ -204,7 +205,7 @@ internal async Task RunAsyncCore(CancellationToken cancellationToken, WebAssembl
             });
 
             await initializationTcs.Task;
-            store.ExistingState.Clear();
+            store.ExistingState = ReadOnlyDictionary<string, byte[]>.Empty;
 
             await tcs.Task;
         }

src/Shared/Components/PrerenderComponentApplicationStore.cs

@@ -16,7 +16,7 @@ internal class PrerenderComponentApplicationStore : IPersistentComponentStateSto
 
     public PrerenderComponentApplicationStore()
     {
-        ExistingState = new();
+        ExistingState = new Dictionary<string, byte[]>();
     }
 
     [UnconditionalSuppressMessage("ReflectionAnalysis", "IL2026:RequiresUnreferencedCode", Justification = "Simple deserialize of primitive types.")]
@@ -41,7 +41,7 @@ protected void DeserializeState(byte[] existingState)
     public string? PersistedState { get; private set; }
 #nullable disable
 
-    public Dictionary<string, byte[]> ExistingState { get; protected set; }
+    public IReadOnlyDictionary<string, byte[]> ExistingState { get; protected internal set; }
 
     public Task<IDictionary<string, byte[]>> GetPersistedStateAsync()
     {

src/Shared/Components/ProtectedPrerenderComponentApplicationStore.cs

@@ -21,10 +21,10 @@ public ProtectedPrerenderComponentApplicationStore(string existingState, IDataPr
         DeserializeState(_protector.Unprotect(Convert.FromBase64String(existingState)));
     }
 
-    public ProtectedPrerenderComponentApplicationStore(Dictionary<string, byte[]> deserializedState, IDataProtectionProvider dataProtectionProvider)
+    public ProtectedPrerenderComponentApplicationStore(IReadOnlyDictionary<string, byte[]> deserializedState, IDataProtectionProvider dataProtectionProvider)
     {
         CreateProtector(dataProtectionProvider);
-        ExistingState = deserializedState;
+        ExistingState = new Dictionary<string, byte[]>(deserializedState);
     }
 
     protected override byte[] SerializeState(IReadOnlyDictionary<string, byte[]> state)