Perf | Reduce async read memory allocations by recycling StateSnapshot packet nodes

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs

@@ -4653,8 +4653,10 @@ internal sealed partial class StateSnapshot
         {
             private sealed partial class PacketData
             {
-                public readonly byte[] Buffer;
-                public readonly int Read;
+                // These are mutable so that instances can be recycled through the snapshot's
+                // free list instead of allocating a new node for every packet that is appended.
+                public byte[] Buffer;
+                public int Read;
 
                 public PacketData NextPacket;
                 public PacketData PrevPacket;
@@ -4666,9 +4668,38 @@ private sealed partial class PacketData
                 public int RunningDataSize;
 
                 public PacketData(byte[] buffer, int read)
+                {
+                    // Delegate so that a freshly allocated node and a recycled one are initialized
+                    // by exactly one piece of code. Keeping these paths separate is how the
+                    // RunningDataSize reset was missed previously.
+                    Initialize(buffer, read);
+                }
+
+                /// <summary>
+                /// Re-initializes an instance taken from the free list so that it describes a new packet.
+                /// </summary>
+                public void Initialize(byte[] buffer, int read)
                 {
                     Buffer = buffer;
                     Read = read;
+                    NextPacket = null;
+                    PrevPacket = null;
+                    RunningDataSize = 0;
+                    ResetDebugState();
+                }
+
+                /// <summary>
+                /// Releases all references held by this instance before it is placed on the free list
+                /// so that recycling a node cannot keep a packet buffer alive.
+                /// </summary>
+                public void Reset()
+                {
+                    Buffer = null;
+                    Read = 0;
+                    NextPacket = null;
+                    PrevPacket = null;
+                    RunningDataSize = 0;
+                    ResetDebugState();
                 }
 
                 public int PacketID => Packet.GetIDFromHeader(Buffer.AsSpan(0, TdsEnums.HEADER_LEN));
@@ -4698,10 +4729,13 @@ internal int GetPacketDataSize()
 
                 internal void CheckDebugDataHash() => CheckDebugDataHashImpl();
 
+                internal void ResetDebugState() => ResetDebugStateImpl();
+
                 partial void SetDebugStackImpl(string value);
                 partial void SetDebugPacketIdImpl(int value);
                 partial void SetDebugDataHashImpl();
                 partial void CheckDebugDataHashImpl();
+                partial void ResetDebugStateImpl();
             }
 
 #if DEBUG
@@ -4845,6 +4879,13 @@ public string Status
 
                 partial void SetDebugPacketIdImpl(int value) => DebugPacketId = value;
 
+                partial void ResetDebugStateImpl()
+                {
+                    DebugPacketId = 0;
+                    Stack = null;
+                    Hash = null;
+                }
+
                 partial void SetDebugDataHashImpl()
                 {
                     if (Buffer != null)
@@ -4990,6 +5031,25 @@ internal void Restore(TdsParserStateObject stateObj)
             private PacketData _current;
             private PacketData _continuePacket;
 
+            /// <summary>
+            /// Head of a free list of <see cref="PacketData"/> instances, linked through
+            /// <see cref="PacketData.NextPacket"/>. A snapshot is taken and released for every async
+            /// continuation, so recycling these nodes avoids allocating one per packet per replay.
+            /// The list is bounded so that an unusually long replay cannot retain nodes indefinitely.
+            /// </summary>
+            private PacketData _sparePackets;
+            private int _sparePacketCount;
+
+            /// <summary>
+            /// Upper bound on the free list. The value is empirical rather than derived: most
+            /// snapshots hold a single packet, so even a one entry list removes the bulk of the
+            /// allocations, and a larger bound only helps reads that span several packets such as
+            /// large rows or XML and binary columns. 16 was the point past which the read
+            /// benchmarks stopped improving. The cost is bounded at roughly 900 bytes retained per
+            /// state object, and the parked nodes hold no buffer references.
+            /// </summary>
+            private const int MaxSparePacketCount = 16;
+
 #if DEBUG
             private int _packetCounter;
             private int _rollingPend = 0;
@@ -5069,7 +5129,7 @@ internal void AppendPacketData(byte[] buffer, int read)
                     }
                 }
 #endif
-                PacketData packetData = new PacketData(buffer, read);
+                PacketData packetData = RentPacket(buffer, read);
 #if DEBUG
                 packetData.SetDebugStack(_stateObj._lastStack);
                 packetData.SetDebugPacketId(Interlocked.Increment(ref _packetCounter));
@@ -5145,9 +5205,11 @@ internal bool MoveToContinue()
 
             internal void CaptureAsStart(TdsParserStateObject stateObj)
             {
-                _firstPacket = null;
-                _lastPacket = null;
-                _current = null;
+                // Note this also clears _continuePacket, which the assignments this replaced did
+                // not. Every path into here runs Clear() first so the chain is already empty in
+                // practice, but routing through ClearPackets keeps the invariant that no field can
+                // still reference a node once one has been parked on the free list.
+                ClearPackets();
 
                 _stateObj = stateObj;
                 _replayStateData ??= new StateObjectData();
@@ -5257,10 +5319,64 @@ internal void Clear()
 
             private void ClearPackets()
             {
+                PacketData current = _firstPacket;
+
                 _firstPacket = null;
                 _lastPacket = null;
                 _continuePacket = null;
                 _current = null;
+
+                // Every field that can reference a node is cleared above before any node is parked,
+                // so nothing on the free list can also be reachable from the live chain.
+                while (current != null)
+                {
+                    PacketData next = current.NextPacket;
+                    if (!ReturnPacket(current))
+                    {
+                        break;
+                    }
+                    current = next;
+                }
+            }
+
+            /// <summary>
+            /// Places a <see cref="PacketData"/> on the free list when there is room for it, clearing
+            /// every field first so a recycled node cannot carry stale state, or a reference to a
+            /// packet buffer, into its next use. Returns false once the list is full, in which case
+            /// the node is left to the garbage collector.
+            /// </summary>
+            private bool ReturnPacket(PacketData packet)
+            {
+                if (_sparePacketCount >= MaxSparePacketCount)
+                {
+                    return false;
+                }
+
+                // Reset must run before the node is linked in: NextPacket is reused as the free
+                // list link, so clearing it afterwards would truncate the list.
+                packet.Reset();
+                packet.NextPacket = _sparePackets;
+                _sparePackets = packet;
+                _sparePacketCount++;
+                return true;
+            }
+
+            /// <summary>
+            /// Takes a <see cref="PacketData"/> from the free list, or allocates one when the list is empty.
+            /// </summary>
+            private PacketData RentPacket(byte[] buffer, int read)
+            {
+                PacketData packetData = _sparePackets;
+                if (packetData == null)
+                {
+                    return new PacketData(buffer, read);
+                }
+
+                // Read the link before Initialize clears NextPacket.
+                _sparePackets = packetData.NextPacket;
+                _sparePacketCount--;
+                packetData.Initialize(buffer, read);
+                return packetData;
             }
 
             private void ClearState()