Add read ahead logic for streams.

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

@@ -83,8 +83,8 @@ private static async ValueTask<TValue> ReadAsync<TValue>(
                 options = JsonSerializerOptions.s_defaultOptions;
             }
 
-            ReadStack state = default;
-            state.Current.Initialize(returnType, options);
+            ReadStack readStack = default;
+            readStack.Current.Initialize(returnType, options);
 
             var readerState = new JsonReaderState(options.GetReaderOptions());
 
@@ -138,10 +138,11 @@ private static async ValueTask<TValue> ReadAsync<TValue>(
                         isFinalBlock,
                         new Span<byte>(buffer, 0, bytesInBuffer),
                         options,
-                        ref state);
+                        ref readStack);
+
+                    Debug.Assert(readStack.BytesConsumed <= bytesInBuffer);
+                    int bytesConsumed = checked((int)readStack.BytesConsumed);
 
-                    Debug.Assert(readerState.BytesConsumed <= bytesInBuffer);
-                    int bytesConsumed = (int)readerState.BytesConsumed;
                     bytesInBuffer -= bytesConsumed;
 
                     if (isFinalBlock)
@@ -183,22 +184,29 @@ private static async ValueTask<TValue> ReadAsync<TValue>(
                 ThrowHelper.ThrowJsonException_DeserializeDataRemaining(totalBytesRead, bytesInBuffer);
             }
 
-            return (TValue)state.Current.ReturnValue;
+            return (TValue)readStack.Current.ReturnValue;
         }
 
         private static void ReadCore(
             ref JsonReaderState readerState,
             bool isFinalBlock,
             Span<byte> buffer,
             JsonSerializerOptions options,
-            ref ReadStack state)
+            ref ReadStack readStack)
         {
             var reader = new Utf8JsonReader(buffer, isFinalBlock, readerState);
 
+            // If we haven't read in the entire stream's payload we'll need to signify that we want
+            // to enable read ahead behaviors to ensure we have complete json objects and arrays
+            // ({}, []) when needed. (Notably to successfully parse JsonElement via JsonDocument
+            // to assign to object and JsonElement properties in the constructed .NET object.)
+            options.ReadAhead = !isFinalBlock;
+            readStack.BytesConsumed = 0;
+
             ReadCore(
                 options,
                 ref reader,
-                ref state);
+                ref readStack);
 
             readerState = reader.CurrentState;
         }

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

@@ -3,8 +3,8 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Buffers;
-using System.Collections;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 
 namespace System.Text.Json.Serialization
 {
@@ -17,101 +17,150 @@ public static partial class JsonSerializer
         private static void ReadCore(
             JsonSerializerOptions options,
             ref Utf8JsonReader reader,
-            ref ReadStack state)
+            ref ReadStack readStack)
         {
             try
             {
-                while (reader.Read())
+                JsonReaderState initialState = default;
+
+                while (true)
                 {
+                    if (options.ReadAhead)
+                    {
+                        // When we're reading ahead we always have to save the state
+                        // as we don't know if the next token is an opening object or
+                        // array brace.
+                        initialState = reader.CurrentState;
+                    }
+
+                    if (!reader.Read())
+                    {
+                        // Need more data
+                        break;
+                    }
+
                     JsonTokenType tokenType = reader.TokenType;
 
                     if (JsonHelpers.IsInRangeInclusive(tokenType, JsonTokenType.String, JsonTokenType.False))
                     {
                         Debug.Assert(tokenType == JsonTokenType.String || tokenType == JsonTokenType.Number || tokenType == JsonTokenType.True || tokenType == JsonTokenType.False);
 
-                        if (HandleValue(tokenType, options, ref reader, ref state))
-                        {
-                            continue;
-                        }
+                        HandleValue(tokenType, options, ref reader, ref readStack);
                     }
                     else if (tokenType == JsonTokenType.PropertyName)
                     {
-                        HandlePropertyName(options, ref reader, ref state);
+                        HandlePropertyName(options, ref reader, ref readStack);
                     }
                     else if (tokenType == JsonTokenType.StartObject)
                     {
-                        if (state.Current.SkipProperty)
+                        if (readStack.Current.SkipProperty)
                         {
-                            state.Push();
-                            state.Current.Drain = true;
+                            readStack.Push();
+                            readStack.Current.Drain = true;
                         }
-                        else if (state.Current.IsProcessingValue)
+                        else if (readStack.Current.IsProcessingValue)
                         {
-                            if (HandleValue(tokenType, options, ref reader, ref state))
+                            if (!HandleObjectAsValue(tokenType, options, ref reader, ref readStack, ref initialState))
                             {
-                                continue;
+                                // Need more data
+                                break;
                             }
                         }
-                        else if (state.Current.IsProcessingDictionary)
+                        else if (readStack.Current.IsProcessingDictionary)
                         {
-                            HandleStartDictionary(options, ref reader, ref state);
+                            HandleStartDictionary(options, ref reader, ref readStack);
                         }
                         else
                         {
-                            HandleStartObject(options, ref reader, ref state);
+                            HandleStartObject(options, ref reader, ref readStack);
                         }
                     }
                     else if (tokenType == JsonTokenType.EndObject)
                     {
-                        if (state.Current.Drain)
+                        if (readStack.Current.Drain)
                         {
-                            state.Pop();
+                            readStack.Pop();
                         }
-                        else if (state.Current.IsProcessingDictionary)
+                        else if (readStack.Current.IsProcessingDictionary)
                         {
-                            HandleEndDictionary(options, ref reader, ref state);
+                            HandleEndDictionary(options, ref reader, ref readStack);
                         }
                         else
                         {
-                            HandleEndObject(options, ref reader, ref state);
+                            HandleEndObject(options, ref reader, ref readStack);
                         }
                     }
                     else if (tokenType == JsonTokenType.StartArray)
                     {
-                        if (!state.Current.IsProcessingValue)
+                        if (!readStack.Current.IsProcessingValue)
                         {
-                            HandleStartArray(options, ref reader, ref state);
+                            HandleStartArray(options, ref reader, ref readStack);
                         }
-                        else if (HandleValue(tokenType, options, ref reader, ref state))
+                        else if (!HandleObjectAsValue(tokenType, options, ref reader, ref readStack, ref initialState))
                         {
-                            continue;
+                            // Need more data
+                            break;
                         }
                     }
                     else if (tokenType == JsonTokenType.EndArray)
                     {
-                        if (HandleEndArray(options, ref reader, ref state))
-                        {
-                            continue;
-                        }
+                        HandleEndArray(options, ref reader, ref readStack);
                     }
                     else if (tokenType == JsonTokenType.Null)
                     {
-                        if (HandleNull(ref reader, ref state, options))
-                        {
-                            continue;
-                        }
+                        HandleNull(ref reader, ref readStack, options);
                     }
                 }
             }
             catch (JsonReaderException e)
             {
                 // Re-throw with Path information.
-                ThrowHelper.ReThrowWithPath(e, state.JsonPath);
+                ThrowHelper.ReThrowWithPath(e, readStack.JsonPath);
             }
 
+            readStack.BytesConsumed += reader.BytesConsumed;
             return;
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static bool HandleObjectAsValue(
+            JsonTokenType tokenType,
+            JsonSerializerOptions options,
+            ref Utf8JsonReader reader,
+            ref ReadStack readStack,
+            ref JsonReaderState initialState)
+        {
+            if (options.ReadAhead)
+            {
+                // Attempt to skip to make sure we have all the data we need.
+                bool complete = reader.TrySkip();
+
+                // We need to restore the state in all cases as we need to be positioned back before
+                // the current token to either attempt to skip again or to actually read the value in
+                // HandleValue below.
+
+                reader = new Utf8JsonReader(
+                    reader.OriginalSpan.Slice(checked((int)initialState.BytesConsumed)),
+                    isFinalBlock: reader.IsFinalBlock,
+                    state: initialState);
+                Debug.Assert(reader.BytesConsumed == 0);
+                readStack.BytesConsumed += initialState.BytesConsumed;
+
+                if (!complete)
+                {
+                    // Couldn't read to the end of the object, exit out to get more data in the buffer.
+                    return false;
+                }
+
+                // Success, requeue the reader to the token for HandleValue.
+                reader.Read();
+                Debug.Assert(tokenType == reader.TokenType);
+            }
+
+            HandleValue(tokenType, options, ref reader, ref readStack);
+            return true;
+        }
+
         private static ReadOnlySpan<byte> GetUnescapedString(ReadOnlySpan<byte> utf8Source, int idx)
         {
             // The escaped name is always longer than the unescaped, so it is safe to use escaped name for the buffer length.

src/System.Text.Json/src/System/Text/Json/Serialization/JsonSerializerOptions.cs

@@ -314,5 +314,10 @@ private void VerifyMutable()
                 ThrowHelper.ThrowInvalidOperationException_SerializerOptionsImmutable();
             }
         }
+
+        /// <summary>
+        /// Internal flag to let us know that we need to read ahead in the inner read loop.
+        /// </summary>
+        internal bool ReadAhead { get; set; }
     }
 }

src/System.Text.Json/src/System/Text/Json/Serialization/ReadStack.cs

@@ -145,5 +145,10 @@ private string GetPropertyName(in ReadStackFrame frame)
 
             return propertyName;
         }
+
+        /// <summary>
+        /// Bytes consumed in the current loop
+        /// </summary>
+        public long BytesConsumed;
     }
 }

src/System.Text.Json/tests/Serialization/JsonElementTests.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.IO;
 using System.Linq;
 using Xunit;
 
@@ -145,5 +146,32 @@ public void Verify()
                 Assert.Equal("Hello", Array[3].ToString());
             }
         }
+
+        [Theory,
+            InlineData(5),
+            InlineData(10),
+            InlineData(20),
+            InlineData(1024)]
+        public void ReadJsonElementFromStream(int defaultBufferSize)
+        {
+            // Streams need to read ahead when they hit objects or arrays that are assigned to JsonElement or object.
+
+            byte[] data = Encoding.UTF8.GetBytes(@"{""Data"":[1,true,{""City"":""MyCity""},null,""foo""]}");
+            MemoryStream stream = new MemoryStream(data);
+            JsonElement obj = JsonSerializer.ReadAsync<JsonElement>(stream, new JsonSerializerOptions { DefaultBufferSize = defaultBufferSize }).Result;
+
+            data = Encoding.UTF8.GetBytes(@"[1,true,{""City"":""MyCity""},null,""foo""]");
+            stream = new MemoryStream(data);
+            obj = JsonSerializer.ReadAsync<JsonElement>(stream, new JsonSerializerOptions { DefaultBufferSize = defaultBufferSize }).Result;
+
+            // Ensure we fail with incomplete data
+            data = Encoding.UTF8.GetBytes(@"{""Data"":[1,true,{""City"":""MyCity""},null,""foo""]");
+            stream = new MemoryStream(data);
+            Assert.Throws<JsonException>(() => JsonSerializer.ReadAsync<JsonElement>(stream, new JsonSerializerOptions { DefaultBufferSize = defaultBufferSize }).Result);
+
+            data = Encoding.UTF8.GetBytes(@"[1,true,{""City"":""MyCity""},null,""foo""");
+            stream = new MemoryStream(data);
+            Assert.Throws<JsonException>(() => JsonSerializer.ReadAsync<JsonElement>(stream, new JsonSerializerOptions { DefaultBufferSize = defaultBufferSize }).Result);
+        }
     }
 }

src/System.Text.Json/tests/Serialization/SpanTests.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Generic;
+using System.IO;
 using Xunit;
 
 namespace System.Text.Json.Serialization.Tests
@@ -24,6 +25,29 @@ public static void Read(Type classType, byte[] data)
             ((ITestClass)obj).Verify();
         }
 
+        [Theory]
+        [MemberData(nameof(ReadSuccessCases))]
+        public static void ReadFromStream(Type classType, byte[] data)
+        {
+            MemoryStream stream = new MemoryStream(data);
+            object obj = JsonSerializer.ReadAsync(
+                stream,
+                classType).Result;
+
+            Assert.IsAssignableFrom(typeof(ITestClass), obj);
+            ((ITestClass)obj).Verify();
+
+            // Try again with a smaller initial buffer size to ensure we handle incomplete data
+            stream = new MemoryStream(data);
+            obj = JsonSerializer.ReadAsync(
+                stream,
+                classType,
+                new JsonSerializerOptions { DefaultBufferSize = 5 }).Result;
+
+            Assert.IsAssignableFrom(typeof(ITestClass), obj);
+            ((ITestClass)obj).Verify();
+        }
+
         [Fact]
         public static void ReadGenericApi()
         {