Implement thread-local caching for Utf8JsonWriter in JsonSerializer. (#…

…73338) * Implement thread-local caching for Utf8JsonWriter in JsonSerializer. * address feedback * Address feedback * Fix nullability warnings in ns2.0 build
dotnet · Aug 5, 2022 · 608e516 · 608e516
1 parent dafc6df
commit 608e516
Show file tree

Hide file tree

Showing 9 changed files with 257 additions and 51 deletions.
diff --git a/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs b/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs
@@ -13,11 +13,19 @@ namespace System.Text.Json
 {
     internal sealed class PooledByteBufferWriter : IBufferWriter<byte>, IDisposable
     {
-        private byte[] _rentedBuffer;
+        // This class allows two possible configurations: if rentedBuffer is not null then
+        // it can be used as an IBufferWriter and holds a buffer that should eventually be
+        // returned to the shared pool. If rentedBuffer is null, then the instance is in a
+        // cleared/disposed state and it must re-rent a buffer before it can be used again.
+        private byte[]? _rentedBuffer;
         private int _index;
 
         private const int MinimumBufferSize = 256;
 
+        private PooledByteBufferWriter()
+        {
+        }
+
         public PooledByteBufferWriter(int initialCapacity)
         {
             Debug.Assert(initialCapacity > 0);
@@ -68,6 +76,16 @@ public void Clear()
             ClearHelper();
         }
 
+        public void ClearAndReturnBuffers()
+        {
+            Debug.Assert(_rentedBuffer != null);
+
+            ClearHelper();
+            byte[] toReturn = _rentedBuffer;
+            _rentedBuffer = null;
+            ArrayPool<byte>.Shared.Return(toReturn);
+        }
+
         private void ClearHelper()
         {
             Debug.Assert(_rentedBuffer != null);
@@ -87,10 +105,21 @@ public void Dispose()
 
             ClearHelper();
             byte[] toReturn = _rentedBuffer;
-            _rentedBuffer = null!;
+            _rentedBuffer = null;
             ArrayPool<byte>.Shared.Return(toReturn);
         }
 
+        public void InitializeEmptyInstance(int initialCapacity)
+        {
+            Debug.Assert(initialCapacity > 0);
+            Debug.Assert(_rentedBuffer is null);
+
+            _rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
+            _index = 0;
+        }
+
+        public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter();
+
         public void Advance(int count)
         {
             Debug.Assert(_rentedBuffer != null);
@@ -125,11 +154,13 @@ internal void WriteToStream(Stream destination)
 #else
         internal Task WriteToStreamAsync(Stream destination, CancellationToken cancellationToken)
         {
+            Debug.Assert(_rentedBuffer != null);
             return destination.WriteAsync(_rentedBuffer, 0, _index, cancellationToken);
         }
 
         internal void WriteToStream(Stream destination)
         {
+            Debug.Assert(_rentedBuffer != null);
             destination.Write(_rentedBuffer, 0, _index);
         }
 #endif

diff --git a/src/libraries/System.Text.Json/src/System.Text.Json.csproj b/src/libraries/System.Text.Json/src/System.Text.Json.csproj
@@ -130,6 +130,7 @@ The System.Text.Json library is built-in as part of the shared framework in .NET
     <Compile Include="System\Text\Json\Serialization\Metadata\JsonTypeInfoKind.cs" />
     <Compile Include="System\Text\Json\Serialization\Metadata\CustomJsonTypeInfoOfT.cs" />
     <Compile Include="System\Text\Json\Serialization\PolymorphicSerializationState.cs" />
+    <Compile Include="System\Text\Json\Writer\Utf8JsonWriterCache.cs" />
     <Compile Include="System\Text\Json\Serialization\ReferenceEqualsWrapper.cs" />
     <Compile Include="System\Text\Json\Serialization\ConverterStrategy.cs" />
     <Compile Include="System\Text\Json\Serialization\ConfigurationList.cs" />

diff --git a/...ies/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.ByteArray.cs b/...ies/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.ByteArray.cs
@@ -120,25 +120,34 @@ private static byte[] WriteBytes<TValue>(in TValue value, JsonTypeInfo<TValue> j
         {
             Debug.Assert(jsonTypeInfo.IsConfigured);
 
-            JsonSerializerOptions options = jsonTypeInfo.Options;
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
-
-            WriteCore(writer, value, jsonTypeInfo);
-            return output.WrittenMemory.ToArray();
+            try
+            {
+                WriteCore(writer, value, jsonTypeInfo);
+                return output.WrittenMemory.ToArray();
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
 
         private static byte[] WriteBytesAsObject(object? value, JsonTypeInfo jsonTypeInfo)
         {
             Debug.Assert(jsonTypeInfo.IsConfigured);
 
-            JsonSerializerOptions options = jsonTypeInfo.Options;
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
-            WriteCoreAsObject(writer, value, jsonTypeInfo);
-            return output.WrittenMemory.ToArray();
+            try
+            {
+                WriteCoreAsObject(writer, value, jsonTypeInfo);
+                return output.WrittenMemory.ToArray();
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
     }
 }
diff --git a/...ries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Document.cs b/...ries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Document.cs
@@ -117,10 +117,17 @@ private static JsonDocument WriteDocument<TValue>(in TValue value, JsonTypeInfo<
             // For performance, share the same buffer across serialization and deserialization.
             // The PooledByteBufferWriter is cleared and returned when JsonDocument.Dispose() is called.
             PooledByteBufferWriter output = new(options.DefaultBufferSize);
-            using Utf8JsonWriter writer = new(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriter(options, output);
 
-            WriteCore(writer, value, jsonTypeInfo);
-            return JsonDocument.ParseRented(output, options.GetDocumentOptions());
+            try
+            {
+                WriteCore(writer, value, jsonTypeInfo);
+                return JsonDocument.ParseRented(output, options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriter(writer);
+            }
         }
 
         private static JsonDocument WriteDocumentAsObject(object? value, JsonTypeInfo jsonTypeInfo)
@@ -131,10 +138,17 @@ private static JsonDocument WriteDocumentAsObject(object? value, JsonTypeInfo js
             // For performance, share the same buffer across serialization and deserialization.
             // The PooledByteBufferWriter is cleared and returned when JsonDocument.Dispose() is called.
             PooledByteBufferWriter output = new(options.DefaultBufferSize);
-            using Utf8JsonWriter writer = new(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriter(options, output);
 
-            WriteCoreAsObject(writer, value, jsonTypeInfo);
-            return JsonDocument.ParseRented(output, options.GetDocumentOptions());
+            try
+            {
+                WriteCoreAsObject(writer, value, jsonTypeInfo);
+                return JsonDocument.ParseRented(output, options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriter(writer);
+            }
         }
     }
 }
diff --git a/...aries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Element.cs b/...aries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Element.cs
@@ -115,25 +115,35 @@ private static JsonElement WriteElement<TValue>(in TValue value, JsonTypeInfo<TV
             Debug.Assert(jsonTypeInfo.IsConfigured);
             JsonSerializerOptions options = jsonTypeInfo.Options;
 
-            // For performance, share the same buffer across serialization and deserialization.
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            WriteCore(writer, value, jsonTypeInfo);
-            return JsonElement.ParseValue(output.WrittenMemory.Span, options.GetDocumentOptions());
+            try
+            {
+                WriteCore(writer, value, jsonTypeInfo);
+                return JsonElement.ParseValue(output.WrittenMemory.Span, options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
 
         private static JsonElement WriteElementAsObject(object? value, JsonTypeInfo jsonTypeInfo)
         {
             JsonSerializerOptions options = jsonTypeInfo.Options;
             Debug.Assert(options != null);
 
-            // For performance, share the same buffer across serialization and deserialization.
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            WriteCoreAsObject(writer, value, jsonTypeInfo);
-            return JsonElement.ParseValue(output.WrittenMemory.Span, options.GetDocumentOptions());
+            try
+            {
+                WriteCoreAsObject(writer, value, jsonTypeInfo);
+                return JsonElement.ParseValue(output.WrittenMemory.Span, options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
     }
 }
diff --git a/...ibraries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Node.cs b/...ibraries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Node.cs
@@ -116,25 +116,35 @@ public static partial class JsonSerializer
             Debug.Assert(jsonTypeInfo.IsConfigured);
             JsonSerializerOptions options = jsonTypeInfo.Options;
 
-            // For performance, share the same buffer across serialization and deserialization.
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            WriteCore(writer, value, jsonTypeInfo);
-            return JsonNode.Parse(output.WrittenMemory.Span, options.GetNodeOptions(), options.GetDocumentOptions());
+            try
+            {
+                WriteCore(writer, value, jsonTypeInfo);
+                return JsonNode.Parse(output.WrittenMemory.Span, options.GetNodeOptions(), options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
 
         private static JsonNode? WriteNodeAsObject(object? value, JsonTypeInfo jsonTypeInfo)
         {
             Debug.Assert(jsonTypeInfo.IsConfigured);
             JsonSerializerOptions options = jsonTypeInfo.Options;
 
-            // For performance, share the same buffer across serialization and deserialization.
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            WriteCoreAsObject(writer, value, jsonTypeInfo);
-            return JsonNode.Parse(output.WrittenMemory.Span, options.GetNodeOptions(), options.GetDocumentOptions());
+            try
+            {
+                WriteCoreAsObject(writer, value, jsonTypeInfo);
+                return JsonNode.Parse(output.WrittenMemory.Span, options.GetNodeOptions(), options.GetDocumentOptions());
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
     }
 }
diff --git a/...raries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.String.cs b/...raries/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.String.cs
@@ -133,26 +133,34 @@ private static string WriteString<TValue>(in TValue value, JsonTypeInfo<TValue>
         {
             Debug.Assert(jsonTypeInfo.IsConfigured);
 
-            JsonSerializerOptions options = jsonTypeInfo.Options;
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
-
-            WriteCore(writer, value, jsonTypeInfo);
-            return JsonReaderHelper.TranscodeHelper(output.WrittenMemory.Span);
+            try
+            {
+                WriteCore(writer, value, jsonTypeInfo);
+                return JsonReaderHelper.TranscodeHelper(output.WrittenMemory.Span);
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
 
         private static string WriteStringAsObject(object? value, JsonTypeInfo jsonTypeInfo)
         {
             Debug.Assert(jsonTypeInfo.IsConfigured);
 
-            JsonSerializerOptions options = jsonTypeInfo.Options;
+            Utf8JsonWriter writer = Utf8JsonWriterCache.RentWriterAndBuffer(jsonTypeInfo.Options, out PooledByteBufferWriter output);
 
-            using var output = new PooledByteBufferWriter(options.DefaultBufferSize);
-            using var writer = new Utf8JsonWriter(output, options.GetWriterOptions());
-
-            WriteCoreAsObject(writer, value, jsonTypeInfo);
-            return JsonReaderHelper.TranscodeHelper(output.WrittenMemory.Span);
+            try
+            {
+                WriteCoreAsObject(writer, value, jsonTypeInfo);
+                return JsonReaderHelper.TranscodeHelper(output.WrittenMemory.Span);
+            }
+            finally
+            {
+                Utf8JsonWriterCache.ReturnWriterAndBuffer(writer, output);
+            }
         }
     }
 }
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Writer/Utf8JsonWriter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Writer/Utf8JsonWriter.cs
@@ -89,6 +89,10 @@ public sealed partial class Utf8JsonWriter : IDisposable, IAsyncDisposable
         /// </summary>
         public int CurrentDepth => _currentDepth & JsonConstants.RemoveFlagsBitMask;
 
+        private Utf8JsonWriter()
+        {
+        }
+
         /// <summary>
         /// Constructs a new <see cref="Utf8JsonWriter"/> instance with a specified <paramref name="bufferWriter"/>.
         /// </summary>
@@ -226,6 +230,29 @@ public void Reset(IBufferWriter<byte> bufferWriter)
             ResetHelper();
         }
 
+        internal void ResetAllStateForCacheReuse()
+        {
+            ResetHelper();
+
+            _stream = null;
+            _arrayBufferWriter = null;
+            _output = null;
+        }
+
+        internal void Reset(IBufferWriter<byte> bufferWriter, JsonWriterOptions options)
+        {
+            Debug.Assert(_output is null && _stream is null && _arrayBufferWriter is null);
+
+            _output = bufferWriter;
+            _options = options;
+            if (_options.MaxDepth == 0)
+            {
+                _options.MaxDepth = JsonWriterOptions.DefaultMaxDepth; // If max depth is not set, revert to the default depth.
+            }
+        }
+
+        internal static Utf8JsonWriter CreateEmptyInstanceForCaching() => new Utf8JsonWriter();
+
         private void ResetHelper()
         {
             BytesPending = default;