Restrict the Long Running option to MARS (#27010)

* Restrict the Long Running option to MARS

* Code review comments

Author: saurabh500

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIHandle.cs

@@ -56,7 +56,7 @@ internal abstract class SNIHandle
         /// </summary>
         /// <param name="packet">SNI packet</param>
         /// <returns>SNI error code</returns>
-        public abstract uint ReceiveAsync(ref SNIPacket packet);
+        public abstract uint ReceiveAsync(ref SNIPacket packet, bool isMars = false);
 
         /// <summary>
         /// Enable SSL

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIMarsConnection.cs

@@ -108,7 +108,7 @@ public uint ReceiveAsync(ref SNIPacket packet)
         {
             lock (this)
             {
-                return _lowerHandle.ReceiveAsync(ref packet);
+                return _lowerHandle.ReceiveAsync(ref packet, isMars: true);
             }
         }
 

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIMarsHandle.cs

@@ -273,7 +273,7 @@ public override uint SendAsync(SNIPacket packet, SNIAsyncCallback callback = nul
         /// </summary>
         /// <param name="packet">SNI packet</param>
         /// <returns>SNI error code</returns>
-        public override uint ReceiveAsync(ref SNIPacket packet)
+        public override uint ReceiveAsync(ref SNIPacket packet, bool isMars = true)
         {
             lock (_receivedPacketQueue)
             {

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNINpHandle.cs

@@ -177,7 +177,7 @@ public override uint Receive(out SNIPacket packet, int timeout)
             }
         }
 
-        public override uint ReceiveAsync(ref SNIPacket packet)
+        public override uint ReceiveAsync(ref SNIPacket packet, bool isMars = false)
         {
             lock (this)
             {
@@ -186,7 +186,7 @@ public override uint ReceiveAsync(ref SNIPacket packet)
 
                 try
                 {
-                    packet.ReadFromStreamAsync(_stream, _receiveCallback);
+                    packet.ReadFromStreamAsync(_stream, _receiveCallback, isMars);
                     return TdsEnums.SNI_SUCCESS_IO_PENDING;
                 }
                 catch (ObjectDisposedException ode)
@@ -356,4 +356,4 @@ public override void KillConnection()
         }
 #endif
     }
-}
+}

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIPacket.cs

@@ -236,9 +236,18 @@ public void Reset()
         /// </summary>
         /// <param name="stream">Stream to read from</param>
         /// <param name="callback">Completion callback</param>
-        public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback)
+        public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback, bool isMars)
         {
             bool error = false;
+            TaskContinuationOptions options = TaskContinuationOptions.DenyChildAttach;
+            // MARS operations during Sync ADO.Net API calls are Sync over Async. Each API call can request 
+            // threads to execute the async reads. MARS operations do not get the threads quickly enough leading to timeout
+            // To fix the MARS thread exhaustion issue LongRunning continuation option is a temporary solution with its own drawbacks, 
+            // and should be removed after evaluating how to fix MARS threading issues efficiently
+            if (isMars)
+            {
+                options |= TaskContinuationOptions.LongRunning;
+            }
 
             stream.ReadAsync(_data, 0, _capacity).ContinueWith(t =>
             {
@@ -267,7 +276,7 @@ public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback)
                 callback(this, error ? TdsEnums.SNI_ERROR : TdsEnums.SNI_SUCCESS);
             },
             CancellationToken.None,
-            TaskContinuationOptions.DenyChildAttach | TaskContinuationOptions.LongRunning,
+            options,
             TaskScheduler.Default);
         }
 

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIProxy.cs

@@ -424,7 +424,7 @@ private SNINpHandle CreateNpHandle(DataSource details, long timerExpire, object
         /// <param name="handle">SNI handle</param>
         /// <param name="packet">Packet</param>
         /// <returns>SNI error status</returns>
-        public uint ReadAsync(SNIHandle handle, out SNIPacket packet)
+        public uint ReadAsync(SNIHandle handle, out SNIPacket packet, bool isMars = false)
         {
             packet = new SNIPacket(null);
 

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNITcpHandle.cs

@@ -577,7 +577,7 @@ public override uint SendAsync(SNIPacket packet, SNIAsyncCallback callback = nul
         /// </summary>
         /// <param name="packet">SNI packet</param>
         /// <returns>SNI error code</returns>
-        public override uint ReceiveAsync(ref SNIPacket packet)
+        public override uint ReceiveAsync(ref SNIPacket packet, bool isMars = false)
         {
             lock (this)
             {
@@ -586,7 +586,7 @@ public override uint ReceiveAsync(ref SNIPacket packet)
 
                 try
                 {
-                    packet.ReadFromStreamAsync(_stream, _receiveCallback);
+                    packet.ReadFromStreamAsync(_stream, _receiveCallback, isMars);
                     return TdsEnums.SNI_SUCCESS_IO_PENDING;
                 }
                 catch (ObjectDisposedException ode)

src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObjectFactory.Windows.cs

@@ -19,7 +19,7 @@ internal sealed class TdsParserStateObjectFactory
         //private static bool shouldUseLegacyNetorking;
         //public static bool UseManagedSNI { get; } = AppContext.TryGetSwitch(UseLegacyNetworkingOnWindows, out shouldUseLegacyNetorking) ? !shouldUseLegacyNetorking : true;
 
-        public static bool UseManagedSNI { get; } = false;
+        public static bool UseManagedSNI { get; } = true;
 
         public EncryptionOptions EncryptionOptions
         {

src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObjectManaged.cs

@@ -160,7 +160,7 @@ internal override uint CheckConnection()
         internal override object ReadAsync(out uint error, ref object handle)
         {
             SNIPacket packet;
-            error = SNIProxy.Singleton.ReadAsync((SNIHandle)handle, out packet);
+            error = SNIProxy.Singleton.ReadAsync((SNIHandle)handle, out packet, isMars: _parser.MARSOn);
             return packet;
         }