tmp

Author: javiercn

src/Shared/Components/PrerenderComponentApplicationStore.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
+using System.IO.Pipelines;
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Components.Lifetime;
@@ -24,7 +25,12 @@ public PrerenderComponentApplicationStore(string existingState)
                 throw new ArgumentNullException(nameof(existingState));
             }
 
-            var state = JsonSerializer.Deserialize<Dictionary<string, byte[]>>(Convert.FromBase64String(existingState));
+            DeserializeState(Convert.FromBase64String(existingState));
+        }
+
+        protected void DeserializeState(byte[] existingState)
+        {
+            var state = JsonSerializer.Deserialize<Dictionary<string, byte[]>>(existingState);
             if (state == null)
             {
                 throw new ArgumentException("Could not deserialize state correctly", nameof(existingState));
@@ -50,48 +56,74 @@ public Task<IDictionary<string, ReadOnlySequence<byte>>> GetPersistedStateAsync(
             return Task.FromResult((IDictionary<string, ReadOnlySequence<byte>>)ExistingState);
         }
 
-        protected virtual byte[] SerializeState(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
+        protected virtual async Task<byte[]> SerializeState(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
         {
             // System.Text.Json doesn't support serializing ReadonlySequence<byte> so we need to buffer
             // 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 pool = MemoryPool<byte>.Shared;
-            var memory = new List<IMemoryOwner<byte>>();
-            var serialization = new Dictionary<string, Memory<byte>>();
-            try
+            var pipe = new Pipe();
+            var pipeWriter = pipe.Writer;
+            var jsonWriter = new Utf8JsonWriter(pipeWriter);
+            jsonWriter.WriteStartObject();
+            foreach (var (key, value) in state)
+            {
+                if (value.IsSingleSegment)
+                {
+                    jsonWriter.WriteBase64String(key, value.First.Span);
+                }
+                else
+                {
+                    WriteMultipleSegments(jsonWriter, key, value);
+                }
+                jsonWriter.Flush();
+            }
+
+            jsonWriter.WriteEndObject();
+            jsonWriter.Flush();
+            pipe.Writer.Complete();
+            var result = await ReadToEnd(pipe.Reader);
+            return result.ToArray();
+
+            async Task<ReadOnlySequence<byte>> ReadToEnd(PipeReader reader)
             {
-                foreach (var (key, value) in state)
+                var result = await reader.ReadAsync();
+                reader.AdvanceTo(result.Buffer.Start, result.Buffer.End);
+                while (!result.IsCompleted)
                 {
-                    IMemoryOwner<byte> buffer = null;
-                    if (value.Length < pool.MaxBufferSize)
-                    {
-                        buffer = pool.Rent((int)value.Length);
-                        memory.Add(buffer);
-                        value.CopyTo(buffer.Memory.Span.Slice(0, (int)value.Length));
-                    }
-
-                    serialization.Add(key, buffer != null ? buffer.Memory.Slice(0, (int)value.Length) : value.ToArray());
+                    // Consume nothing, just wait for everything
+                    result = await reader.ReadAsync();
+                    reader.AdvanceTo(result.Buffer.Start, result.Buffer.End);
                 }
 
-                return JsonSerializer.SerializeToUtf8Bytes(serialization, JsonSerializerOptionsProvider.Options);
+                return result.Buffer;
             }
-            finally
+
+            static void WriteMultipleSegments(Utf8JsonWriter jsonWriter, string key, ReadOnlySequence<byte> value)
             {
-                serialization.Clear();
-                foreach (var item in memory)
+                byte[] unescapedArray = null;
+                var valueLenght = (int)value.Length;
+
+                Span<byte> valueSegment = value.Length <= 256 ?
+                    stackalloc byte[valueLenght] :
+                    (unescapedArray = ArrayPool<byte>.Shared.Rent(valueLenght)).AsSpan().Slice(0, valueLenght);
+
+                value.CopyTo(valueSegment);
+                jsonWriter.WriteBase64String(key, valueSegment);
+
+                if (unescapedArray != null)
                 {
-                    item.Dispose();
+                    valueSegment.Clear();
+                    ArrayPool<byte>.Shared.Return(unescapedArray);
                 }
             }
         }
 
-        public Task PersistStateAsync(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
+        public async Task PersistStateAsync(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
         {
-            var bytes = SerializeState(state);
+            var bytes = await SerializeState(state);
 
             var result = Convert.ToBase64String(bytes);
             PersistedState = result;
-            return Task.CompletedTask;
         }
     }
 }

src/Shared/Components/ProtectedPrerenderComponentApplicationStore.cs

@@ -5,6 +5,7 @@
 using System.Buffers;
 using System.Collections.Generic;
 using System.Text.Json;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.DataProtection;
 
 namespace Microsoft.AspNetCore.Components
@@ -21,12 +22,12 @@ public ProtectedPrerenderComponentApplicationStore(IDataProtectionProvider dataP
         public ProtectedPrerenderComponentApplicationStore(string existingState, IDataProtectionProvider dataProtectionProvider)
         {
             CreateProtector(dataProtectionProvider);
-            ExistingState = JsonSerializer.Deserialize<Dictionary<string, ReadOnlySequence<byte>>>(_protector.Unprotect(Convert.FromBase64String(existingState)));
+            DeserializeState(_protector.Unprotect(Convert.FromBase64String(existingState)));
         }
 
-        protected override byte[] SerializeState(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
+        protected async override Task<byte[]> SerializeState(IReadOnlyDictionary<string, ReadOnlySequence<byte>> state)
         {
-            var bytes = base.SerializeState(state);
+            var bytes = await base.SerializeState(state);
             if (_protector != null)
             {
                 bytes = _protector.Protect(bytes);