Async multi packet fixes

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -13055,6 +13055,8 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
             char[] temp = null;
             bool buffIsRented = false;
             int startOffset = 0;
+
+            stateObj.RequestContinue(true);
             (bool canContinue, bool isStarting, bool isContinuing) = stateObj.GetSnapshotStatuses();
 
             if (canContinue)
@@ -13149,7 +13151,7 @@ bool writeDataSizeToSnapshot
             if (stateObj._longlen == 0)
             {
                 Debug.Assert(stateObj._longlenleft == 0);
-                totalCharsRead = 0;
+                totalCharsRead = startOffsetByteCount / 2;
                 return TdsOperationStatus.Done;       // No data
             }
 
@@ -13169,14 +13171,15 @@ bool writeDataSizeToSnapshot
             //  later needing to repeatedly allocate new target buffers and copy data as we discover new data
             if (buff == null && stateObj._longlen != TdsEnums.SQL_PLP_UNKNOWNLEN && stateObj._longlen < (int.MaxValue >> 1))
             {
-                if (supportRentedBuff && stateObj._longlen < 1073741824) // 1 Gib
+                int stateLen = (int)stateObj._longlen >> 1;
+                if (supportRentedBuff && stateLen < 1073741824) // 1 Gib
                 {
-                    buff = ArrayPool<char>.Shared.Rent((int)Math.Min((int)stateObj._longlen, len));
+                    buff = ArrayPool<char>.Shared.Rent(Math.Min(stateLen, len));
                     rentedBuff = true;
                 }
                 else
                 {
-                    buff = new char[(int)Math.Min((int)stateObj._longlen, len)];
+                    buff = new char[Math.Min(stateLen, len)];
                     rentedBuff = false;
                 }
             }

src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -13312,7 +13312,7 @@ bool writeDataSizeToSnapshot
             if (stateObj._longlen == 0)
             {
                 Debug.Assert(stateObj._longlenleft == 0);
-                totalCharsRead = 0;
+                totalCharsRead = startOffsetByteCount / 2;
                 return TdsOperationStatus.Done;       // No data
             }
 
@@ -13327,19 +13327,20 @@ bool writeDataSizeToSnapshot
             );
             charsLeft = len;
 
-            // If total length is known up front, the length isn't specified as unknown 
-            // and the caller doesn't pass int.max/2 indicating that it doesn't know the length
-            // allocate the whole buffer in one shot instead of realloc'ing and copying over each time
+            // If total data length is known up front from the plp header by being not SQL_PLP_UNKNOWNLEN
+            //  and the number of chars required is less than int.max/2 allocate the entire buffer now to avoid
+            //  later needing to repeatedly allocate new target buffers and copy data as we discover new data
             if (buff == null && stateObj._longlen != TdsEnums.SQL_PLP_UNKNOWNLEN && stateObj._longlen < (int.MaxValue >> 1))
             {
-                if (supportRentedBuff && stateObj._longlen < 1073741824) // 1 Gib
+                int stateLen = (int)stateObj._longlen >> 1;
+                if (supportRentedBuff && stateLen < 1073741824) // 1 Gib
                 {
-                    buff = ArrayPool<char>.Shared.Rent((int)Math.Min((int)stateObj._longlen, len));
+                    buff = ArrayPool<char>.Shared.Rent(Math.Min(stateLen, len));
                     rentedBuff = true;
                 }
                 else
                 {
-                    buff = new char[(int)Math.Min((int)stateObj._longlen, len)];
+                    buff = new char[Math.Min(stateLen, len)];
                     rentedBuff = false;
                 }
             }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCachedBuffer.cs

@@ -39,10 +39,10 @@ internal static TdsOperationStatus TryCreate(SqlMetaDataPriv metadata, TdsParser
         {
             buffer = null;
 
-            (bool canContinue, bool isStarting, _) = stateObj.GetSnapshotStatuses();
+            (bool canContinue, bool isStarting, bool isContinuing) = stateObj.GetSnapshotStatuses();
 
             List<byte[]> cachedBytes = null;
-            if (canContinue)
+            if (canContinue && isContinuing)
             {
                 cachedBytes = stateObj.TryTakeSnapshotStorage() as List<byte[]>;
                 if (isStarting)
@@ -81,7 +81,7 @@ internal static TdsOperationStatus TryCreate(SqlMetaDataPriv metadata, TdsParser
                     result = stateObj.TryReadPlpBytes(ref byteArr, 0, cb, out cb, canContinue, writeDataSizeToSnapshot: false, compatibilityMode: false);
                     if (result != TdsOperationStatus.Done)
                     {
-                        if (result == TdsOperationStatus.NeedMoreData && canContinue && cb == byteArr.Length)
+                        if (result == TdsOperationStatus.NeedMoreData && canContinue && cb == byteArr.Length && (isContinuing || !isStarting))
                         {
                             // succeeded in getting the data but failed to find the next plp length
                             returnAfterAdd = true;

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlDataReader.cs

@@ -4238,9 +4238,7 @@ private TdsOperationStatus TryResetBlobState()
             else
             {
                 Debug.Assert(
-                    (_sharedState._columnDataBytesRemaining == 0 || _sharedState._columnDataBytesRemaining == -1)
-                    &&
-                    (_stateObj._longlen == 0 || _stateObj.IsSnapshotContinuing()),
+                    (_sharedState._columnDataBytesRemaining == 0 || _sharedState._columnDataBytesRemaining == -1),
                     "Haven't read header yet, but column is partially read?"
                 );
             }
@@ -5395,6 +5393,10 @@ private static Task<T> GetFieldValueAsyncExecute<T>(Task task, object state)
                     {
                         return Task.FromResult<T>(reader.GetFieldValueFromSqlBufferInternal<T>(reader._data[columnIndex], reader._metaData[columnIndex], isAsync: true));
                     }
+                    else
+                    {
+                        return reader.ExecuteAsyncCall(context);
+                    }
                 }
             }
 

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

@@ -41,6 +41,10 @@ public void ProcessSniPacket(PacketHandle packet, uint error)
                 if (PartialPacketContainsCompletePacket())
                 {
                     Packet partialPacket = _partialPacket;
+                    // the partial packet can contain more than a single packet worth of data so to consume the 
+                    //  partial packet we must use the CurrentLength not just the RequiredLength and then later
+                    //  the multiplexer will split out the complete packet for consumption and maintain the
+                    //  additional data
                     SetBuffer(partialPacket.Buffer, 0, partialPacket.CurrentLength);
                     ClearPartialPacket();
                     getDataError = TdsEnums.SNI_SUCCESS;
@@ -50,7 +54,7 @@ public void ProcessSniPacket(PacketHandle packet, uint error)
                 {
                     if (_inBytesRead != 0)
                     {
-                        SetBuffer(new byte[_inBuff.Length], 0, 0);
+                        NewBuffer(_inBuff.Length);
                     }
                     getDataError = GetSniPacket(packet, ref dataSize);
                 }
@@ -76,7 +80,7 @@ public void ProcessSniPacket(PacketHandle packet, uint error)
                     {
                         if (recurse && appended)
                         {
-                            SetBuffer(new byte[_inBuff.Length], 0, 0);
+                            NewBuffer(_inBuff.Length);
                             appended = false;
                         }
                         MultiplexPackets(
@@ -95,16 +99,19 @@ out recurse
                         // if a partial packet was reconstructed it must be handled first
                         if (consumePartialPacket)
                         {
+                            // the partial packet has been processed by the multiplexer and should now have 
+                            //  only data from a single packet in it so we should use RequiredLength which
+                            //  is defined by the packet header here not CurrentLength
+                            Debug.Assert(PartialPacket != null && PartialPacket.RequiredLength == PartialPacket.CurrentLength);
                             if (_snapshot != null)
                             {
-                                _snapshot.AppendPacketData(PartialPacket.Buffer, PartialPacket.CurrentLength);
-                                SetBuffer(new byte[_inBuff.Length], 0, 0);
+                                _snapshot.AppendPacketData(PartialPacket.Buffer, PartialPacket.RequiredLength);
+                                NewBuffer(_inBuff.Length);
                                 appended = true;
                             }
                             else
                             {
-                                SetBuffer(PartialPacket.Buffer, 0, PartialPacket.CurrentLength);
-
+                                SetBuffer(PartialPacket.Buffer, 0, PartialPacket.RequiredLength);
                             }
                             bufferIsPartialCompleted = true;
                             ClearPartialPacket();
@@ -125,7 +132,7 @@ out recurse
                                 // if we SetBuffer here to clear the packet buffer we will break the attention handling which relies
                                 // on the attention containing packet remaining in the active buffer even if we're appending to the
                                 // snapshot so we will have to use the appended variable to prevent the same buffer being added again
-                                //// SetBuffer(new byte[_inBuff.Length], 0, 0);
+                                //// NewBuffer(_inBuff.Length);
                                 appended = true;
                             }
                             else
@@ -141,19 +148,28 @@ out recurse
                             // we don't process it
                             if (!bufferIsPartialCompleted)
                             {
-                                SetBuffer(_inBuff, 0, 0);
+                                NewBuffer(_inBuff.Length);
                             }
                         }
 
                         // if there is a remainder it must be last
                         if (remainderPacketProduced)
                         {
                             SetPartialPacket(remainderPacket);
+                            if (appended && recurse)
+                            {
+                                // if we've appended to the snapshot already we can't recurse and append to it again because the
+                                //   snapshot might be cleared by the async cleanup functions
+                                // the only way to get a recurse output from the multiplexer is if it has produced a remainder packet so
+                                //  assert that this is the case and the put the remainder packet in the partial packet so that it
+                                //  can be picked up in another call.
+                                recurse = false;
+                            }
                             if (!bufferIsPartialCompleted)
                             {
                                 // we are keeping the partial packet buffer so replace it with a new one
                                 // unless we have already set the buffer to the partial packet buffer
-                                SetBuffer(new byte[_inBuff.Length], 0, 0);
+                                NewBuffer(_inBuff.Length);
                             }
                         }
 
@@ -301,7 +317,7 @@ out bool recurse
                         remainderPacket = new Packet
                         {
                             Buffer = new byte[dataBuffer.Length],
-                            CurrentLength = remainderLength,
+                            CurrentLength = remainderLength
                         };
                         remainderPacket.SetCreatedBy(1);