Add async snapshot continue capability for multipacket fields (#3161)

Author: Wraith2

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

@@ -3766,7 +3766,7 @@ private TdsOperationStatus TryNextResult(out bool more)
                             if (result != TdsOperationStatus.Done)
                             {
                                 more = false;
-                                return TdsOperationStatus.Done;
+                                return result;
                             }
 
                             // In the case of not closing the reader, null out the metadata AFTER
@@ -4525,7 +4525,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/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -4524,7 +4524,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/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs

@@ -21,6 +21,7 @@ private enum Tristate : byte
         internal const string UseMinimumLoginTimeoutString = @"Switch.Microsoft.Data.SqlClient.UseOneSecFloorInTimeoutCalculationDuringLogin";
         internal const string LegacyVarTimeZeroScaleBehaviourString = @"Switch.Microsoft.Data.SqlClient.LegacyVarTimeZeroScaleBehaviour";
         internal const string UseCompatibilityProcessSniString = @"Switch.Microsoft.Data.SqlClient.UseCompatibilityProcessSni";
+        internal const string UseCompatibilityAsyncBehaviourString = @"Switch.Microsoft.Data.SqlClient.UseCompatibilityAsyncBehaviour";
 
         // this field is accessed through reflection in tests and should not be renamed or have the type changed without refactoring NullRow related tests
         private static Tristate s_legacyRowVersionNullBehavior;
@@ -30,6 +31,7 @@ private enum Tristate : byte
         // this field is accessed through reflection in Microsoft.Data.SqlClient.Tests.SqlParameterTests and should not be renamed or have the type changed without refactoring related tests
         private static Tristate s_legacyVarTimeZeroScaleBehaviour;
         private static Tristate s_useCompatProcessSni;
+        private static Tristate s_useCompatAsyncBehaviour;
 
 #if NET
         static LocalAppContextSwitches()
@@ -85,6 +87,12 @@ public static bool DisableTNIRByDefault
             }
         }
 #endif
+        /// <summary>
+        /// In TdsParser the ProcessSni function changed significantly when the packet
+        /// multiplexing code needed for high speed multi-packet column values was added.
+        /// In case of compatibility problems this switch will change TdsParser to use
+        /// the previous version of the function.
+        /// </summary>
         public static bool UseCompatibilityProcessSni
         {
             get
@@ -104,6 +112,42 @@ public static bool UseCompatibilityProcessSni
             }
         }
 
+        /// <summary>
+        /// In TdsParser the async multi-packet column value fetch behaviour is capable of
+        /// using a continue snapshot state in addition to the original replay from start
+        /// logic.
+        /// This switch disables use of the continue snapshot state. This switch will always
+        /// return true if <see cref="UseCompatibilityProcessSni"/> is enabled because the 
+        /// continue state is not stable without the multiplexer.
+        /// </summary>
+        public static bool UseCompatibilityAsyncBehaviour
+        {
+            get
+            {
+                if (UseCompatibilityProcessSni)
+                {
+                    // If ProcessSni compatibility mode has been enabled then the packet
+                    // multiplexer has been disabled. The new async behaviour using continue
+                    // point capture is only stable if the multiplexer is enabled so we must
+                    // return true to enable compatibility async behaviour using only restarts.
+                    return true;
+                }
+
+                if (s_useCompatAsyncBehaviour == Tristate.NotInitialized)
+                {
+                    if (AppContext.TryGetSwitch(UseCompatibilityAsyncBehaviourString, out bool returnedValue) && returnedValue)
+                    {
+                        s_useCompatAsyncBehaviour = Tristate.True;
+                    }
+                    else
+                    {
+                        s_useCompatAsyncBehaviour = Tristate.False;
+                    }
+                }
+                return s_useCompatAsyncBehaviour == Tristate.True;
+            }
+        }
+
         /// <summary>
         /// When using Encrypt=false in the connection string, a security warning is output to the console if the TLS version is 1.2 or lower.
         /// This warning can be suppressed by enabling this AppContext switch.

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

@@ -37,18 +37,34 @@ private SqlCachedBuffer(List<byte[]> cachedBytes)
         /// </summary>
         internal static TdsOperationStatus TryCreate(SqlMetaDataPriv metadata, TdsParser parser, TdsParserStateObject stateObj, out SqlCachedBuffer buffer)
         {
-            byte[] byteArr;
-
-            List<byte[]> cachedBytes = new();
             buffer = null;
 
+            (bool isAvailable, bool isStarting, bool isContinuing) = stateObj.GetSnapshotStatuses();
+
+            List<byte[]> cachedBytes = null;
+            if (isAvailable)
+            {
+                cachedBytes = stateObj.TryTakeSnapshotStorage() as List<byte[]>;
+                if (cachedBytes != null && !isStarting && !isContinuing) 
+                {
+                    stateObj.SetSnapshotStorage(null);
+                }
+            }
+ 
+            if (cachedBytes == null)
+            {
+                cachedBytes = new List<byte[]>();
+            }
+
+
             // the very first length is already read.
             TdsOperationStatus result = parser.TryPlpBytesLeft(stateObj, out ulong plplength);
             if (result != TdsOperationStatus.Done)
             {
                 return result;
             }
 
+
             // For now we  only handle Plp data from the parser directly.
             Debug.Assert(metadata.metaType.IsPlp, "SqlCachedBuffer call on a non-plp data");
             do
@@ -59,13 +75,25 @@ internal static TdsOperationStatus TryCreate(SqlMetaDataPriv metadata, TdsParser
                 }
                 do
                 {
+                    bool returnAfterAdd = false;
                     int cb = (plplength > (ulong)MaxChunkSize) ? MaxChunkSize : (int)plplength;
-                    byteArr = new byte[cb];
-                    result = stateObj.TryReadPlpBytes(ref byteArr, 0, cb, out cb);
+                    byte[] byteArr = new byte[cb];
+                    // pass false for the writeDataSizeToSnapshot parameter because we want to only take data
+                    // from the current packet and not try to do a continue-capable multi packet read
+                    result = stateObj.TryReadPlpBytes(ref byteArr, 0, cb, out cb, writeDataSizeToSnapshot: false, compatibilityMode: false);
                     if (result != TdsOperationStatus.Done)
                     {
-                        return result;
+                        if (result == TdsOperationStatus.NeedMoreData && isAvailable && cb == byteArr.Length)
+                        {
+                            // succeeded in getting the data but failed to find the next plp length
+                            returnAfterAdd = true;
+                        }
+                        else
+                        {
+                            return result;
+                        }
                     }
+
                     Debug.Assert(cb == byteArr.Length);
                     if (cachedBytes.Count == 0)
                     {
@@ -74,6 +102,16 @@ internal static TdsOperationStatus TryCreate(SqlMetaDataPriv metadata, TdsParser
                     }
                     cachedBytes.Add(byteArr);
                     plplength -= (ulong)cb;
+
+                    if (returnAfterAdd)
+                    {
+                        if (isStarting || isContinuing)
+                        {
+                            stateObj.SetSnapshotStorage(cachedBytes);
+                        }
+                        return result;
+                    }
+
                 } while (plplength > 0);
 
                 result = parser.TryPlpBytesLeft(stateObj, out plplength);