Cleanup | AsyncHelpers Generic State

Directory.Packages.props

@@ -31,6 +31,7 @@
     <PackageVersion Include="Microsoft.SqlServer.SqlManagementObjects" Version="172.76.0" />
     <PackageVersion Include="Microsoft.SqlServer.Types" Version="160.1000.6" />
     <PackageVersion Include="Microsoft.VisualStudio.Azure.Containers.Tools.Targets" Version="1.19.5" />
+    <PackageVersion Include="Moq" Version="4.20.72" />
     <PackageVersion Include="Newtonsoft.Json" Version="13.0.3" />
     <PackageVersion Include="System.Data.Odbc" Version="9.0.9" />
     <PackageVersion Include="System.Security.Principal.Windows" Version="5.0.0" />

src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj

@@ -938,6 +938,9 @@
     <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\SSPI\SspiContextProvider.cs">
       <Link>Microsoft\Data\SqlClient\SSPI\SspiContextProvider.cs</Link>
     </Compile>
+    <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\Utilities\AsyncHelper.cs">
+      <Link>Microsoft\Data\SqlClient\Utilities\AsyncHelper.cs</Link>
+    </Compile>
     <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\Utilities\ObjectPool.cs">
       <Link>Microsoft\Data\SqlClient\Utilities\ObjectPool.cs</Link>
     </Compile>

src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj

@@ -966,6 +966,9 @@
     <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\TransactionRequest.cs">
       <Link>Microsoft\Data\SqlClient\TransactionRequest.cs</Link>
     </Compile>
+    <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\Utilities\AsyncHelper.cs">
+      <Link>Microsoft\Data\SqlClient\Utilities\AsyncHelper.cs</Link>
+    </Compile>
     <Compile Include="$(CommonSourceRoot)Microsoft\Data\SqlClient\Utilities\BufferWriterExtensions.netfx.cs">
       <Link>Microsoft\Data\SqlClient\Utilities\BufferWriterExtensions.netfx.cs</Link>
     </Compile>

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientEventSource.cs

@@ -330,6 +330,19 @@ private string GetFormattedMessage(string className, string memberName, string e
         #region Trace
 
         #region Traces without if statements
+
+        [NonEvent]
+        internal void TraceEvent(string message)
+        {
+            Trace(message);
+        }
+
+        [NonEvent]
+        internal void TraceEvent<T0>(string message, T0 args0)
+        {
+            Trace(string.Format(message, args0?.ToString() ?? NullStr));
+        }
+
         [NonEvent]
         internal void TraceEvent<T0, T1>(string message, T0 args0, T1 args1)
         {

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs

@@ -12,6 +12,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Data.Common;
+using Microsoft.Data.SqlClient.Utilities;
 
 namespace Microsoft.Data.SqlClient
 {
@@ -250,24 +251,22 @@ private SqlDataReader GetParameterEncryptionDataReader(
             bool isRetry)
         {
             returnTask = AsyncHelper.CreateContinuationTaskWithState(
-                task: fetchInputParameterEncryptionInfoTask,
+                taskToContinue: fetchInputParameterEncryptionInfoTask,
                 state: this,
-                onSuccess: state =>
+                onSuccess: this2 =>
                 {
-                    SqlCommand command = (SqlCommand)state;
                     bool processFinallyBlockAsync = true;
                     bool decrementAsyncCountInFinallyBlockAsync = true;
 
                     try
                     {
                         // Check for any exceptions on network write, before reading.
-                        command.CheckThrowSNIException();
+                        this2.CheckThrowSNIException();
 
                         // If it is async, then TryFetchInputParameterEncryptionInfo ->
                         // RunExecuteReaderTds would have incremented the async count. Decrement it
                         // when we are about to complete async execute reader.
-                        SqlInternalConnectionTds internalConnectionTds =
-                            command._activeConnection.GetOpenTdsConnection();
+                        SqlInternalConnectionTds internalConnectionTds = this2._activeConnection.GetOpenTdsConnection();
                         if (internalConnectionTds is not null)
                         {
                             internalConnectionTds.DecrementAsyncCount();
@@ -276,13 +275,13 @@ private SqlDataReader GetParameterEncryptionDataReader(
 
                         // Complete executereader.
                         // @TODO: If we can remove this reference, this could be a static lambda
-                        describeParameterEncryptionDataReader = command.CompleteAsyncExecuteReader(
+                        describeParameterEncryptionDataReader = this2.CompleteAsyncExecuteReader(
                             isInternal: false,
                             forDescribeParameterEncryption: true);
-                        Debug.Assert(command._stateObj is null, "non-null state object in PrepareForTransparentEncryption.");
+                        Debug.Assert(this2._stateObj is null, "non-null state object in PrepareForTransparentEncryption.");
 
                         // Read the results of describe parameter encryption.
-                        command.ReadDescribeEncryptionParameterResults(
+                        this2.ReadDescribeEncryptionParameterResults(
                             describeParameterEncryptionDataReader,
                             describeParameterEncryptionRpcOriginalRpcMap,
                             isRetry);
@@ -302,7 +301,7 @@ private SqlDataReader GetParameterEncryptionDataReader(
                     }
                     finally
                     {
-                        command.PrepareTransparentEncryptionFinallyBlock(
+                        this2.PrepareTransparentEncryptionFinallyBlock(
                             closeDataReader: processFinallyBlockAsync,
                             decrementAsyncCount: decrementAsyncCountInFinallyBlockAsync,
                             clearDataStructures: processFinallyBlockAsync,
@@ -311,10 +310,9 @@ private SqlDataReader GetParameterEncryptionDataReader(
                             describeParameterEncryptionDataReader: describeParameterEncryptionDataReader);
                     }
                 },
-                onFailure: static (exception, state) =>
+                onFailure: static (this2, exception) =>
                 {
-                    SqlCommand command = (SqlCommand)state;
-                    command.CachedAsyncState?.ResetAsyncState();
+                    this2.CachedAsyncState?.ResetAsyncState();
 
                     if (exception is not null)
                     {

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs

@@ -9,6 +9,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Data.Common;
+using Microsoft.Data.SqlClient.Utilities;
 
 #if NETFRAMEWORK
 using System.Security.Permissions;
@@ -219,14 +220,12 @@ private IAsyncResult BeginExecuteNonQueryInternal(
                     if (execNonQuery is not null)
                     {
                         AsyncHelper.ContinueTaskWithState(
-                            task: execNonQuery,
-                            completion: localCompletion,
-                            state: Tuple.Create(this, localCompletion),
-                            onSuccess: static state =>
-                            {
-                                var parameters = (Tuple<SqlCommand, TaskCompletionSource<object>>)state;
-                                parameters.Item1.BeginExecuteNonQueryInternalReadStage(parameters.Item2);
-                            });
+                            taskToContinue: execNonQuery,
+                            taskCompletionSource: localCompletion,
+                            state1: this,
+                            state2: localCompletion,
+                            onSuccess: static (this2, localCompletion2) =>
+                                this2.BeginExecuteNonQueryInternalReadStage(localCompletion2));
                     }
                     else
                     {
@@ -871,8 +870,8 @@ private void RunExecuteNonQueryTdsSetupReconnnectContinuation(
                 timeoutCts.Token);
 
             AsyncHelper.ContinueTask(
-                reconnectTask,
-                completion,
+                taskToContinue: reconnectTask,
+                taskCompletionSource: completion,
                 onSuccess: () =>
                 {
                     if (completion.Task.IsCompleted)
@@ -896,10 +895,10 @@ private void RunExecuteNonQueryTdsSetupReconnnectContinuation(
                     else
                     {
                         AsyncHelper.ContinueTaskWithState(
-                            subTask,
-                            completion,
+                            taskToContinue: subTask,
+                            taskCompletionSource: completion,
                             state: completion,
-                            onSuccess: static state => ((TaskCompletionSource<object>)state).SetResult(null));
+                            onSuccess: static state => state.SetResult(null));
                     }
                 });
         }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs

@@ -10,10 +10,10 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Data.Common;
+using Microsoft.Data.SqlClient.Utilities;
 
 #if NETFRAMEWORK
 using System.Security.Permissions;
-using Microsoft.Data.SqlClient.Utilities;
 #endif
 
 namespace Microsoft.Data.SqlClient
@@ -308,14 +308,12 @@ private IAsyncResult BeginExecuteReaderInternal(
                 if (writeTask is not null)
                 {
                     AsyncHelper.ContinueTaskWithState(
-                        writeTask,
-                        localCompletion,
-                        state: Tuple.Create(this, localCompletion),
-                        onSuccess: static state =>
-                        {
-                            var parameters = (Tuple<SqlCommand, TaskCompletionSource<object>>)state;
-                            parameters.Item1.BeginExecuteReaderInternalReadStage(parameters.Item2);
-                        });
+                        taskToContinue: writeTask,
+                        taskCompletionSource: localCompletion,
+                        state1: this,
+                        state2: localCompletion,
+                        onSuccess: static (this2, localCompletion2) =>
+                            this2.BeginExecuteReaderInternalReadStage(localCompletion2));
                 }
                 else
                 {
@@ -1605,21 +1603,19 @@ private Task RunExecuteReaderTdsSetupContinuation(
             string optionSettings,
             Task writeTask)
         {
-            // @TODO: Why use the state version if we can't make this a static helper?
             return AsyncHelper.CreateContinuationTaskWithState(
-                task: writeTask,
-                state: _activeConnection,
-                onSuccess: state =>
+                taskToContinue: writeTask,
+                state1: this,
+                state2: Tuple.Create(ds, runBehavior, optionSettings),
+                onSuccess: static (this2, parameters) =>
                 {
                     // This will throw if the connection is closed.
                     // @TODO: So... can we have something that specifically does that?
-                    ((SqlConnection)state).GetOpenTdsConnection();
-                    CachedAsyncState.SetAsyncReaderState(ds, runBehavior, optionSettings);
+                    this2._activeConnection.GetOpenTdsConnection();
+                    this2.CachedAsyncState.SetAsyncReaderState(parameters.Item1, parameters.Item2, parameters.Item3);
                 },
-                onFailure: static (exception, state) =>
-                {
-                    ((SqlConnection)state).GetOpenTdsConnection().DecrementAsyncCount();
-                });
+                onFailure: static (this2, _, _) =>
+                    this2._activeConnection.GetOpenTdsConnection().DecrementAsyncCount());
         }
 
         // @TODO: This is way too many parameters being shoveled back and forth. We can do better.
@@ -1640,13 +1636,12 @@ private void RunExecuteReaderTdsSetupReconnectContinuation(
             AsyncHelper.SetTimeoutException(
                 completion,
                 timeout,
-                onFailure: static () => SQL.CR_ReconnectTimeout(),
+                onTimeout: static () => SQL.CR_ReconnectTimeout(),
                 timeoutCts.Token);
 
-            // @TODO: With an object to pass around we can use the state-based version
             AsyncHelper.ContinueTask(
-                reconnectTask,
-                completion,
+                taskToContinue: reconnectTask,
+                taskCompletionSource: completion,
                 onSuccess: () =>
                 {
                     if (completion.Task.IsCompleted)
@@ -1675,10 +1670,10 @@ private void RunExecuteReaderTdsSetupReconnectContinuation(
                     else
                     {
                         AsyncHelper.ContinueTaskWithState(
-                            subTask,
-                            completion,
+                            taskToContinue: subTask,
+                            taskCompletionSource: completion,
                             state: completion,
-                            onSuccess: static state => ((TaskCompletionSource<object>)state).SetResult(null));
+                            onSuccess: static completion2 => completion2.SetResult(null));
                     }
                 });
         }
@@ -1711,14 +1706,13 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption(
                 // @TODO: This is a prime candidate for proper async-await execution
                 TaskCompletionSource<object> completion = new TaskCompletionSource<object>();
                 AsyncHelper.ContinueTaskWithState(
-                    task: describeParameterEncryptionTask,
-                    completion: completion,
+                    taskToContinue: describeParameterEncryptionTask,
+                    taskCompletionSource: completion,
                     state: this,
-                    onSuccess: state =>
+                    onSuccess: this2 =>
                     {
-                        SqlCommand command = (SqlCommand)state;
-                        command.GenerateEnclavePackage();
-                        command.RunExecuteReaderTds(
+                        this2.GenerateEnclavePackage();
+                        this2.RunExecuteReaderTds(
                             cmdBehavior,
                             runBehavior,
                             returnStream,
@@ -1737,24 +1731,22 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption(
                         else
                         {
                             AsyncHelper.ContinueTaskWithState(
-                                task: subTask,
-                                completion: completion,
+                                taskToContinue: subTask,
+                                taskCompletionSource: completion,
                                 state: completion,
-                                onSuccess: static state => ((TaskCompletionSource<object>)state).SetResult(null));
+                                onSuccess: static state => state.SetResult(null));
                         }
                     },
-                    onFailure: static (exception, state) =>
+                    onFailure: static (this2, exception) =>
                     {
-                        ((SqlCommand)state).CachedAsyncState?.ResetAsyncState();
+                        this2.CachedAsyncState?.ResetAsyncState();
                         if (exception is not null)
                         {
+                            // @TODO: This doesn't do anything, afaik.
                             throw exception;
                         }
                     },
-                    onCancellation: static state =>
-                    {
-                        ((SqlCommand)state).CachedAsyncState?.ResetAsyncState();
-                    });
+                    onCancellation: static this2 => this2.CachedAsyncState?.ResetAsyncState());
 
                 task = completion.Task;
                 return ds;

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs

@@ -10,6 +10,7 @@
 using System.Xml;
 using Microsoft.Data.Common;
 using Microsoft.Data.SqlClient.Server;
+using Microsoft.Data.SqlClient.Utilities;
 
 #if NETFRAMEWORK
 using System.Security.Permissions;
@@ -269,14 +270,12 @@ private IAsyncResult BeginExecuteXmlReaderInternal(
                 if (writeTask is not null)
                 {
                     AsyncHelper.ContinueTaskWithState(
-                        task: writeTask,
-                        completion: localCompletion,
-                        state: Tuple.Create(this, localCompletion),
-                        onSuccess: static state =>
-                        {
-                            var parameters = (Tuple<SqlCommand, TaskCompletionSource<object>>)state;
-                            parameters.Item1.BeginExecuteXmlReaderInternalReadStage(parameters.Item2);
-                        });
+                        taskToContinue: writeTask,
+                        taskCompletionSource: localCompletion,
+                        state1: this,
+                        state2: localCompletion,
+                        onSuccess: static (this2, localCompletion2) =>
+                            this2.BeginExecuteXmlReaderInternalReadStage(localCompletion2));
                 }
                 else
                 {

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnection.cs

@@ -23,7 +23,8 @@
 using Microsoft.Data.SqlClient.Connection;
 using Microsoft.Data.SqlClient.ConnectionPool;
 using Microsoft.Data.SqlClient.Diagnostics;
-using Microsoft.SqlServer.Server;
+using Microsoft.Data.SqlClient.Utilities;
+
 #if NETFRAMEWORK
 using System.Runtime.CompilerServices;
 using System.Security.Permissions;

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs

@@ -19,6 +19,7 @@
 using Microsoft.Data.ProviderBase;
 using Microsoft.Data.SqlClient.Connection;
 using Microsoft.Data.SqlClient.ConnectionPool;
+using Microsoft.Data.SqlClient.Utilities;
 using Microsoft.Identity.Client;
 
 namespace Microsoft.Data.SqlClient