Cleanups

Author: javiercn

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

@@ -64,7 +64,9 @@ public Task PersistStateAsync(IPersistentComponentStateStore store, Renderer ren
 
             async Task PauseAndPersistState()
             {
+                State.PersistingState = true;
                 await PauseAsync();
+                State.PersistingState = false;
 
                 var data = new Dictionary<string, ReadOnlySequence<byte>>(StringComparer.Ordinal);
                 foreach (var (key, value) in _currentState)

src/Components/Components/src/PersistentComponentState.cs

@@ -5,7 +5,6 @@
 using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
-using System.IO.Pipelines;
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Components.Infrastructure;
@@ -20,6 +19,7 @@ public class PersistentComponentState
     {
         private IDictionary<string, ReadOnlySequence<byte>>? _existingState;
         private IDictionary<string, PooledByteBufferWriter> _currentState;
+
         private readonly List<Func<Task>> _registeredCallbacks;
 
         internal PersistentComponentState(
@@ -30,6 +30,8 @@ internal PersistentComponentState(
             _registeredCallbacks = pauseCallbacks;
         }
 
+        internal bool PersistingState { get; set; }
+
         internal void InitializeExistingState(IDictionary<string, ReadOnlySequence<byte>> existingState)
         {
             if (_existingState != null)
@@ -110,6 +112,11 @@ public void Persist(string key, Action<IBufferWriter<byte>> valueWriter)
                 throw new ArgumentNullException(nameof(valueWriter));
             }
 
+            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}'");
@@ -138,6 +145,11 @@ public void Persist(string key, Action<IBufferWriter<byte>> valueWriter)
                 throw new ArgumentNullException(nameof(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}'");

src/Shared/Components/PooledByteBufferWritter.cs

@@ -27,42 +27,15 @@ public PooledByteBufferWriter(byte[] existingBuffer)
             _owned = true;
         }
 
-        public ReadOnlyMemory<byte> WrittenMemory
-        {
-            get
-            {
-                return _currentBuffer.AsMemory(0, _index);
-            }
-        }
+        public ReadOnlyMemory<byte> WrittenMemory => _currentBuffer.AsMemory(0, _index);
 
-        public int WrittenCount
-        {
-            get
-            {
-                return _index;
-            }
-        }
+        public int WrittenCount => _index;
 
-        public int Capacity
-        {
-            get
-            {
-                return _currentBuffer.Length;
-            }
-        }
+        public int Capacity => _currentBuffer.Length;
 
-        public int FreeCapacity
-        {
-            get
-            {
-                return _currentBuffer.Length - _index;
-            }
-        }
+        public int FreeCapacity => _currentBuffer.Length - _index;
 
-        public void Clear()
-        {
-            ClearHelper();
-        }
+        public void Clear() => ClearHelper();
 
         private void ClearHelper()
         {

src/Shared/Components/PrerenderComponentApplicationStore.cs

@@ -66,25 +66,33 @@ protected virtual PooledByteBufferWriter SerializeState(IReadOnlyDictionary<stri
             // the data with a memory pool here. We will change our serialization strategy in the future here
             // so that we can avoid this step.
             var buffer = new PooledByteBufferWriter();
-            var jsonWriter = new Utf8JsonWriter(buffer);
-            jsonWriter.WriteStartObject();
-            foreach (var (key, value) in state)
+            try
             {
-                if (value.IsSingleSegment)
-                {
-                    jsonWriter.WriteBase64String(key, value.First.Span);
-                }
-                else
+                var jsonWriter = new Utf8JsonWriter(buffer);
+                jsonWriter.WriteStartObject();
+                foreach (var (key, value) in state)
                 {
-                    WriteMultipleSegments(jsonWriter, key, value);
+                    if (value.IsSingleSegment)
+                    {
+                        jsonWriter.WriteBase64String(key, value.First.Span);
+                    }
+                    else
+                    {
+                        WriteMultipleSegments(jsonWriter, key, value);
+                    }
+                    jsonWriter.Flush();
                 }
-                jsonWriter.Flush();
-            }
 
-            jsonWriter.WriteEndObject();
-            jsonWriter.Flush();
-            return buffer;
+                jsonWriter.WriteEndObject();
+                jsonWriter.Flush();
+                return buffer;
 
+            }
+            catch
+            {
+                buffer.Dispose();
+                throw;
+            }
             static void WriteMultipleSegments(Utf8JsonWriter jsonWriter, string key, ReadOnlySequence<byte> value)
             {
                 byte[] unescapedArray = null;