Fix JsonChunkAssembler bugs; add streaming JSON tests & benchmarks

Bug 1 — only 1 element returned per chunk:
  ProcessBuffer had needBreak=true after finding one element, causing it
  to stop immediately. Removed the flag; the reader now processes ALL
  complete elements in a single pass without early exit.

Bug 2 — JsonReaderException on wide rows (FOR JSON splits mid-string):
  The outer loop was compacting the buffer even while elemStart >= 0
  (mid-element), which shifted rawBuf bytes and made elemStart point into
  the middle of a string instead of its opening '{'. Fixed: CompactBuffer
  is now only called when consumedTo > 0, and absolute positions
  (parsedTo, elemStart) are adjusted by the compaction offset.

Also added inner do-while loop so one chunk worth of data is fully drained
before awaiting the next chunk (buffer may contain many complete elements).

Integration tests (MsSqlIntegrationTests.cs):
  - QueryJsonStreamAsync_SmallDataset_YieldsCorrectCount  — 20 rows, all properties validated
  - QueryJsonStreamAsync_LargeDataset_YieldsAllRows       — 5 000 rows; first-element latency logged
  - QueryJsonStreamAsync_VeryLargeRows_HandlesChunkBoundaries — 50 rows × 2500-char Description
    forcing SQL Server to split a single JSON object across multiple FOR JSON rows

Benchmark (MsSqlBenchmarks.cs):
  - CosmoSQL  Warm  FOR JSON streamed (object-by-object)
  - CosmoSQL  Warm  FOR JSON buffered (full array)
  - ADO.NET   Warm  FOR JSON buffered (full array)
  Controlled by BENCH_JSON_QUERY env var (default: SELECT TOP 10000 … FOR JSON PATH)

All 10/10 integration tests pass; 5 000-row stream: first element in 5.3 ms, all in 17.6 ms.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: vkuttyp

Benchmarks/CosmoSQLBenchmarks/MsSqlBenchmarks.cs

@@ -184,6 +184,49 @@ public async Task Ado_Warm_ToJson()
               .ToList());
     }
 
+    // ── FOR JSON streaming vs buffered ──────────────────────────────────────────
+
+    private static readonly string JsonQuery =
+        Env("BENCH_JSON_QUERY",
+            "SELECT TOP 10000 AccountNo, AccountName FROM Accounts FOR JSON PATH");
+
+    [Benchmark(Description = "CosmoSQL  Warm  FOR JSON streamed (object-by-object)")]
+    public async Task Cosmo_Warm_ForJson_Streamed()
+    {
+        // Process each JSON object as it arrives — never holds the full array in RAM.
+        int count = 0;
+        await foreach (var elem in _cosmoConn!.QueryJsonStreamAsync(JsonQuery))
+            count++;
+        _ = count;
+    }
+
+    [Benchmark(Description = "CosmoSQL  Warm  FOR JSON buffered (full array)")]
+    public async Task Cosmo_Warm_ForJson_Buffered()
+    {
+        // Traditional approach: collect all FOR JSON rows, concat, then parse.
+        var rows = await _cosmoConn!.QueryAsync(JsonQuery);
+        var sb   = new System.Text.StringBuilder();
+        foreach (var row in rows)
+            sb.Append(row[0].AsString());
+        using var doc  = System.Text.Json.JsonDocument.Parse(sb.ToString());
+        var       root = doc.RootElement;
+        _ = root.GetArrayLength();
+    }
+
+    [Benchmark(Description = "ADO.NET   Warm  FOR JSON buffered (full array)")]
+    public async Task Ado_Warm_ForJson_Buffered()
+    {
+        // ADO.NET has no streaming FOR JSON — must buffer everything.
+        var sb = new System.Text.StringBuilder();
+        await using var cmd = new SqlCommand(JsonQuery, _adoConn);
+        await using var rdr = await cmd.ExecuteReaderAsync();
+        while (await rdr.ReadAsync())
+            sb.Append(rdr.GetString(0));
+        using var doc  = System.Text.Json.JsonDocument.Parse(sb.ToString());
+        var       root = doc.RootElement;
+        _ = root.GetArrayLength();
+    }
+
     // ── helpers ─────────────────────────────────────────────────────────────────
 
     private static string Env(string key, string fallback)

src/SqlDotnetty.MsSql/JsonChunkAssembler.cs

@@ -48,45 +48,66 @@ public static async IAsyncEnumerable<JsonElement> AssembleJsonObjectsAsync(
         IAsyncEnumerable<string> chunks,
         [EnumeratorCancellation] CancellationToken ct = default)
     {
-        // Growing byte buffer for all received chunk bytes.
         using var buf = new MemoryStream(32768);
 
-        var  state    = default(JsonReaderState);
-        long parsedTo = 0;   // absolute byte offset consumed so far
-        int  depth    = 0;   // JSON nesting depth (1 = inside the top-level array)
-        long elemStart = -1; // absolute offset of the current element's opening byte
-        bool inArray  = false;
+        var  state     = default(JsonReaderState);
+        long parsedTo  = 0;   // absolute position in rawBuf where reader stopped
+        int  depth     = 0;
+        long elemStart = -1;  // absolute position of current in-progress element's '{', or -1
+        bool inArray   = false;
 
         await foreach (var chunk in chunks.WithCancellation(ct))
         {
             if (string.IsNullOrEmpty(chunk)) continue;
 
-            // Append chunk bytes.
+            // Append chunk bytes to the rolling buffer.
             buf.Position = buf.Length;
             var chunkBytes = Encoding.UTF8.GetBytes(chunk);
             buf.Write(chunkBytes, 0, chunkBytes.Length);
 
-            // Process the buffer SYNCHRONOUSLY (Utf8JsonReader is a ref struct and
-            // cannot cross yield/await boundaries — so we collect results then yield).
-            var result = ProcessBuffer(
-                buf.GetBuffer(), (int)parsedTo, (int)buf.Length, state,
-                depth, elemStart, inArray);
-
-            // Update mutable state from synchronous result.
-            state     = result.State;
-            depth     = result.Depth;
-            inArray   = result.InArray;
-            elemStart = result.ElemStart;
-
-            // Compact: if elements were yielded, discard consumed prefix bytes.
-            if (result.ConsumedAbsolute > parsedTo)
-                CompactBuffer(buf, result.ConsumedAbsolute, out parsedTo);
-            else
-                parsedTo += result.BytesConsumedFromOffset;
-
-            // Yield complete elements OUTSIDE the Utf8JsonReader scope.
-            foreach (var elem in result.Elements)
-                yield return elem;
+            // Inner loop: drain as many complete elements as possible from the current
+            // buffer (one ProcessBuffer call may not find all of them if the result
+            // set is large, but more commonly runs 1-2 times).
+            bool madeProgress;
+            do
+            {
+                madeProgress = false;
+
+                // ProcessBuffer is synchronous (Utf8JsonReader is a ref struct and cannot
+                // cross yield/await boundaries).
+                var result = ProcessBuffer(
+                    buf.GetBuffer(), (int)parsedTo, (int)buf.Length,
+                    state, depth, elemStart, inArray);
+
+                state    = result.State;
+                depth    = result.Depth;
+                inArray  = result.InArray;
+
+                long newParsedTo  = parsedTo + result.BytesConsumedFromOffset;
+                long consumedTo   = result.ConsumedTo;   // absolute; safe to discard
+                long newElemStart = result.ElemStart;    // absolute, or -1
+
+                if (result.BytesConsumedFromOffset > 0 || result.Elements.Count > 0)
+                    madeProgress = true;
+
+                // Compact: discard bytes [0..consumedTo). This is always safe because
+                // consumedTo only advances after a complete element (elemStart resets
+                // to -1 at that point), so newElemStart is either -1 or >= consumedTo.
+                if (consumedTo > 0)
+                {
+                    CompactBuffer(buf, consumedTo);
+                    newParsedTo  -= consumedTo;
+                    if (newElemStart >= 0) newElemStart -= consumedTo;
+                }
+
+                parsedTo  = newParsedTo;
+                elemStart = newElemStart;
+
+                // Yield complete elements OUTSIDE the Utf8JsonReader scope.
+                foreach (var elem in result.Elements)
+                    yield return elem;
+
+            } while (madeProgress);
         }
     }
 
@@ -104,26 +125,35 @@ public static async Task<string> BufferJsonAsync(
         return sb.ToString();
     }
 
-    // ── Synchronous processing (no async/yield — safe for ref struct) ─────────
+    // ── Synchronous inner processor (no async/yield — safe for ref struct) ───
 
     private readonly ref struct BufferResult
     {
-        public IReadOnlyList<JsonElement> Elements            { get; init; }
-        public JsonReaderState             State              { get; init; }
-        public int                         Depth              { get; init; }
-        public long                        ElemStart          { get; init; }
-        public bool                        InArray            { get; init; }
-        /// <summary>Bytes consumed from the start of the slice passed in.</summary>
+        public IReadOnlyList<JsonElement> Elements               { get; init; }
+        public JsonReaderState             State                  { get; init; }
+        public int                         Depth                  { get; init; }
+        public long                        ElemStart              { get; init; }
+        public bool                        InArray                { get; init; }
+        /// <summary>Bytes the reader consumed from parsedToInt (may be 0 if all data is partial).</summary>
         public long                        BytesConsumedFromOffset { get; init; }
-        /// <summary>Absolute offset to compact to (= original parsedTo + BytesConsumedFromOffset).</summary>
-        public long                        ConsumedAbsolute   { get; init; }
+        /// <summary>
+        /// Absolute position through which bytes are safe to discard (= after the last
+        /// complete element extracted this call, or 0 if no element was completed).
+        /// </summary>
+        public long                        ConsumedTo             { get; init; }
     }
 
+    /// <summary>
+    /// Parse as many complete top-level JSON array elements as possible from
+    /// <c>rawBuf[parsedToInt..bufLength]</c> using <see cref="Utf8JsonReader"/> in
+    /// incremental mode.  Never breaks early — processes ALL complete elements.
+    /// </summary>
     private static BufferResult ProcessBuffer(
         byte[] rawBuf, int parsedToInt, int bufLength,
         JsonReaderState state, int depth, long elemStart, bool inArray)
     {
-        var elements = new List<JsonElement>();
+        var elements  = new List<JsonElement>();
+        long consumedTo = 0;   // absolute; nothing safe to discard yet
 
         var available = bufLength - parsedToInt;
         if (available <= 0)
@@ -132,20 +162,18 @@ private static BufferResult ProcessBuffer(
             {
                 Elements = elements, State = state, Depth = depth,
                 ElemStart = elemStart, InArray = inArray,
-                BytesConsumedFromOffset = 0, ConsumedAbsolute = parsedToInt,
+                BytesConsumedFromOffset = 0, ConsumedTo = 0,
             };
         }
 
         var span   = new ReadOnlySpan<byte>(rawBuf, parsedToInt, available);
         var reader = new Utf8JsonReader(span, isFinalBlock: false, state);
 
-        long consumedAbsolute = parsedToInt;
-        bool needBreak        = false;
-
-        while (!needBreak && reader.Read())
+        while (reader.Read())
         {
             var tt = reader.TokenType;
 
+            // Detect the outer array wrapper (first token of a FOR JSON result).
             if (tt == JsonTokenType.StartArray && !inArray && depth == 0)
             {
                 inArray = true;
@@ -157,6 +185,7 @@ private static BufferResult ProcessBuffer(
 
             if (tt is JsonTokenType.StartObject or JsonTokenType.StartArray)
             {
+                // Mark where a new top-level element begins (only at array depth 1).
                 if (depth == 1)
                     elemStart = parsedToInt + reader.TokenStartIndex;
                 depth++;
@@ -167,20 +196,18 @@ private static BufferResult ProcessBuffer(
 
                 if (depth == 1 && elemStart >= 0)
                 {
-                    // Complete element — extract and parse bytes.
-                    long end = parsedToInt + reader.BytesConsumed;
-                    int  len = (int)(end - elemStart);
+                    // A complete top-level element just finished.
+                    long end  = parsedToInt + reader.BytesConsumed;
+                    int  len  = (int)(end - elemStart);
                     var  elemBytes = new byte[len];
                     Array.Copy(rawBuf, (int)elemStart, elemBytes, 0, len);
 
                     using var doc = JsonDocument.Parse(elemBytes);
                     elements.Add(doc.RootElement.Clone());
 
-                    elemStart        = -1;
-                    consumedAbsolute = end;
-
-                    // Signal that the outer loop should compact and break after this.
-                    needBreak = true;
+                    elemStart  = -1;
+                    consumedTo = end;   // bytes up to here are now safe to discard
+                    // ── DO NOT break — continue to collect further elements ──
                 }
                 else if (depth == 0)
                 {
@@ -189,55 +216,28 @@ private static BufferResult ProcessBuffer(
             }
         }
 
-        long bytesConsumed = parsedToInt + reader.BytesConsumed - parsedToInt;
-
         return new BufferResult
         {
-            Elements = elements,
-            State    = reader.CurrentState,
-            Depth    = depth,
-            ElemStart = elemStart,
-            InArray  = inArray,
+            Elements                = elements,
+            State                   = reader.CurrentState,
+            Depth                   = depth,
+            ElemStart               = elemStart,
+            InArray                 = inArray,
             BytesConsumedFromOffset = reader.BytesConsumed,
-            ConsumedAbsolute        = needBreak ? consumedAbsolute : parsedToInt + reader.BytesConsumed,
+            ConsumedTo              = consumedTo,
         };
     }
 
     /// <summary>
-    /// Discard fully-consumed prefix bytes from <paramref name="buf"/> to cap memory usage.
+    /// Discard the first <paramref name="consumedAbsolute"/> bytes from the buffer by
+    /// block-copying the remaining bytes to the front.  Keeps memory usage bounded.
     /// </summary>
-    private static void CompactBuffer(MemoryStream buf, long consumedAbsolute, out long newParsedTo)
+    private static void CompactBuffer(MemoryStream buf, long consumedAbsolute)
     {
         long remaining = buf.Length - consumedAbsolute;
         var  rawBuf    = buf.GetBuffer();
         if (remaining > 0)
             Buffer.BlockCopy(rawBuf, (int)consumedAbsolute, rawBuf, 0, (int)remaining);
         buf.SetLength(remaining);
-        newParsedTo = 0;
     }
 }
-
-
-/// <summary>
-/// Assembles fragmented JSON chunks (as produced by SQL Server's <c>FOR JSON</c> queries)
-/// into complete, independently-parsed <see cref="JsonElement"/> objects.
-/// </summary>
-/// <remarks>
-/// <para>
-/// SQL Server's <c>FOR JSON PATH</c> / <c>FOR JSON AUTO</c> does not respect JSON value
-/// boundaries when splitting output across rows — a single row may contain half a property
-/// name, and the next row may start mid-value.  Naively streaming row-by-row therefore
-/// yields unusable partial JSON.
-/// </para>
-/// <para>
-/// This class solves that by feeding chunks through .NET's <see cref="Utf8JsonReader"/>
-/// with <see cref="JsonReaderState"/> persistence.  The reader correctly handles escaped
-/// strings, nested objects and arrays, and resumes from exactly the right position when
-/// more data arrives — so <c>{</c> characters inside string values are never misread as
-/// object boundaries.
-/// </para>
-/// <para>
-/// For a <c>FOR JSON</c> result like <c>[{"Id":1},{"Id":2},…]</c>, each top-level array
-/// element is yielded as a separate <see cref="JsonElement"/> the moment the closing <c>}</c>
-/// of that element is received, without waiting for the entire array to arrive.
-/// </para>