SqlClient Fix managed command cancellation

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

@@ -8,6 +8,7 @@
 using System.Net.Security;
 using System.Security.Authentication;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Data.SqlClient.SNI
@@ -22,7 +23,7 @@ internal sealed class SNINpHandle : SNIHandle
 
         private readonly string _targetServer;
         private readonly object _callbackObject;
-
+        private readonly object _sendSync;
         private Stream _stream;
         private NamedPipeClientStream _pipeStream;
         private SslOverTdsStream _sslOverTdsStream;
@@ -37,6 +38,7 @@ internal sealed class SNINpHandle : SNIHandle
 
         public SNINpHandle(string serverName, string pipeName, long timerExpire, object callbackObject)
         {
+            _sendSync = new object();
             _targetServer = serverName;
             _callbackObject = callbackObject;
 
@@ -193,20 +195,48 @@ public override uint ReceiveAsync(ref SNIPacket packet)
 
         public override uint Send(SNIPacket packet)
         {
-            lock (this)
+            bool releaseLock = false;
+            try
             {
-                try
+                // is the packet is marked out out-of-band (attention packets only) it must be
+                // sent immediately even if a send of recieve operation is already in progress
+                // because out of band packets are used to cancel ongoing operations
+                // so try to take the lock if possible but continue even if it can't be taken
+                if (packet.IsOutOfBand)
                 {
-                    packet.WriteToStream(_stream);
-                    return TdsEnums.SNI_SUCCESS;
+                    Monitor.TryEnter(this, ref releaseLock);
                 }
-                catch (ObjectDisposedException ode)
+                else
                 {
-                    return ReportErrorAndReleasePacket(packet, ode);
+                    Monitor.Enter(this);
+                    releaseLock = true;
                 }
-                catch (IOException ioe)
+
+                // this lock ensures that two packets are not being written to the transport at the same time
+                // so that sending a standard and an out-of-band packet are both written atomically no data is 
+                // interleaved
+                lock (_sendSync)
                 {
-                    return ReportErrorAndReleasePacket(packet, ioe);
+                    try
+                    {
+                        packet.WriteToStream(_stream);
+                        return TdsEnums.SNI_SUCCESS;
+                    }
+                    catch (ObjectDisposedException ode)
+                    {
+                        return ReportErrorAndReleasePacket(packet, ode);
+                    }
+                    catch (IOException ioe)
+                    {
+                        return ReportErrorAndReleasePacket(packet, ioe);
+                    }
+                }
+            }
+            finally
+            {
+                if (releaseLock)
+                {
+                    Monitor.Exit(this);
                 }
             }
         }

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

@@ -40,6 +40,15 @@ public SNIPacket(int headerSize, int dataSize)
         /// </summary>
         public bool IsInvalid => _data is null;
 
+        /// <summary>
+        /// Indicates that the packet should be sent out of band bypassing the normal send-recieve lock
+        /// </summary>
+        public bool IsOutOfBand { get; set; }
+
+        /// <summary>
+        /// Indicates the amonut of reserved header space available at the start of this packet.
+        /// Used with SMUX packet headers, see <see cref="GetHeaderBuffer"/> and <see cref="SetHeaderActive"/>
+        /// </summary>
         public int ReservedHeaderSize => _headerLength;
 
         /// <summary>
@@ -74,7 +83,7 @@ private void Allocate(int headerLength, int dataLength)
         }
 
         /// <summary>
-        /// Read packet data into a buffer without removing it from the packet
+        /// Get packet data
         /// </summary>
         /// <param name="buffer">Buffer</param>
         /// <param name="dataSize">Number of bytes read from the packet into the buffer</param>

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

@@ -23,6 +23,7 @@ internal sealed class SNITCPHandle : SNIHandle
     {
         private readonly string _targetServer;
         private readonly object _callbackObject;
+        private readonly object _sendSync;
         private readonly Socket _socket;
         private NetworkStream _tcpStream;
 
@@ -100,6 +101,7 @@ public override uint Status
         /// <param name="callbackObject">Callback object</param>
         public SNITCPHandle(string serverName, int port, long timerExpire, object callbackObject, bool parallel)
         {
+            _sendSync = new object();
             _callbackObject = callbackObject;
             _targetServer = serverName;
 
@@ -420,24 +422,52 @@ public override void SetBufferSize(int bufferSize)
         /// <returns>SNI error code</returns>
         public override uint Send(SNIPacket packet)
         {
-            lock (this)
+            bool releaseLock = false;
+            try
             {
-                try
+                // is the packet is marked out out-of-band (attention packets only) it must be
+                // sent immediately even if a send of recieve operation is already in progress
+                // because out of band packets are used to cancel ongoing operations
+                // so try to take the lock if possible but continue even if it can't be taken
+                if (packet.IsOutOfBand)
                 {
-                    packet.WriteToStream(_stream);
-                    return TdsEnums.SNI_SUCCESS;
+                    Monitor.TryEnter(this, ref releaseLock);
                 }
-                catch (ObjectDisposedException ode)
+                else
                 {
-                    return ReportTcpSNIError(ode);
+                    Monitor.Enter(this);
+                    releaseLock = true;
                 }
-                catch (SocketException se)
+
+                // this lock ensures that two packets are not being written to the transport at the same time
+                // so that sending a standard and an out-of-band packet are both written atomically no data is 
+                // interleaved
+                lock (_sendSync)
                 {
-                    return ReportTcpSNIError(se);
+                    try
+                    {
+                        packet.WriteToStream(_stream);
+                        return TdsEnums.SNI_SUCCESS;
+                    }
+                    catch (ObjectDisposedException ode)
+                    {
+                        return ReportTcpSNIError(ode);
+                    }
+                    catch (SocketException se)
+                    {
+                        return ReportTcpSNIError(se);
+                    }
+                    catch (IOException ioe)
+                    {
+                        return ReportTcpSNIError(ioe);
+                    }
                 }
-                catch (IOException ioe)
+            }
+            finally
+            {
+                if (releaseLock)
                 {
-                    return ReportTcpSNIError(ioe);
+                    Monitor.Exit(this);
                 }
             }
         }

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

@@ -144,6 +144,7 @@ internal override PacketHandle CreateAndSetAttentionPacket()
         {
             PacketHandle packetHandle = GetResetWritePacket(TdsEnums.HEADER_LEN);
             SetPacketData(packetHandle, SQL.AttentionHeader, TdsEnums.HEADER_LEN);
+            packetHandle.ManagedPacket.IsOutOfBand = true;
             return packetHandle;
         }
 

src/System.Data.SqlClient/tests/ManualTests/SQL/CommandCancelTest/CommandCancelTest.cs

@@ -13,25 +13,25 @@ public static class CommandCancelTest
         // Shrink the packet size - this should make timeouts more likely
         private static readonly string s_connStr = (new SqlConnectionStringBuilder(DataTestUtility.TcpConnStr) { PacketSize = 512 }).ConnectionString;
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void PlainCancelTest()
         {
             PlainCancel(s_connStr);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void PlainMARSCancelTest()
         {
             PlainCancel((new SqlConnectionStringBuilder(s_connStr) { MultipleActiveResultSets = true }).ConnectionString);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void PlainCancelTestAsync()
         {
             PlainCancelAsync(s_connStr);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void PlainMARSCancelTestAsync()
         {
             PlainCancelAsync((new SqlConnectionStringBuilder(s_connStr) { MultipleActiveResultSets = true }).ConnectionString);
@@ -87,31 +87,31 @@ private static void PlainCancelAsync(string connString)
             }
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void MultiThreadedCancel_NonAsync()
         {
             MultiThreadedCancel(s_connStr, false);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void MultiThreadedCancel_Async()
         {
             MultiThreadedCancel(s_connStr, true);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void TimeoutCancel()
         {
             TimeoutCancel(s_connStr);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void CancelAndDisposePreparedCommand()
         {
             CancelAndDisposePreparedCommand(s_connStr);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void TimeOutDuringRead()
         {
             TimeOutDuringRead(s_connStr);
@@ -135,7 +135,7 @@ private static void MultiThreadedCancel(string constr, bool async)
                 tasks[1].Start();
 
                 Task.WaitAll(tasks, 15 * 1000);
-                
+
                 CommandCancelTest.VerifyConnection(command);
             }
         }
@@ -224,7 +224,7 @@ private static void ExecuteCommandCancelExpected(object state)
                     } while (r.NextResult());
                 }
             }, new string[] { errorMessage, errorMessageSevereFailure });
-            
+
         }
 
         private static void CancelSharedCommand(object state)
@@ -279,5 +279,65 @@ private static void TimeOutDuringRead(string constr)
                 throw;
             }
         }
+
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static void CancelDoesNotWait()
+        {
+            const int delaySeconds = 30;
+            const int cancelSeconds = 1;
+
+            using (SqlConnection conn = new SqlConnection(s_connStr))
+            using (var cmd = new SqlCommand($"WAITFOR DELAY '00:00:{delaySeconds:D2}'", conn))
+            {
+                conn.Open();
+
+                Task.Delay(TimeSpan.FromSeconds(cancelSeconds))
+                    .ContinueWith(t => cmd.Cancel());
+
+                DateTime started = DateTime.UtcNow;
+                DateTime ended = default;
+                Exception exception = null;
+                try
+                {
+                    cmd.ExecuteNonQuery();
+                }
+                catch (Exception ex)
+                {
+                    exception = ex;
+                }
+                ended = DateTime.UtcNow;
+
+                Assert.NotNull(exception);
+                Assert.InRange((ended - started).TotalSeconds, cancelSeconds, delaySeconds-1);
+            }
+        }
+
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static async Task AsyncCancelDoesNotWait()
+        {
+            const int delaySeconds = 30;
+            const int cancelSeconds = 1;
+
+            using (SqlConnection conn = new SqlConnection(s_connStr))
+            using (var cmd = new SqlCommand($"WAITFOR DELAY '00:00:{delaySeconds:D2}'", conn))
+            {
+                await conn.OpenAsync();
+
+                DateTime started = DateTime.UtcNow;
+                Exception exception = null;
+                try
+                {
+                    await cmd.ExecuteNonQueryAsync(new CancellationTokenSource(1000).Token);
+                }
+                catch (Exception ex)
+                {
+                    exception = ex;
+                }
+                DateTime ended = DateTime.UtcNow;
+
+                Assert.NotNull(exception);
+                Assert.InRange((ended - started).TotalSeconds, cancelSeconds, delaySeconds - 1);
+            }
+        }
     }
 }