WIP | Prevent WriteAsync collision

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SNI/SNICommon.cs

@@ -3,8 +3,11 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.IO;
 using System.Net.Security;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Data.SqlClient.SNI
 {
@@ -99,6 +102,26 @@ internal enum SNISMUXFlags
         SMUX_DATA = 8       // SMUX data packet
     }
 
+    internal class SslStreamProxy : SslStream
+    {
+        private Task _currentTask;
+
+        public SslStreamProxy(Stream innerStream, bool leaveInnerStreamOpen, RemoteCertificateValidationCallback userCertificateValidationCallback) 
+            : base(innerStream, leaveInnerStreamOpen, userCertificateValidationCallback)
+        { }
+
+        // Prevent the WriteAsync's collision
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            if (_currentTask != null && _currentTask.Status != TaskStatus.RanToCompletion)
+            {
+                _currentTask.Wait(cancellationToken);
+            }
+            _currentTask = base.WriteAsync(buffer, offset, count, cancellationToken);
+            return _currentTask;
+        }
+    }
+
     internal class SNICommon
     {
         // Each error number maps to SNI_ERROR_* in String.resx

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SNI/SNINpHandle.cs

@@ -93,7 +93,7 @@ public SNINpHandle(string serverName, string pipeName, long timerExpire, object
                 }
 
                 _sslOverTdsStream = new SslOverTdsStream(_pipeStream);
-                _sslStream = new SslStream(_sslOverTdsStream, true, new RemoteCertificateValidationCallback(ValidateServerCertificate), null);
+                _sslStream = new SslStreamProxy(_sslOverTdsStream, true, new RemoteCertificateValidationCallback(ValidateServerCertificate));
 
                 _stream = _pipeStream;
                 _status = TdsEnums.SNI_SUCCESS;
@@ -325,12 +325,17 @@ public override uint Send(SNIPacket packet)
         public override uint SendAsync(SNIPacket packet, SNIAsyncCallback callback = null)
         {
             long scopeID = SqlClientEventSource.Log.TrySNIScopeEnterEvent("<sc.SNI.SNINpHandle.SendAsync |SNI|SCOPE>");
+            SNIAsyncCallback cb = callback ?? _sendCallback;
             try
             {
-                SNIAsyncCallback cb = callback ?? _sendCallback;
                 packet.WriteToStreamAsync(_stream, cb, SNIProviders.NP_PROV);
                 return TdsEnums.SNI_SUCCESS_IO_PENDING;
             }
+            catch (Exception e) when (e is ObjectDisposedException || e is InvalidOperationException || e is IOException)
+            {
+                SNIPacket errorPacket = packet;
+                return ReportErrorAndReleasePacket(errorPacket, e);
+            }
             finally
             {
                 SqlClientEventSource.Log.TrySNIScopeLeaveEvent(scopeID);

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SNI/SNIPacket.cs

@@ -304,16 +304,21 @@ public void WriteToStream(Stream stream)
         public async void WriteToStreamAsync(Stream stream, SNIAsyncCallback callback, SNIProviders provider)
         {
             uint status = TdsEnums.SNI_SUCCESS;
-            try
-            {
-                await stream.WriteAsync(_data, 0, _dataLength, CancellationToken.None).ConfigureAwait(false);
-            }
-            catch (Exception e)
+
+            await stream.WriteAsync(_data, 0, _dataLength, CancellationToken.None).ContinueWith(t =>
             {
-                SNILoadHandle.SingletonInstance.LastError = new SNIError(provider, SNICommon.InternalExceptionError, e);
-                status = TdsEnums.SNI_ERROR;
-            }
-            callback(this, status);
+                Exception e = t.Exception?.InnerException;
+                if (e != null)
+                {
+                    SNILoadHandle.SingletonInstance.LastError = new SNIError(provider, SNICommon.InternalExceptionError, e);
+                    status = TdsEnums.SNI_ERROR;
+                    Release();
+                }
+                callback(this, status);
+            },
+            CancellationToken.None,
+            TaskContinuationOptions.DenyChildAttach,
+            TaskScheduler.Default);
         }
     }
 }

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SNI/SNITcpHandle.cs

@@ -226,7 +226,7 @@ public SNITCPHandle(string serverName, int port, long timerExpire, object callba
                 _tcpStream = new NetworkStream(_socket, true);
 
                 _sslOverTdsStream = new SslOverTdsStream(_tcpStream);
-                _sslStream = new SslStream(_sslOverTdsStream, true, new RemoteCertificateValidationCallback(ValidateServerCertificate), null);
+                _sslStream = new SslStreamProxy(_sslOverTdsStream, true, new RemoteCertificateValidationCallback(ValidateServerCertificate));
             }
             catch (SocketException se)
             {
@@ -331,7 +331,7 @@ private static Socket Connect(string serverName, int port, TimeSpan timeout, boo
             }
 
             CancellationTokenSource cts = null;
-            
+
             void Cancel()
             {
                 for (int i = 0; i < sockets.Length; ++i)
@@ -355,7 +355,7 @@ void Cancel()
             }
 
             Socket availableSocket = null;
-            try 
+            try
             {
                 for (int i = 0; i < sockets.Length; ++i)
                 {
@@ -706,12 +706,22 @@ public override void SetAsyncCallbacks(SNIAsyncCallback receiveCallback, SNIAsyn
         /// <returns>SNI error code</returns>
         public override uint SendAsync(SNIPacket packet, SNIAsyncCallback callback = null)
         {
+            long scopeID = SqlClientEventSource.Log.TrySNIScopeEnterEvent("<sc.SNI.SNITcpHandle.SendAsync |SNI|SCOPE>");
             SNIAsyncCallback cb = callback ?? _sendCallback;
-            lock (this)
+            try
             {
                 packet.WriteToStreamAsync(_stream, cb, SNIProviders.TCP_PROV);
+                return TdsEnums.SNI_SUCCESS_IO_PENDING;
+            }
+            catch (Exception e) when (e is ObjectDisposedException || e is InvalidOperationException || e is IOException)
+            {
+                SNIPacket errorPacket = packet;
+                return ReportErrorAndReleasePacket(errorPacket, e);
+            }
+            finally
+            {
+                SqlClientEventSource.Log.TrySNIScopeLeaveEvent(scopeID);
             }
-            return TdsEnums.SNI_SUCCESS_IO_PENDING;
         }
 
         /// <summary>

src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj

@@ -66,6 +66,7 @@
     <Compile Include="SQL\SqlBulkCopyTest\OrderHintMissingTargetColumn.cs" />
     <Compile Include="SQL\SqlBulkCopyTest\OrderHintTransaction.cs" />
     <Compile Include="SQL\DataClassificationTest\DataClassificationTest.cs" />
+    <Compile Include="SQL\SqlCommand\SqlCommandExecuteTest.cs" />
     <Compile Include="TracingTests\EventSourceTest.cs" />
     <Compile Include="SQL\AdapterTest\AdapterTest.cs" />
     <Compile Include="SQL\AsyncTest\BeginExecAsyncTest.cs" />

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/SqlCommand/SqlCommandCompletedTest.cs

@@ -1,4 +1,8 @@
-using System.Data;
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Data;
 using Xunit;
 
 namespace Microsoft.Data.SqlClient.ManualTesting.Tests

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/SqlCommand/SqlCommandExecuteTest.cs

@@ -0,0 +1,141 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    public static class SqlCommandExecuteTest
+    {
+        [Theory]
+        [MemberData(nameof(GetConnectionStrings))]
+        public static void ExecuteReaderAsyncTest(string connectionString)
+        {
+            int counter = 100;
+            while (counter-- > 0)
+            {
+                ExecuteReaderAsync(connectionString, CancellationToken.None).GetAwaiter().GetResult();
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(GetConnectionStrings))]
+        public static void ExecuteScalarAsyncTest(string connectionString)
+        {
+            int counter = 100;
+            while (counter-- > 0)
+            {
+                ExecuteScalarAsync(connectionString, CancellationToken.None).GetAwaiter().GetResult();
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(GetConnectionStrings))]
+        public static void ExecuteNonQueryAsyncTest(string connectionString)
+        {
+            int counter = 100;
+            while (counter-- > 0)
+            {
+                ExecuteNonQueryAsync(connectionString, CancellationToken.None).GetAwaiter().GetResult();
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(GetConnectionStrings))]
+        public static void ExecuteXmlReaderAsyncTest(string connectionString)
+        {
+            int counter = 100;
+            while (counter-- > 0)
+            {
+                ExecuteXmlReaderAsync(connectionString, CancellationToken.None).GetAwaiter().GetResult();
+            }
+        }
+
+        #region Execute Async
+        private static async Task ExecuteReaderAsync(string connectionString, CancellationToken token)
+        {
+            using (var connection = new SqlConnection(connectionString))
+            using (var cmd = GetCommand(connection))
+            using (var r = await cmd.ExecuteReaderAsync(token))
+            {
+                while (await r.ReadAsync(token))
+                {
+                    await r.GetFieldValueAsync<string>(0);
+                    await r.GetFieldValueAsync<string>(1);
+                    await r.GetFieldValueAsync<string>(2);
+                }
+            }
+        }
+
+        private static async Task ExecuteScalarAsync(string connectionString, CancellationToken token)
+        {
+            using (var connection = new SqlConnection(connectionString))
+            using (var cmd = GetCommand(connection))
+            {
+                await cmd.ExecuteScalarAsync(token);
+            }
+        }
+
+        private static async Task ExecuteNonQueryAsync(string connectionString, CancellationToken token)
+        {
+            using (var connection = new SqlConnection(connectionString))
+            using (var cmd = GetCommand(connection))
+            {
+                await cmd.ExecuteNonQueryAsync(token);
+            }
+        }
+
+        private static async Task ExecuteXmlReaderAsync(string connectionString, CancellationToken token)
+        {
+            using (var connection = new SqlConnection(connectionString))
+            using (var cmd = GetCommand(connection))
+            {
+                cmd.CommandText += " FOR XML AUTO, XMLDATA";
+                var r = await cmd.ExecuteXmlReaderAsync(token);
+            }
+        }
+
+        private static SqlCommand GetCommand(SqlConnection cnn)
+        {
+            string aRecord = "('2455cf1b-ebcf-418d-8cce-88e21e1683e3', 'something', 'updated'),";
+            string query = "SELECT * FROM (VALUES"
+                + string.Concat(Enumerable.Repeat(aRecord, 200)).Substring(0, (aRecord.Length * 200) - 1)
+                + ") tbl_A ([Id], [Name], [State])";
+            cnn.Open();
+            var cmd = cnn.CreateCommand();
+            cmd.CommandText = query;
+            return cmd;
+        }
+        #endregion
+
+        public static IEnumerable<object[]> GetConnectionStrings()
+        {
+            SqlConnectionStringBuilder builder;
+            foreach (var item in DataTestUtility.ConnectionStrings)
+            {
+                builder = new SqlConnectionStringBuilder(item)
+                {
+                    TrustServerCertificate = true,
+                    Encrypt = true,
+                    MultipleActiveResultSets = false,
+                    ConnectTimeout = 10,
+                    ConnectRetryCount = 3,
+                    ConnectRetryInterval = 10,
+                    LoadBalanceTimeout = 60,
+                    MaxPoolSize = 10,
+                    MinPoolSize = 0
+                };
+                yield return new object[] { builder.ConnectionString };
+
+                builder.TrustServerCertificate = false;
+                builder.Encrypt = false;
+                yield return new object[] { builder.ConnectionString };
+            }
+        }
+    }
+}

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/SqlCommand/SqlCommandSetTest.cs

@@ -1,8 +1,10 @@
-using System;
-using System.Collections.Generic;
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
 using System.Data.SqlTypes;
 using System.Reflection;
-using System.Text;
 using Xunit;
 
 namespace Microsoft.Data.SqlClient.ManualTesting.Tests