Fix writer depth issue on large stream async (#39417) (#39444)

Author: wtgodbe

src/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Helpers.cs

@@ -134,7 +134,7 @@ private static void WriteCore(Utf8JsonWriter writer, PooledByteBufferWriter outp
                 state.Current.Initialize(type, options);
                 state.Current.CurrentValue = value;
 
-                Write(writer, -1, options, ref state);
+                Write(writer, writer.CurrentDepth, flushThreshold: -1, options, ref state);
             }
 
             writer.Flush();

src/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.Stream.cs

@@ -74,13 +74,13 @@ private static async Task WriteAsyncCore(Stream utf8Json, object value, Type typ
                 state.Current.CurrentValue = value;
 
                 bool isFinalBlock;
-
                 int flushThreshold;
+
                 do
                 {
                     flushThreshold = (int)(bufferWriter.Capacity * .9); //todo: determine best value here
 
-                    isFinalBlock = Write(writer, flushThreshold, options, ref state);
+                    isFinalBlock = Write(writer, originalWriterDepth: 0, flushThreshold, options, ref state);
                     writer.Flush();
 
                     await bufferWriter.WriteToStreamAsync(utf8Json, cancellationToken).ConfigureAwait(false);

src/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializer.Write.cs

@@ -14,12 +14,12 @@ public static partial class JsonSerializer
         // 3) The object is an element in an enumerable.
         private static bool Write(
             Utf8JsonWriter writer,
+            int originalWriterDepth,
             int flushThreshold,
             JsonSerializerOptions options,
             ref WriteStack state)
         {
             bool finishedSerializing;
-            int currentDepth = writer.CurrentDepth;
 
             try
             {
@@ -53,7 +53,7 @@ private static bool Write(
 
                     if (finishedSerializing)
                     {
-                        if (writer.CurrentDepth == 0 || writer.CurrentDepth == currentDepth)
+                        if (writer.CurrentDepth == 0 || writer.CurrentDepth == originalWriterDepth)
                         {
                             break;
                         }

src/System.Text.Json/tests/Serialization/Stream.WriteTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
 using Xunit;
@@ -149,6 +150,112 @@ public static async Task WritePrimitivesAsync()
             int i = await JsonSerializer.DeserializeAsync<int>(stream, options);
             Assert.Equal(1, i);
         }
+
+        private class Session
+        {
+            public int Id { get; set; }
+            public string Title { get; set; }
+            public virtual string Abstract { get; set; }
+            public virtual DateTimeOffset? StartTime { get; set; }
+            public virtual DateTimeOffset? EndTime { get; set; }
+            public TimeSpan Duration => EndTime?.Subtract(StartTime ?? EndTime ?? DateTimeOffset.MinValue) ?? TimeSpan.Zero;
+            public int? TrackId { get; set; }
+        }
+
+        private class SessionResponse : Session
+        {
+            public Track Track { get; set; }
+            public List<Speaker> Speakers { get; set; } = new List<Speaker>();
+        }
+
+        private class Track
+        {
+            public int Id { get; set; }
+            public string Name { get; set; }
+        }
+
+        private class Speaker
+        {
+            public int Id { get; set; }
+            public string Name { get; set; }
+            public string Bio { get; set; }
+            public virtual string WebSite { get; set; }
+        }
+
+        [Theory]
+        [InlineData(0)]
+        [InlineData(10)]
+        [InlineData(1000)]
+        [InlineData(4000)]
+        [InlineData(8000)]
+        [InlineData(16000)]
+        public static async Task LargeJsonFile(int bufferSize)
+        {
+            // Build up a large list to serialize.
+            var list = new List<SessionResponse>();
+            for (int i = 0; i < 100; i++)
+            {
+                SessionResponse response = new SessionResponse
+                {
+                    Id = i,
+                    Abstract = new string('A', i * 2),
+                    Title = new string('T', i),
+                    StartTime = new DateTime(i),
+                    EndTime = new DateTime(i * 10000),
+                    TrackId = i,
+                    Track = new Track()
+                    {
+                        Id = i,
+                        Name = new string('N', i),
+                    },
+                };
+
+                for (int j = 0; j < 5; j++)
+                {
+                    response.Speakers.Add(new Speaker()
+                    {
+                        Bio = new string('B', 50),
+                        Id = j,
+                        Name = new string('N', i),
+                        WebSite = new string('W', 20),
+                    });
+                }
+
+                list.Add(response);
+            }
+
+            // Adjust buffer length to encourage buffer flusing at several levels.
+            JsonSerializerOptions options = new JsonSerializerOptions();
+            if (bufferSize != 0)
+            {
+                options.DefaultBufferSize = bufferSize;
+            }
+
+            string json = JsonSerializer.Serialize(list, options);
+            Assert.True(json.Length > 100_000); // Verify data is large and will cause buffer flushing.
+            Assert.True(json.Length < 200_000); // But not too large for memory considerations.
+
+            // Sync case.
+            {
+                List<SessionResponse> deserializedList = JsonSerializer.Deserialize<List<SessionResponse>>(json, options);
+                Assert.Equal(100, deserializedList.Count);
+
+                string jsonSerialized = JsonSerializer.Serialize(deserializedList, options);
+                Assert.Equal(json, jsonSerialized);
+            }
+
+            // Async case.
+            using (var memoryStream = new MemoryStream())
+            {
+                await JsonSerializer.SerializeAsync(memoryStream, list, options);
+                string jsonSerialized = Encoding.UTF8.GetString(memoryStream.ToArray());
+                Assert.Equal(json, jsonSerialized);
+
+                memoryStream.Position = 0;
+                List<SessionResponse> deserializedList = await JsonSerializer.DeserializeAsync<List<SessionResponse>>(memoryStream, options);
+                Assert.Equal(100, deserializedList.Count);
+            }
+        }
     }
 
     public sealed class TestStream : Stream