address asserts and refine TryReadPlpBytes

Author: Wraith2

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

@@ -4413,7 +4413,12 @@ private TdsOperationStatus TryResetBlobState()
 #if DEBUG
             else
             {
-                Debug.Assert((_sharedState._columnDataBytesRemaining == 0 || _sharedState._columnDataBytesRemaining == -1) && _stateObj._longlen == 0, "Haven't read header yet, but column is partially read?");
+                Debug.Assert(
+                    (_sharedState._columnDataBytesRemaining == 0 || _sharedState._columnDataBytesRemaining == -1)
+                    &&
+                    (_stateObj._longlen == 0 || _stateObj.IsSnapshotContinuing()),
+                    "Haven't read header yet, but column is partially read?"
+                );
             }
 #endif
 

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

@@ -1860,7 +1860,8 @@ internal int ReadPlpBytesChunk(byte[] buff, int offset, int len)
 
         internal TdsOperationStatus TryReadPlpBytes(ref byte[] buff, int offset, int len, out int totalBytesRead)
         {
-            return TryReadPlpBytes(ref buff, offset, len, out totalBytesRead, IsSnapshotAvailable());
+            ( _, bool isStarting, bool isContinuing) = GetSnapshotStatuses();
+            return TryReadPlpBytes(ref buff, offset, len, out totalBytesRead, isStarting || isContinuing);
         }
         // Reads the requested number of bytes from a plp data stream, or the entire data if
         // requested length is -1 or larger than the actual length of data. First call to this method
@@ -1904,6 +1905,11 @@ internal TdsOperationStatus TryReadPlpBytes(ref byte[] buff, int offset, int len
                         totalBytesRead = offset;
                     }
                 }
+                else if (_snapshot != null)
+                {
+                    // legacy replay path perf optimization
+                    buff = (byte[])TryTakeSnapshotStorage();
+                }
 
                 if ((ulong)(buff?.Length ?? 0) != _longlen)
                 {
@@ -1935,7 +1941,6 @@ internal TdsOperationStatus TryReadPlpBytes(ref byte[] buff, int offset, int len
 
             while (bytesLeft > 0)
             {
-                bool stored = false;
                 int bytesToRead = (int)Math.Min(_longlenleft, (ulong)bytesLeft);
                 if (buff.Length < (offset + bytesToRead))
                 {
@@ -1961,20 +1966,15 @@ internal TdsOperationStatus TryReadPlpBytes(ref byte[] buff, int offset, int len
                         // a partial read has happened so store the target buffer in the snapshot
                         // so it can be re-used when another packet arrives and we read again
                         SetSnapshotStorage(buff);
+                        SetSnapshotDataSize(bytesRead);
+                        
                     }
-                    return result;
-                }
-                else
-                {
-                    if (writeDataSizeToSnapshot)
+                    else if (_snapshot != null)
                     {
+                        // legacy replay path perf optimization
                         SetSnapshotStorage(buff);
-                        if (_snapshot.ContinueEnabled)
-                        {
-                            SetSnapshotDataSize(bytesRead);
-                            stored = true;
-                        }
                     }
+                    return result;
                 }
 
                 if (_longlenleft == 0)
@@ -1983,9 +1983,18 @@ internal TdsOperationStatus TryReadPlpBytes(ref byte[] buff, int offset, int len
                     result = TryReadPlpLength(false, out _);
                     if (result != TdsOperationStatus.Done)
                     {
-                        if (result == TdsOperationStatus.NeedMoreData && writeDataSizeToSnapshot && !stored)
+                        if (result == TdsOperationStatus.NeedMoreData)
                         {
-                            SetSnapshotDataSize(bytesRead);
+                            if (writeDataSizeToSnapshot)
+                            {
+                                SetSnapshotStorage(buff);
+                                SetSnapshotDataSize(bytesRead);
+                            }
+                            else if (_snapshot != null)
+                            {
+                                // legacy replay path perf optimization
+                                SetSnapshotStorage(buff);
+                            }
                         }
                         return result;
                     }
@@ -2767,7 +2776,7 @@ internal int GetSnapshotStorageLength<T>()
 
         internal object TryTakeSnapshotStorage()
         {
-            Debug.Assert(_snapshot != null && _snapshot.ContinueEnabled, "should not access snapshot accessor functions without first checking that the snapshot is present");
+            Debug.Assert(_snapshot != null, "should not access snapshot accessor functions without first checking that the snapshot is present");
             object buffer = null;
             if (_snapshot != null)
             {
@@ -2779,7 +2788,7 @@ internal object TryTakeSnapshotStorage()
 
         internal void SetSnapshotStorage(object buffer)
         {
-            Debug.Assert(_snapshot != null && _snapshot.ContinueEnabled, "should not access snapshot accessor functions without first checking that the snapshot is available");
+            Debug.Assert(_snapshot != null, "should not access snapshot accessor functions without first checking that the snapshot is available");
             Debug.Assert(_snapshot._storage == null, "should not overwrite snapshot stored buffer");
             if (_snapshot != null)
             {