Periodically flush when using PipeWriter in Json

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

src/libraries/System.Text.Json/src/Resources/Strings.resx

@@ -731,8 +731,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>

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" />

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;

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;

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;

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;
                 }

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;

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;

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;

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;

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;

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;

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;

src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Object/ObjectDefaultConverter.cs

@@ -401,7 +401,7 @@ internal sealed override bool OnTryWrite(
                         state.Current.EndProperty();
                         state.Current.EnumeratorIndex++;
 
-                        if (ShouldFlush(writer, ref state))
+                        if (ShouldFlush(ref state))
                         {
                             return false;
                         }
@@ -430,7 +430,7 @@ internal sealed override bool OnTryWrite(
                         state.Current.EndProperty();
                         state.Current.EnumeratorIndex++;
 
-                        if (ShouldFlush(writer, ref state))
+                        if (ShouldFlush(ref state))
                         {
                             return false;
                         }

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);

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);