Periodically flush when using PipeWriter in Json (dotnet#102541)

BrennanConroy · Ruihan-Yin · commit 3a824a83b441 · 2024-05-30T10:31:56.000-07:00
diff --git a/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs b/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs
@@ -5,20 +5,22 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
+using System.IO.Pipelines;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Text.Json
 {
-    internal sealed class PooledByteBufferWriter : IBufferWriter<byte>, IDisposable
+    internal sealed class PooledByteBufferWriter : PipeWriter, IDisposable
     {
         // 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 readonly Stream? _stream;
 
         private const int MinimumBufferSize = 256;
 
@@ -41,6 +43,11 @@ public PooledByteBufferWriter(int initialCapacity) : this()
             _index = 0;
         }
 
+        public PooledByteBufferWriter(int initialCapacity, Stream stream) : this(initialCapacity)
+        {
+            _stream = stream;
+        }
+
         public ReadOnlyMemory<byte> WrittenMemory
         {
             get
@@ -127,43 +134,32 @@ public void InitializeEmptyInstance(int initialCapacity)
 
         public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter();
 
-        public void Advance(int count)
+        public override void Advance(int count)
         {
             Debug.Assert(_rentedBuffer != null);
             Debug.Assert(count >= 0);
             Debug.Assert(_index <= _rentedBuffer.Length - count);
             _index += count;
         }
 
-        public Memory<byte> GetMemory(int sizeHint = MinimumBufferSize)
+        public override Memory<byte> GetMemory(int sizeHint = MinimumBufferSize)
         {
             CheckAndResizeBuffer(sizeHint);
             return _rentedBuffer.AsMemory(_index);
         }
 
-        public Span<byte> GetSpan(int sizeHint = MinimumBufferSize)
+        public override Span<byte> GetSpan(int sizeHint = MinimumBufferSize)
         {
             CheckAndResizeBuffer(sizeHint);
             return _rentedBuffer.AsSpan(_index);
         }
 
 #if NET
-        internal ValueTask WriteToStreamAsync(Stream destination, CancellationToken cancellationToken)
-        {
-            return destination.WriteAsync(WrittenMemory, cancellationToken);
-        }
-
         internal void WriteToStream(Stream destination)
         {
             destination.Write(WrittenMemory.Span);
         }
 #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);
@@ -217,6 +213,28 @@ private void CheckAndResizeBuffer(int sizeHint)
             Debug.Assert(_rentedBuffer.Length - _index > 0);
             Debug.Assert(_rentedBuffer.Length - _index >= sizeHint);
         }
+
+        public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
+        {
+            Debug.Assert(_stream is not null);
+#if NET
+            await _stream.WriteAsync(WrittenMemory, cancellationToken).ConfigureAwait(false);
+#else
+            Debug.Assert(_rentedBuffer != null);
+            await _stream.WriteAsync(_rentedBuffer, 0, _index, cancellationToken).ConfigureAwait(false);
+#endif
+            Clear();
+
+            return new FlushResult(isCanceled: false, isCompleted: false);
+        }
+
+        public override bool CanGetUnflushedBytes => true;
+        public override long UnflushedBytes => _index;
+
+        // This type is used internally in JsonSerializer to help buffer and flush bytes to the underlying Stream.
+        // It's only pretending to be a PipeWriter and doesn't need Complete or CancelPendingFlush for the internal usage.
+        public override void CancelPendingFlush() => throw new NotImplementedException();
+        public override void Complete(Exception? exception = null) => throw new NotImplementedException();
     }
 
     internal static partial class ThrowHelper
diff --git a/src/libraries/System.Text.Json/src/Resources/Strings.resx b/src/libraries/System.Text.Json/src/Resources/Strings.resx
@@ -734,8 +734,8 @@
   <data name="PipeWriterCanceled" xml:space="preserve">
     <value>PipeWriter.FlushAsync was canceled.</value>
   </data>
-  <data name="PipeWriterCompleted" xml:space="preserve">
-    <value>PipeWriter has been completed, nothing more can be written to it.</value>
+  <data name="PipeWriter_DoesNotImplementUnflushedBytes" xml:space="preserve">
+    <value>The PipeWriter '{0}' does not implement PipeWriter.UnflushedBytes.</value>
   </data>
   <data name="InvalidNewLine" xml:space="preserve">
     <value>New line can be only "\n" or "\r\n".</value>
diff --git a/src/libraries/System.Text.Json/src/System.Text.Json.csproj b/src/libraries/System.Text.Json/src/System.Text.Json.csproj
@@ -102,7 +102,6 @@ The System.Text.Json library is built-in as part of the shared framework in .NET
     <Compile Include="System\Text\Json\Reader\Utf8JsonReader.TryGet.cs" />
     <Compile Include="System\Text\Json\Serialization\Arguments.cs" />
     <Compile Include="System\Text\Json\Serialization\ArgumentState.cs" />
-    <Compile Include="System\Text\Json\Serialization\AsyncSerializationBufferWriterContext.cs" />
     <Compile Include="System\Text\Json\Serialization\Attributes\JsonObjectCreationHandlingAttribute.cs" />
     <Compile Include="System\Text\Json\Serialization\Attributes\JsonConstructorAttribute.cs" />
     <Compile Include="System\Text\Json\Serialization\Attributes\JsonConverterAttribute.cs" />
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/AsyncSerializationBufferWriterContext.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/AsyncSerializationBufferWriterContext.cs
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ArrayConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ArrayConverter.cs
@@ -56,7 +56,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TElement[] array, J
 
                     state.Current.EndCollectionElement();
 
-                    if (ShouldFlush(writer, ref state))
+                    if (ShouldFlush(ref state))
                     {
                         state.Current.EnumeratorIndex = ++index;
                         return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/DictionaryDefaultConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/DictionaryDefaultConverter.cs
@@ -45,7 +45,7 @@ protected internal override bool OnWriteResume(
 
             do
             {
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     state.Current.CollectionEnumerator = enumerator;
                     return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/DictionaryOfTKeyTValueConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/DictionaryOfTKeyTValueConverter.cs
@@ -61,7 +61,7 @@ protected internal override bool OnWriteResume(
             {
                 do
                 {
-                    if (ShouldFlush(writer, ref state))
+                    if (ShouldFlush(ref state))
                     {
                         state.Current.CollectionEnumerator = enumerator;
                         return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IAsyncEnumerableOfTConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IAsyncEnumerableOfTConverter.cs
@@ -107,7 +107,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TAsyncEnumerable va
                     return true;
                 }
 
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     return false;
                 }
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IDictionaryConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IDictionaryConverter.cs
@@ -60,7 +60,7 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar
 
             do
             {
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     state.Current.CollectionEnumerator = enumerator;
                     return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableConverter.cs
@@ -59,7 +59,7 @@ protected override bool OnWriteResume(
             JsonConverter<object?> converter = GetElementConverter(ref state);
             do
             {
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     state.Current.CollectionEnumerator = enumerator;
                     return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableDefaultConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableDefaultConverter.cs
@@ -37,7 +37,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,
             JsonConverter<TElement> converter = GetElementConverter(ref state);
             do
             {
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     state.Current.CollectionEnumerator = enumerator;
                     return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IListConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IListConverter.cs
@@ -65,7 +65,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,
 
                     state.Current.EndCollectionElement();
 
-                    if (ShouldFlush(writer, ref state))
+                    if (ShouldFlush(ref state))
                     {
                         state.Current.EnumeratorIndex = ++index;
                         return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ListOfTConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ListOfTConverter.cs
@@ -62,7 +62,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,
 
                     state.Current.EndCollectionElement();
 
-                    if (ShouldFlush(writer, ref state))
+                    if (ShouldFlush(ref state))
                     {
                         state.Current.EnumeratorIndex = ++index;
                         return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ReadOnlyMemoryConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/ReadOnlyMemoryConverter.cs
@@ -74,7 +74,7 @@ internal static bool OnWriteResume(Utf8JsonWriter writer, ReadOnlySpan<T> value,
 
                     state.Current.EndCollectionElement();
 
-                    if (ShouldFlush(writer, ref state))
+                    if (ShouldFlush(ref state))
                     {
                         state.Current.EnumeratorIndex = ++index;
                         return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/StackOrQueueConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/StackOrQueueConverter.cs
@@ -59,7 +59,7 @@ protected sealed override bool OnWriteResume(Utf8JsonWriter writer, TCollection
             JsonConverter<object?> converter = GetElementConverter(ref state);
             do
             {
-                if (ShouldFlush(writer, ref state))
+                if (ShouldFlush(ref state))
                 {
                     state.Current.CollectionEnumerator = enumerator;
                     return false;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Object/ObjectDefaultConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Object/ObjectDefaultConverter.cs
@@ -403,7 +403,7 @@ internal sealed override bool OnTryWrite(
                         state.Current.EndProperty();
                         state.Current.EnumeratorIndex++;
 
-                        if (ShouldFlush(writer, ref state))
+                        if (ShouldFlush(ref state))
                         {
                             return false;
                         }
@@ -432,7 +432,7 @@ internal sealed override bool OnTryWrite(
                         state.Current.EndProperty();
                         state.Current.EnumeratorIndex++;
 
-                        if (ShouldFlush(writer, ref state))
+                        if (ShouldFlush(ref state))
                         {
                             return false;
                         }
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/JsonConverter.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/JsonConverter.cs
@@ -171,10 +171,17 @@ internal JsonConverter<TTarget> CreateCastingConverter<TTarget>()
         /// </summary>
         internal bool IsInternalConverterForNumberType { get; init; }
 
-        internal static bool ShouldFlush(Utf8JsonWriter writer, ref WriteStack state)
+        internal static bool ShouldFlush(ref WriteStack state)
         {
-            // If surpassed flush threshold then return false which will flush stream.
-            return (state.FlushThreshold > 0 && writer.BytesPending > state.FlushThreshold);
+            Debug.Assert(state.FlushThreshold == 0 || (state.PipeWriter is { CanGetUnflushedBytes: true }),
+                "ShouldFlush should only be called by resumable serializers, all of which use the PipeWriter abstraction with CanGetUnflushedBytes == true.");
+            // If surpassed flush threshold then return true which will flush stream.
+            if (state.PipeWriter is { } pipeWriter)
+            {
+                return state.FlushThreshold > 0 && pipeWriter.UnflushedBytes > state.FlushThreshold;
+            }
+
+            return false;
         }
 
         internal abstract object? ReadAsObject(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options);
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfo.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfo.cs
@@ -990,8 +990,7 @@ public JsonPropertyInfo CreateJsonPropertyInfo(Type propertyType, string name)
 
         // Untyped, root-level serialization methods
         internal abstract void SerializeAsObject(Utf8JsonWriter writer, object? rootValue);
-        internal abstract Task SerializeAsObjectAsync<TSerializationContext>(TSerializationContext serializationContext, object? rootValue, CancellationToken cancellationToken)
-            where TSerializationContext : struct, IAsyncSerializationBufferWriterContext;
+        internal abstract Task SerializeAsObjectAsync(PipeWriter pipeWriter, object? rootValue, int flushThreshold, CancellationToken cancellationToken);
         internal abstract Task SerializeAsObjectAsync(Stream utf8Json, object? rootValue, CancellationToken cancellationToken);
         internal abstract Task SerializeAsObjectAsync(PipeWriter utf8Json, object? rootValue, CancellationToken cancellationToken);
         internal abstract void SerializeAsObject(Stream utf8Json, object? rootValue);
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfoOfT.WriteHelpers.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfoOfT.WriteHelpers.cs
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/WriteStack.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/WriteStack.cs
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.Serialization.cs b/src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.Serialization.cs
diff --git a/src/libraries/System.Text.Json/tests/System.Text.Json.Tests/Serialization/Pipe.WriteTests.cs b/src/libraries/System.Text.Json/tests/System.Text.Json.Tests/Serialization/Pipe.WriteTests.cs
diff --git a/src/libraries/System.Text.Json/tests/System.Text.Json.Tests/Serialization/Stream.WriteTests.cs b/src/libraries/System.Text.Json/tests/System.Text.Json.Tests/Serialization/Stream.WriteTests.cs

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TElement[] array, J`
`56`	`56`
`57`	`57`	`state.Current.EndCollectionElement();`
`58`	`58`
`59`		`- if (ShouldFlush(writer, ref state))`
	`59`	`+ if (ShouldFlush(ref state))`
`60`	`60`	`{`
`61`	`61`	`state.Current.EnumeratorIndex = ++index;`
`62`	`62`	`return false;`
Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ protected internal override bool OnWriteResume(`
`45`	`45`
`46`	`46`	`do`
`47`	`47`	`{`
`48`		`- if (ShouldFlush(writer, ref state))`
	`48`	`+ if (ShouldFlush(ref state))`
`49`	`49`	`{`
`50`	`50`	`state.Current.CollectionEnumerator = enumerator;`
`51`	`51`	`return false;`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ protected internal override bool OnWriteResume(`
`61`	`61`	`{`
`62`	`62`	`do`
`63`	`63`	`{`
`64`		`- if (ShouldFlush(writer, ref state))`
	`64`	`+ if (ShouldFlush(ref state))`
`65`	`65`	`{`
`66`	`66`	`state.Current.CollectionEnumerator = enumerator;`
`67`	`67`	`return false;`
Original file line number	Diff line number	Diff line change
`@@ -107,7 +107,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TAsyncEnumerable va`
`107`	`107`	`return true;`
`108`	`108`	`}`
`109`	`109`
`110`		`- if (ShouldFlush(writer, ref state))`
	`110`	`+ if (ShouldFlush(ref state))`
`111`	`111`	`{`
`112`	`112`	`return false;`
`113`	`113`	`}`
Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar`
`60`	`60`
`61`	`61`	`do`
`62`	`62`	`{`
`63`		`- if (ShouldFlush(writer, ref state))`
	`63`	`+ if (ShouldFlush(ref state))`
`64`	`64`	`{`
`65`	`65`	`state.Current.CollectionEnumerator = enumerator;`
`66`	`66`	`return false;`
Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ protected override bool OnWriteResume(`
`59`	`59`	`JsonConverter<object?> converter = GetElementConverter(ref state);`
`60`	`60`	`do`
`61`	`61`	`{`
`62`		`- if (ShouldFlush(writer, ref state))`
	`62`	`+ if (ShouldFlush(ref state))`
`63`	`63`	`{`
`64`	`64`	`state.Current.CollectionEnumerator = enumerator;`
`65`	`65`	`return false;`