diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj index b69dc00a90..8537b27a1e 100644 --- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj @@ -588,6 +588,9 @@ Microsoft\Data\SqlClient\SqlCommand.Scalar.cs + + Microsoft\Data\SqlClient\SqlCommand.Xml.cs + Microsoft\Data\SqlClient\SqlCommandSet.cs diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.netcore.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.netcore.cs index 90979eeec9..b6a4ff8405 100644 --- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.netcore.cs +++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.netcore.cs @@ -58,30 +58,6 @@ protected override void AfterCleared(SqlCommand owner) } } - internal sealed class ExecuteXmlReaderAsyncCallContext : AAsyncCallContext - { - public Guid OperationID; - - public SqlCommand Command => _owner; - public TaskCompletionSource TaskCompletionSource => _source; - - public void Set(SqlCommand command, TaskCompletionSource source, CancellationTokenRegistration disposable, Guid operationID) - { - base.Set(command, source, disposable); - OperationID = operationID; - } - - protected override void Clear() - { - OperationID = default; - } - - protected override void AfterCleared(SqlCommand owner) - { - owner?.SetCachedCommandExecuteXmlReaderContext(this); - } - } - /// /// Indicates if the column encryption setting was set at-least once in the batch rpc mode, when using AddBatchCommand. /// @@ -604,315 +580,11 @@ private void ThrowIfReconnectionHasBeenCanceled() } } - /// - public XmlReader ExecuteXmlReader() - { - // Reset _pendingCancel upon entry into any Execute - used to synchronize state - // between entry into Execute* API and the thread obtaining the stateObject. - _pendingCancel = false; - - using (DiagnosticScope diagnosticScope = s_diagnosticListener.CreateCommandScope(this, _transaction)) - using (TryEventScope.Create("SqlCommand.ExecuteXmlReader | API | Object Id {0}", ObjectID)) - { - SqlStatistics statistics = null; - bool success = false; - int? sqlExceptionNumber = null; - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.ExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - - try - { - statistics = SqlStatistics.StartTimer(Statistics); - WriteBeginExecuteEvent(); - - // use the reader to consume metadata - SqlDataReader ds = IsProviderRetriable ? - RunExecuteReaderWithRetry(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true) : - RunExecuteReader(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true); - success = true; - return CompleteXmlReader(ds); - } - catch (Exception ex) - { - diagnosticScope.SetException(ex); - if (ex is SqlException sqlException) - { - sqlExceptionNumber = sqlException.Number; - } - throw; - } - finally - { - SqlStatistics.StopTimer(statistics); - WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: true); - } - } - } - - /// - public IAsyncResult BeginExecuteXmlReader() - { - // BeginExecuteXmlReader will track executiontime - return BeginExecuteXmlReader(null, null); - } - - /// - public IAsyncResult BeginExecuteXmlReader(AsyncCallback callback, object stateObject) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.BeginExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - return BeginExecuteXmlReaderInternal(CommandBehavior.SequentialAccess, callback, stateObject, 0, isRetry: false); - } - - private IAsyncResult BeginExecuteXmlReaderAsync(AsyncCallback callback, object stateObject) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.BeginExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - return BeginExecuteXmlReaderInternal(CommandBehavior.SequentialAccess, callback, stateObject, CommandTimeout, isRetry: false, asyncWrite: true); - } - - private IAsyncResult BeginExecuteXmlReaderInternal(CommandBehavior behavior, AsyncCallback callback, object stateObject, int timeout, bool isRetry, bool asyncWrite = false) - { - TaskCompletionSource globalCompletion = new TaskCompletionSource(stateObject); - TaskCompletionSource localCompletion = new TaskCompletionSource(stateObject); - - if (!isRetry) - { - // Reset _pendingCancel upon entry into any Execute - used to synchronize state - // between entry into Execute* API and the thread obtaining the stateObject. - _pendingCancel = false; - - // Special case - done outside of try/catches to prevent putting a stateObj back into pool when we should not. - ValidateAsyncCommand(); - } - - SqlStatistics statistics = null; - try - { - if (!isRetry) - { - statistics = SqlStatistics.StartTimer(Statistics); - WriteBeginExecuteEvent(); - } - - bool usedCache; - Task writeTask; - try - { - // InternalExecuteNonQuery already has reliability block, but if failure will not put stateObj back into pool. - RunExecuteReader( - behavior, - RunBehavior.ReturnImmediately, - returnStream: true, - localCompletion, - timeout, - out writeTask, - out usedCache, - asyncWrite, - isRetry); - } - catch (Exception e) - { - if (!ADP.IsCatchableOrSecurityExceptionType(e)) - { - // If not catchable - the connection has already been caught and doomed in RunExecuteReader. - throw; - } - - // For async, RunExecuteReader will never put the stateObj back into the pool, so do so now. - ReliablePutStateObject(); - throw; - } - - if (writeTask != null) - { - AsyncHelper.ContinueTaskWithState(writeTask, localCompletion, - state: Tuple.Create(this, localCompletion), - onSuccess: state => - { - var parameters = (Tuple>)state; - parameters.Item1.BeginExecuteXmlReaderInternalReadStage(parameters.Item2); - } - ); - } - else - { - BeginExecuteXmlReaderInternalReadStage(localCompletion); - } - - // When we use query caching for parameter encryption we need to retry on specific errors. - // In these cases finalize the call internally and trigger a retry when needed. - if ( - !TriggerInternalEndAndRetryIfNecessary( - behavior, - stateObject, - timeout, - usedCache, - isRetry, - asyncWrite, - globalCompletion, - localCompletion, - endFunc: static (SqlCommand command, IAsyncResult asyncResult, bool isInternal, string endMethod) => - { - return command.InternalEndExecuteReader(asyncResult, isInternal, endMethod); - }, - retryFunc: static (SqlCommand command, CommandBehavior behavior, AsyncCallback callback, object stateObject, int timeout, bool isRetry, bool asyncWrite) => - { - return command.BeginExecuteXmlReaderInternal(behavior, callback, stateObject, timeout, isRetry, asyncWrite); - }, - endMethod: nameof(EndExecuteXmlReader))) - { - globalCompletion = localCompletion; - } - - // Add callback after work is done to avoid overlapping Begin/End methods - if (callback != null) - { - localCompletion.Task.ContinueWith( - static (Task task, object state) => ((AsyncCallback)state)(task), - state: callback - ); - } - return localCompletion.Task; - } - finally - { - SqlStatistics.StopTimer(statistics); - } - } - - private void BeginExecuteXmlReaderInternalReadStage(TaskCompletionSource completion) - { - Debug.Assert(completion != null, "Completion source should not be null"); - // Read SNI does not have catches for async exceptions, handle here. - try - { - // must finish caching information before ReadSni which can activate the callback before returning - CachedAsyncState.SetActiveConnectionAndResult(completion, nameof(EndExecuteXmlReader), _activeConnection); - _stateObj.ReadSni(completion); - } - // @TODO: CER Exception Handling was removed here (see GH#3581) - catch (Exception e) - { - // Similarly, if an exception occurs put the stateObj back into the pool. - // and reset async cache information to allow a second async execute - CachedAsyncState?.ResetAsyncState(); - ReliablePutStateObject(); - completion.TrySetException(e); - } - } - - /// - public XmlReader EndExecuteXmlReader(IAsyncResult asyncResult) - { - try - { - return EndExecuteXmlReaderInternal(asyncResult); - } - finally - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.EndExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - } - } - - private XmlReader EndExecuteXmlReaderAsync(IAsyncResult asyncResult) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.EndExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - Debug.Assert(!_internalEndExecuteInitiated || _stateObj == null); - - Exception asyncException = ((Task)asyncResult).Exception; - if (asyncException != null) - { - CachedAsyncState?.ResetAsyncState(); - ReliablePutStateObject(); - throw asyncException.InnerException; - } - else - { - ThrowIfReconnectionHasBeenCanceled(); - // lock on _stateObj prevents races with close/cancel. - // If we have already initiate the End call internally, we have already done that, so no point doing it again. - if (!_internalEndExecuteInitiated) - { - lock (_stateObj) - { - return EndExecuteXmlReaderInternal(asyncResult); - } - } - else - { - return EndExecuteXmlReaderInternal(asyncResult); - } - } - } - - private XmlReader EndExecuteXmlReaderInternal(IAsyncResult asyncResult) - { - bool success = false; - int? sqlExceptionNumber = null; - try - { - success = true; - return CompleteXmlReader(InternalEndExecuteReader(asyncResult, false, nameof(EndExecuteXmlReader)), true); - } - catch (Exception e) - { - if (e is SqlException) - { - SqlException ex = (SqlException)e; - sqlExceptionNumber = ex.Number; - } - if (CachedAsyncState != null) - { - CachedAsyncState.ResetAsyncState(); - }; - if (ADP.IsCatchableExceptionType(e)) - { - ReliablePutStateObject(); - }; - throw; - } - finally - { - WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: false); - } - } - - private XmlReader CompleteXmlReader(SqlDataReader ds, bool isAsync = false) - { - XmlReader xr = null; - - SmiExtendedMetaData[] md = ds.GetInternalSmiMetaData(); - bool isXmlCapable = (md != null && md.Length == 1 && (md[0].SqlDbType == SqlDbType.NText - || md[0].SqlDbType == SqlDbType.NVarChar - || md[0].SqlDbType == SqlDbType.Xml)); - - if (isXmlCapable) - { - try - { - SqlStream sqlBuf = new SqlStream(ds, true /*addByteOrderMark*/, (md[0].SqlDbType == SqlDbType.Xml) ? false : true /*process all rows*/); - xr = sqlBuf.ToXmlReader(isAsync); - } - catch (Exception e) - { - if (ADP.IsCatchableExceptionType(e)) - { - ds.Close(); - } - throw; - } - } - if (xr == null) - { - ds.Close(); - throw SQL.NonXmlResult(); - } - return xr; - } - - /// + /// public IAsyncResult BeginExecuteReader() => BeginExecuteReader(callback: null, stateObject: null, CommandBehavior.Default); - /// + /// public IAsyncResult BeginExecuteReader(AsyncCallback callback, object stateObject) => BeginExecuteReader(callback, stateObject, CommandBehavior.Default); @@ -1605,122 +1277,6 @@ private void SetCachedCommandExecuteReaderAsyncContext(ExecuteReaderAsyncCallCon } } - private void SetCachedCommandExecuteXmlReaderContext(ExecuteXmlReaderAsyncCallContext instance) - { - if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) - { - Interlocked.CompareExchange(ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, instance, null); - } - } - - /// - public Task ExecuteXmlReaderAsync() => - ExecuteXmlReaderAsync(CancellationToken.None); - - /// - public Task ExecuteXmlReaderAsync(CancellationToken cancellationToken) => - IsProviderRetriable - ? InternalExecuteXmlReaderWithRetryAsync(cancellationToken) - : InternalExecuteXmlReaderAsync(cancellationToken); - - private Task InternalExecuteXmlReaderWithRetryAsync(CancellationToken cancellationToken) => - RetryLogicProvider.ExecuteAsync( - sender: this, - () => InternalExecuteXmlReaderAsync(cancellationToken), - cancellationToken); - - private Task InternalExecuteXmlReaderAsync(CancellationToken cancellationToken) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.InternalExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - Guid operationId = s_diagnosticListener.WriteCommandBefore(this, _transaction); - - // connection can be used as state in RegisterForConnectionCloseNotification continuation - // to avoid an allocation so use it as the state value if possible but it can be changed if - // you need it for a more important piece of data that justifies the tuple allocation later - TaskCompletionSource source = new TaskCompletionSource(_activeConnection); - - CancellationTokenRegistration registration = new CancellationTokenRegistration(); - if (cancellationToken.CanBeCanceled) - { - if (cancellationToken.IsCancellationRequested) - { - source.SetCanceled(); - return source.Task; - } - registration = cancellationToken.Register(s_cancelIgnoreFailure, this); - } - - ExecuteXmlReaderAsyncCallContext context = null; - if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) - { - context = Interlocked.Exchange(ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, null); - } - if (context is null) - { - context = new ExecuteXmlReaderAsyncCallContext(); - } - context.Set(this, source, registration, operationId); - - Task returnedTask = source.Task; - try - { - returnedTask = RegisterForConnectionCloseNotification(returnedTask); - - Task.Factory.FromAsync( - beginMethod: static (AsyncCallback callback, object stateObject) => // with c# 10/NET6 add [StackTraceHidden] to this - { - return ((ExecuteXmlReaderAsyncCallContext)stateObject).Command.BeginExecuteXmlReaderAsync(callback, stateObject); - }, - endMethod: static (IAsyncResult asyncResult) => // with c# 10/NET6 add [StackTraceHidden] to this - { - return ((ExecuteXmlReaderAsyncCallContext)asyncResult.AsyncState).Command.EndExecuteXmlReaderAsync(asyncResult); - }, - state: context - ).ContinueWith( - static (Task task) => - { - ExecuteXmlReaderAsyncCallContext context = (ExecuteXmlReaderAsyncCallContext)task.AsyncState; - SqlCommand command = context.Command; - Guid operationId = context.OperationID; - TaskCompletionSource source = context.TaskCompletionSource; - context.Dispose(); - - command.CleanupAfterExecuteXmlReaderAsync(task, source, operationId); - }, - TaskScheduler.Default - ); - } - catch (Exception e) - { - s_diagnosticListener.WriteCommandError(operationId, this, _transaction, e); - source.SetException(e); - } - - return returnedTask; - } - - private void CleanupAfterExecuteXmlReaderAsync(Task task, TaskCompletionSource source, Guid operationId) - { - if (task.IsFaulted) - { - Exception e = task.Exception.InnerException; - s_diagnosticListener.WriteCommandError(operationId, this, _transaction, e); - source.SetException(e); - } - else - { - if (task.IsCanceled) - { - source.SetCanceled(); - } - else - { - source.SetResult(task.Result); - } - s_diagnosticListener.WriteCommandAfter(operationId, this, _transaction); - } - } - /// public void RegisterColumnEncryptionKeyStoreProvidersOnCommand(IDictionary customProviders) { @@ -2588,7 +2144,7 @@ private SqlDataReader GetParameterEncryptionDataReaderAsync(out Task returnTask, { bool processFinallyBlockAsync = true; bool decrementAsyncCountInFinallyBlockAsync = true; - + try { // Check for any exceptions on network write, before reading. diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj index b4c4654150..4673ae9551 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj @@ -756,6 +756,9 @@ Microsoft\Data\SqlClient\SqlCommand.Scalar.cs + + Microsoft\Data\SqlClient\SqlCommand.Xml.cs + Microsoft\Data\SqlClient\SqlCommandBuilder.cs diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.netfx.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.netfx.cs index 24c02fd0ae..a19511a825 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.netfx.cs +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.netfx.cs @@ -61,30 +61,6 @@ protected override void AfterCleared(SqlCommand owner) } } - internal sealed class ExecuteXmlReaderAsyncCallContext : AAsyncCallContext - { - public Guid OperationID; - - public SqlCommand Command => _owner; - public TaskCompletionSource TaskCompletionSource => _source; - - public void Set(SqlCommand command, TaskCompletionSource source, CancellationTokenRegistration disposable, Guid operationID) - { - base.Set(command, source, disposable); - OperationID = operationID; - } - - protected override void Clear() - { - OperationID = default; - } - - protected override void AfterCleared(SqlCommand owner) - { - owner?.SetCachedCommandExecuteXmlReaderContext(this); - } - } - /// /// Indicates if the column encryption setting was set at-least once in the batch rpc mode, when using AddBatchCommand. /// @@ -606,318 +582,12 @@ private void ThrowIfReconnectionHasBeenCanceled() } } - /// - public XmlReader ExecuteXmlReader() - { - SqlConnection.ExecutePermission.Demand(); - - // Reset _pendingCancel upon entry into any Execute - used to synchronize state - // between entry into Execute* API and the thread obtaining the stateObject. - _pendingCancel = false; - - SqlStatistics statistics = null; - - using (TryEventScope.Create("SqlCommand.ExecuteXmlReader | API | Object Id {0}", ObjectID)) - { - bool success = false; - int? sqlExceptionNumber = null; - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.ExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - - try - { - statistics = SqlStatistics.StartTimer(Statistics); - WriteBeginExecuteEvent(); - - // use the reader to consume metadata - SqlDataReader ds = IsProviderRetriable - ? RunExecuteReaderWithRetry(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true) - : RunExecuteReader(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true); - XmlReader result = CompleteXmlReader(ds); - success = true; - return result; - } - catch (SqlException ex) - { - sqlExceptionNumber = ex.Number; - throw; - } - finally - { - SqlStatistics.StopTimer(statistics); - WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: true); - } - } - } - - /// - [System.Security.Permissions.HostProtectionAttribute(ExternalThreading = true)] - public IAsyncResult BeginExecuteXmlReader() - { - // BeginExecuteXmlReader will track executiontime - return BeginExecuteXmlReader(null, null); - } - - /// - [System.Security.Permissions.HostProtectionAttribute(ExternalThreading = true)] - public IAsyncResult BeginExecuteXmlReader(AsyncCallback callback, object stateObject) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.BeginExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - SqlConnection.ExecutePermission.Demand(); - return BeginExecuteXmlReaderInternal(CommandBehavior.SequentialAccess, callback, stateObject, 0, isRetry: false); - } - - private IAsyncResult BeginExecuteXmlReaderAsync(AsyncCallback callback, object stateObject) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.BeginExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - return BeginExecuteXmlReaderInternal(CommandBehavior.SequentialAccess, callback, stateObject, CommandTimeout, isRetry: false, asyncWrite: true); - } - - private IAsyncResult BeginExecuteXmlReaderInternal(CommandBehavior behavior, AsyncCallback callback, object stateObject, int timeout, bool isRetry, bool asyncWrite = false) - { - TaskCompletionSource globalCompletion = new TaskCompletionSource(stateObject); - TaskCompletionSource localCompletion = new TaskCompletionSource(stateObject); - - if (!isRetry) - { - // Reset _pendingCancel upon entry into any Execute - used to synchronize state - // between entry into Execute* API and the thread obtaining the stateObject. - _pendingCancel = false; - - // Special case - done outside of try/catches to prevent putting a stateObj back into pool when we should not. - ValidateAsyncCommand(); - } - - SqlStatistics statistics = null; - try - { - if (!isRetry) - { - statistics = SqlStatistics.StartTimer(Statistics); - WriteBeginExecuteEvent(); - } - - bool usedCache; - Task writeTask; - try - { - // InternalExecuteNonQuery already has reliability block, but if failure will not put stateObj back into pool. - RunExecuteReader( - behavior, - RunBehavior.ReturnImmediately, - returnStream: true, - localCompletion, - timeout, - out writeTask, - out usedCache, - asyncWrite, - isRetry); - } - catch (Exception e) - { - if (!ADP.IsCatchableOrSecurityExceptionType(e)) - { - // If not catchable - the connection has already been caught and doomed in RunExecuteReader. - throw; - } - - // For async, RunExecuteReader will never put the stateObj back into the pool, so do so now. - ReliablePutStateObject(); - throw; - } - - if (writeTask != null) - { - AsyncHelper.ContinueTaskWithState(writeTask, localCompletion, this, (object state) => ((SqlCommand)state).BeginExecuteXmlReaderInternalReadStage(localCompletion)); - } - else - { - BeginExecuteXmlReaderInternalReadStage(localCompletion); - } - - // When we use query caching for parameter encryption we need to retry on specific errors. - // In these cases finalize the call internally and trigger a retry when needed. - if ( - !TriggerInternalEndAndRetryIfNecessary( - behavior, - stateObject, - timeout, - usedCache, - isRetry, - asyncWrite, - globalCompletion, - localCompletion, - endFunc: static (SqlCommand command, IAsyncResult asyncResult, bool isInternal, string endMethod) => - { - return command.InternalEndExecuteReader(asyncResult, isInternal, endMethod); - }, - retryFunc: static (SqlCommand command, CommandBehavior behavior, AsyncCallback callback, object stateObject, int timeout, bool isRetry, bool asyncWrite) => - { - return command.BeginExecuteXmlReaderInternal(behavior, callback, stateObject, timeout, isRetry, asyncWrite); - }, - endMethod: nameof(EndExecuteXmlReader))) - { - globalCompletion = localCompletion; - } - - // Add callback after work is done to avoid overlapping Begin/End methods - if (callback != null) - { - globalCompletion.Task.ContinueWith((t) => callback(t), TaskScheduler.Default); - } - return globalCompletion.Task; - } - finally - { - SqlStatistics.StopTimer(statistics); - } - } - - private void BeginExecuteXmlReaderInternalReadStage(TaskCompletionSource completion) - { - Debug.Assert(completion != null, "Completion source should not be null"); - // Read SNI does not have catches for async exceptions, handle here. - try - { - // must finish caching information before ReadSni which can activate the callback before returning - CachedAsyncState.SetActiveConnectionAndResult(completion, nameof(EndExecuteXmlReader), _activeConnection); - _stateObj.ReadSni(completion); - } - // @TODO: CER Exception Handling was removed here (see GH#3581) - catch (Exception e) - { - // Similarly, if an exception occurs put the stateObj back into the pool. - // and reset async cache information to allow a second async execute - CachedAsyncState?.ResetAsyncState(); - ReliablePutStateObject(); - completion.TrySetException(e); - } - } - - /// - public XmlReader EndExecuteXmlReader(IAsyncResult asyncResult) - { - try - { - return EndExecuteXmlReaderInternal(asyncResult); - } - finally - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.EndExecuteXmlReader | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - } - } - - private XmlReader EndExecuteXmlReaderAsync(IAsyncResult asyncResult) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.EndExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - Debug.Assert(!_internalEndExecuteInitiated || _stateObj == null); - - Exception asyncException = ((Task)asyncResult).Exception; - if (asyncException != null) - { - CachedAsyncState?.ResetAsyncState(); - ReliablePutStateObject(); - throw asyncException.InnerException; - } - else - { - ThrowIfReconnectionHasBeenCanceled(); - // lock on _stateObj prevents races with close/cancel. - // If we have already initiate the End call internally, we have already done that, so no point doing it again. - if (!_internalEndExecuteInitiated) - { - lock (_stateObj) - { - return EndExecuteXmlReaderInternal(asyncResult); - } - } - else - { - return EndExecuteXmlReaderInternal(asyncResult); - } - } - } - - private XmlReader EndExecuteXmlReaderInternal(IAsyncResult asyncResult) - { - bool success = false; - int? sqlExceptionNumber = null; - try - { - XmlReader result = CompleteXmlReader( - InternalEndExecuteReader(asyncResult, isInternal: false, nameof(EndExecuteXmlReader)), - isAsync: true); - success = true; - return result; - } - catch (SqlException e) - { - sqlExceptionNumber = e.Number; - if (CachedAsyncState != null) - { - CachedAsyncState.ResetAsyncState(); - }; - - // SqlException is always catchable - ReliablePutStateObject(); - throw; - } - catch (Exception e) - { - if (CachedAsyncState != null) - { - CachedAsyncState.ResetAsyncState(); - }; - if (ADP.IsCatchableExceptionType(e)) - { - ReliablePutStateObject(); - }; - throw; - } - finally - { - WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: false); - } - } - - private XmlReader CompleteXmlReader(SqlDataReader ds, bool isAsync = false) - { - XmlReader xr = null; - - SmiExtendedMetaData[] md = ds.GetInternalSmiMetaData(); - bool isXmlCapable = (md != null && md.Length == 1 && (md[0].SqlDbType == SqlDbType.NText - || md[0].SqlDbType == SqlDbType.NVarChar - || md[0].SqlDbType == SqlDbType.Xml)); - - if (isXmlCapable) - { - try - { - SqlStream sqlBuf = new SqlStream(ds, true /*addByteOrderMark*/, (md[0].SqlDbType == SqlDbType.Xml) ? false : true /*process all rows*/); - xr = sqlBuf.ToXmlReader(isAsync); - } - catch (Exception e) - { - if (ADP.IsCatchableExceptionType(e)) - { - ds.Close(); - } - throw; - } - } - if (xr == null) - { - ds.Close(); - throw SQL.NonXmlResult(); - } - return xr; - } - - /// + /// [HostProtection(ExternalThreading = true)] public IAsyncResult BeginExecuteReader() => BeginExecuteReader(callback: null, stateObject: null, CommandBehavior.Default); - /// + /// [HostProtection(ExternalThreading = true)] public IAsyncResult BeginExecuteReader(AsyncCallback callback, object stateObject) => BeginExecuteReader(callback, stateObject, CommandBehavior.Default); @@ -1560,120 +1230,6 @@ private void SetCachedCommandExecuteReaderAsyncContext(ExecuteReaderAsyncCallCon } } - private void SetCachedCommandExecuteXmlReaderContext(ExecuteXmlReaderAsyncCallContext instance) - { - if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) - { - Interlocked.CompareExchange(ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, instance, null); - } - } - - /// - public Task ExecuteXmlReaderAsync() => - ExecuteXmlReaderAsync(CancellationToken.None); - - /// - public Task ExecuteXmlReaderAsync(CancellationToken cancellationToken) => - IsProviderRetriable - ? InternalExecuteXmlReaderWithRetryAsync(cancellationToken) - : InternalExecuteXmlReaderAsync(cancellationToken); - - private Task InternalExecuteXmlReaderWithRetryAsync(CancellationToken cancellationToken) => - RetryLogicProvider.ExecuteAsync( - sender: this, - () => InternalExecuteXmlReaderAsync(cancellationToken), - cancellationToken); - - private Task InternalExecuteXmlReaderAsync(CancellationToken cancellationToken) - { - SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.InternalExecuteXmlReaderAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); - SqlConnection.ExecutePermission.Demand(); - Guid operationId = Guid.Empty; - - // connection can be used as state in RegisterForConnectionCloseNotification continuation - // to avoid an allocation so use it as the state value if possible but it can be changed if - // you need it for a more important piece of data that justifies the tuple allocation later - TaskCompletionSource source = new TaskCompletionSource(_activeConnection); - - CancellationTokenRegistration registration = new CancellationTokenRegistration(); - if (cancellationToken.CanBeCanceled) - { - if (cancellationToken.IsCancellationRequested) - { - source.SetCanceled(); - return source.Task; - } - registration = cancellationToken.Register(s_cancelIgnoreFailure, this); - } - - ExecuteXmlReaderAsyncCallContext context = null; - if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) - { - context = Interlocked.Exchange(ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, null); - } - if (context is null) - { - context = new ExecuteXmlReaderAsyncCallContext(); - } - context.Set(this, source, registration, operationId); - - Task returnedTask = source.Task; - try - { - returnedTask = RegisterForConnectionCloseNotification(returnedTask); - - Task.Factory.FromAsync( - beginMethod: static (AsyncCallback callback, object stateObject) => - { - return ((ExecuteXmlReaderAsyncCallContext)stateObject).Command.BeginExecuteXmlReaderAsync(callback, stateObject); - }, - endMethod: static (IAsyncResult asyncResult) => - { - return ((ExecuteXmlReaderAsyncCallContext)asyncResult.AsyncState).Command.EndExecuteXmlReaderAsync(asyncResult); - }, - state: context - ).ContinueWith( - static (Task task) => - { - ExecuteXmlReaderAsyncCallContext context = (ExecuteXmlReaderAsyncCallContext)task.AsyncState; - SqlCommand command = context.Command; - Guid operationId = context.OperationID; - TaskCompletionSource source = context.TaskCompletionSource; - context.Dispose(); - - command.CleanupAfterExecuteXmlReaderAsync(task, source, operationId); - }, - TaskScheduler.Default - ); - } - catch (Exception e) - { - source.SetException(e); - } - - return returnedTask; - } - - private void CleanupAfterExecuteXmlReaderAsync(Task task, TaskCompletionSource source, Guid operationId) - { - if (task.IsFaulted) - { - Exception e = task.Exception.InnerException; - source.SetException(e); - } - else - { - if (task.IsCanceled) - { - source.SetCanceled(); - } - else - { - source.SetResult(task.Result); - } - } - } - /// public void RegisterColumnEncryptionKeyStoreProvidersOnCommand(IDictionary customProviders) { diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs index 1377666777..af30865850 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs @@ -226,7 +226,7 @@ private IAsyncResult BeginExecuteNonQueryInternal( task: execNonQuery, completion: localCompletion, state: Tuple.Create(this, localCompletion), - onSuccess: state => + onSuccess: static state => { var parameters = (Tuple>)state; parameters.Item1.BeginExecuteNonQueryInternalReadStage(parameters.Item2); @@ -253,8 +253,6 @@ private IAsyncResult BeginExecuteNonQueryInternal( throw; } - // When we use query caching for parameter encryption we need to retry on specific errors. - // In these cases finalize the call internally and trigger a retry when needed. // When we use query caching for parameter encryption we need to retry on specific errors. // In these cases finalize the call internally and trigger a retry when needed. // @TODO: store this method call in a variable, it's faaaaar too big to be used in an if statement @@ -276,18 +274,17 @@ private IAsyncResult BeginExecuteNonQueryInternal( { return command.BeginExecuteNonQueryInternal(behavior, callback, stateObject, timeout, isRetry, asyncWrite); }, - nameof(EndExecuteNonQuery))) + endMethod: nameof(EndExecuteNonQuery))) { globalCompletion = localCompletion; } // Add callback after work is done to avoid overlapping Begin/End methods - if (callback != null) + if (callback is not null) { globalCompletion.Task.ContinueWith( static (task, state) => ((AsyncCallback)state)(task), - state: callback - ); + state: callback); } return globalCompletion.Task; diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs new file mode 100644 index 0000000000..3391c300d7 --- /dev/null +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs @@ -0,0 +1,610 @@ +// 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; +using System.Diagnostics; +using System.Threading; +using System.Threading.Tasks; +using System.Xml; +using Microsoft.Data.Common; +using Microsoft.Data.SqlClient.Server; + +#if NETFRAMEWORK +using System.Security.Permissions; +#endif + +namespace Microsoft.Data.SqlClient +{ + /// + /// This partial contains the methods related to ExecuteXmlReader public API. + /// + public sealed partial class SqlCommand + { + #region Public/Internal Methods + + /// + #if NETFRAMEWORK + [HostProtection(ExternalThreading = true)] + #endif + public IAsyncResult BeginExecuteXmlReader() => + BeginExecuteXmlReader(callback: null, stateObject: null); + + /// + #if NETFRAMEWORK + [HostProtection(ExternalThreading = true)] + #endif + public IAsyncResult BeginExecuteXmlReader(AsyncCallback callback, object stateObject) + { + #if NETFRAMEWORK + SqlConnection.ExecutePermission.Demand(); + #endif + + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.BeginExecuteXmlReader | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection?.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + + return BeginExecuteXmlReaderInternal( + CommandBehavior.SequentialAccess, + callback, + stateObject, + timeout: 0, + isRetry: false, + asyncWrite: false); + } + + /// + public XmlReader EndExecuteXmlReader(IAsyncResult asyncResult) + { + try + { + return EndExecuteXmlReaderInternal(asyncResult); + } + finally + { + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.EndExecuteXmlReader | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + } + } + + /// + public XmlReader ExecuteXmlReader() + { + #if NETFRAMEWORK + SqlConnection.ExecutePermission.Demand(); + #endif + + // Reset _pendingCancel upon entry into any Execute - used to synchronize state + // between entry into Execute* API and the thread obtaining the stateObject. + _pendingCancel = false; + + #if NET + using var diagnosticScope = s_diagnosticListener.CreateCommandScope(this, _transaction); + #endif + + using var eventScope = TryEventScope.Create($"SqlCommand.ExecuteXmlReader | API | Object Id {ObjectID}"); + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.ExecuteXmlReader | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection?.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + + SqlStatistics statistics = null; + bool success = false; + int? sqlExceptionNumber = null; + try + { + statistics = SqlStatistics.StartTimer(Statistics); + WriteBeginExecuteEvent(); + + // Use the reader to consume metadata + SqlDataReader reader = IsProviderRetriable + ? RunExecuteReaderWithRetry(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true) + : RunExecuteReader(CommandBehavior.SequentialAccess, RunBehavior.ReturnImmediately, returnStream: true); + success = true; + + return CompleteXmlReader(reader, isAsync: false); + } + catch (Exception ex) + { + #if NET + diagnosticScope.SetException(ex); + #endif + + if (ex is SqlException sqlException) + { + sqlExceptionNumber = sqlException.Number; + } + + throw; + } + finally + { + SqlStatistics.StopTimer(statistics); + WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: true); + } + } + + /// + public Task ExecuteXmlReaderAsync() => + ExecuteXmlReaderAsync(CancellationToken.None); + + /// + public Task ExecuteXmlReaderAsync(CancellationToken cancellationToken) => + IsProviderRetriable + ? InternalExecuteXmlReaderWithRetryAsync(cancellationToken) + : InternalExecuteXmlReaderAsync(cancellationToken); + + #endregion + + #region Private Methods + + private static XmlReader CompleteXmlReader(SqlDataReader dataReader, bool isAsync) + { + XmlReader xmlReader = null; + + SmiExtendedMetaData[] metaData = dataReader.GetInternalSmiMetaData(); + bool isXmlCapable = metaData?.Length == 1 && + metaData[0].SqlDbType is SqlDbType.NVarChar or SqlDbType.NText or SqlDbType.Xml; + + if (isXmlCapable) + { + try + { + SqlStream sqlStream = new SqlStream( + dataReader, + addByteOrderMark: true, + processAllRows: metaData[0].SqlDbType is not SqlDbType.Xml); + xmlReader = sqlStream.ToXmlReader(isAsync); + } + catch (Exception e) + { + if (ADP.IsCatchableExceptionType(e)) + { + dataReader.Close(); + } + + throw; + } + } + + if (xmlReader is null) + { + dataReader.Close(); + throw SQL.NonXmlResult(); + } + + return xmlReader; + } + + private IAsyncResult BeginExecuteXmlReaderAsync(AsyncCallback callback, object stateObject) + { + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.BeginExecuteXmlReaderAsync | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection?.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + + return BeginExecuteXmlReaderInternal( + CommandBehavior.SequentialAccess, + callback, + stateObject, + CommandTimeout, + isRetry: false, + asyncWrite: true); + } + + private IAsyncResult BeginExecuteXmlReaderInternal( + CommandBehavior behavior, + AsyncCallback callback, + object stateObject, + int timeout, + bool isRetry, + bool asyncWrite) + { + TaskCompletionSource globalCompletion = new TaskCompletionSource(stateObject); + TaskCompletionSource localCompletion = new TaskCompletionSource(stateObject); + + if (!isRetry) + { + // Reset _pendingCancel upon entry into any Execute - used to synchronize state + // between entry into Execute* API and the thread obtaining the stateObject. + _pendingCancel = false; + + // Special case - done outside of try/catches to prevent putting a stateObj back + // into pool when we should not. + ValidateAsyncCommand(); + } + + SqlStatistics statistics = null; + try + { + if (!isRetry) + { + statistics = SqlStatistics.StartTimer(Statistics); + WriteBeginExecuteEvent(); + } + + bool usedCache; // @TODO: Is this used *usefully* anywhere? + Task writeTask; + try + { + // RunExecuteReader already has reliability block, but if it fails, it + // will not put stateObj back into pool. + RunExecuteReader( + behavior, + RunBehavior.ReturnImmediately, + returnStream: true, + localCompletion, + timeout, + out writeTask, + out usedCache, + asyncWrite, + isRetry); + + // @TODO: NonQuery pathway has the continueTaskWithState block inside this try. One or the other seems wrong + } + catch (Exception e) + { + // @TODO: Invert + if (!ADP.IsCatchableOrSecurityExceptionType(e)) + { + // If not catchable - the connection has already been caught and doomed in + // RunExecuteReader. + throw; + } + + // For async, RunExecuteReader will never put the stateObj back into the pool, + // so, do so now. + ReliablePutStateObject(); + throw; + } + + if (writeTask is not null) + { + AsyncHelper.ContinueTaskWithState( + task: writeTask, + completion: localCompletion, + state: Tuple.Create(this, localCompletion), + onSuccess: static state => + { + var parameters = (Tuple>)state; + parameters.Item1.BeginExecuteXmlReaderInternalReadStage(parameters.Item2); + }); + } + else + { + BeginExecuteXmlReaderInternalReadStage(localCompletion); + } + + // When we use query caching for parameter encryption we need to retry on specific + // errors. In these cases finalize the call internally and trigger a retry when needed. + // @TODO: store this method call in a variable, it's faaaaar too big to be used in an if statement + if ( + !TriggerInternalEndAndRetryIfNecessary( + behavior, + stateObject, + timeout, + usedCache, + isRetry, + asyncWrite, + globalCompletion, + localCompletion, + endFunc: static (SqlCommand command, IAsyncResult asyncResult, bool isInternal, string endMethod) => + { + return command.InternalEndExecuteReader(asyncResult, isInternal, endMethod); + }, + retryFunc: static (SqlCommand command, CommandBehavior behavior, AsyncCallback callback, object stateObject, int timeout, bool isRetry, bool asyncWrite) => + { + return command.BeginExecuteXmlReaderInternal(behavior, callback, stateObject, timeout, isRetry, asyncWrite); + }, + endMethod: nameof(EndExecuteXmlReader))) + { + // @TODO: globalCompletion isn't being used past here! Why are we doing this?? + globalCompletion = localCompletion; + } + + // Add callback after work is done to avoid overlapping Begin/End methods + if (callback is not null) + { + localCompletion.Task.ContinueWith( + static (task, state) => ((AsyncCallback)state)(task), + state: callback); + } + + return localCompletion.Task; + } + finally + { + SqlStatistics.StopTimer(statistics); + } + } + + private void BeginExecuteXmlReaderInternalReadStage(TaskCompletionSource completion) + { + // Read SNI does not have catches for async exceptions, handle here. + + try + { + // Must finish caching information before ReadSni which can activate the callback + // before returning. + CachedAsyncState.SetActiveConnectionAndResult( + completion, + nameof(EndExecuteXmlReader), + _activeConnection); + _stateObj.ReadSni(completion); + } + // @TODO: CER Exception Handling was removed here (see GH#3581) + catch (Exception e) + { + // Similarly, if an exception occurs put the stateObj back into the pool. + // and reset async cache information to allow a second async execute + CachedAsyncState?.ResetAsyncState(); + ReliablePutStateObject(); + completion.TrySetException(e); + } + } + + // @TODO: This is basically identical to non-query (and probably the reader) versions. Can be refactored to be generic. + private void CleanupAfterExecuteXmlReaderAsync( + Task task, + TaskCompletionSource source, + Guid operationId) + { + if (task.IsFaulted) + { + Exception e = task.Exception?.InnerException; + + #if NET + s_diagnosticListener.WriteCommandError(operationId, this, _transaction, e); + #endif + + source.SetException(e); + } + else if (task.IsCanceled) + { + #if NET + s_diagnosticListener.WriteCommandAfter(operationId, this, _transaction); + #endif + + source.SetCanceled(); + } + else + { + #if NET + s_diagnosticListener.WriteCommandAfter(operationId, this, _transaction); + #endif + + source.SetResult(task.Result); + } + } + + private XmlReader EndExecuteXmlReaderAsync(IAsyncResult asyncResult) + { + Debug.Assert(!_internalEndExecuteInitiated || _stateObj is null); + + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.EndExecuteXmlReaderAsync | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection?.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + + Exception asyncException = ((Task)asyncResult).Exception; + if (asyncException is not null) + { + CachedAsyncState?.ResetAsyncState(); + ReliablePutStateObject(); + throw asyncException.InnerException; + } + + ThrowIfReconnectionHasBeenCanceled(); + + // Locking _stateObj prevents races with close/cancel. + // If we have already initiated the End call internally, we have already done that, so + // no point doing it again. + if (!_internalEndExecuteInitiated) + { + lock (_stateObj) + { + return EndExecuteXmlReaderInternal(asyncResult); + } + } + + return EndExecuteXmlReaderInternal(asyncResult); + } + + private XmlReader EndExecuteXmlReaderInternal(IAsyncResult asyncResult) + { + bool success = false; + int? sqlExceptionNumber = null; + + try + { + SqlDataReader dataReader = InternalEndExecuteReader( + asyncResult, + isInternal: false, + endMethod: nameof(EndExecuteXmlReader)); + XmlReader result = CompleteXmlReader(dataReader, isAsync: true); + + success = true; + return result; + } + catch (Exception e) + { + if (e is SqlException sqlException) + { + sqlExceptionNumber = sqlException.Number; + } + + CachedAsyncState?.ResetAsyncState(); + + if (ADP.IsCatchableExceptionType(e)) + { + ReliablePutStateObject(); + } + + throw; + } + finally + { + WriteEndExecuteEvent(success, sqlExceptionNumber, synchronous: false); + } + } + + private Task InternalExecuteXmlReaderAsync(CancellationToken cancellationToken) + { + #if NETFRAMEWORK + SqlConnection.ExecutePermission.Demand(); + #endif + + SqlClientEventSource.Log.TryCorrelationTraceEvent( + "SqlCommand.InternalExecuteXmlReaderAsync | API | Correlation | " + + $"Object Id {ObjectID}, " + + $"Activity Id {ActivityCorrelator.Current}, " + + $"Client Connection Id {_activeConnection?.ClientConnectionId}, " + + $"Command Text '{CommandText}'"); + + #if NET + Guid operationId = s_diagnosticListener.WriteCommandBefore(this, _transaction); + #else + Guid operationId = Guid.Empty; + #endif + + // Connection can be used as state in RegisterForConnectionCloseNotification continuation + // to avoid an allocation so use it as the state value if possible but it can be changed if + // you need it for a more important piece of data that justifies the tuple allocation later + TaskCompletionSource source = new TaskCompletionSource(_activeConnection); + + CancellationTokenRegistration registration = new CancellationTokenRegistration(); + if (cancellationToken.CanBeCanceled) + { + if (cancellationToken.IsCancellationRequested) + { + source.SetCanceled(); + return source.Task; + } + + registration = cancellationToken.Register(callback: s_cancelIgnoreFailure, state: this); + } + + // @TODO: This can be cleaned up to lines if InnerConnection is always SqlInternalConnection + ExecuteXmlReaderAsyncCallContext context = null; + if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) + { + context = Interlocked.Exchange( + ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, + null); + } + + context ??= new ExecuteXmlReaderAsyncCallContext(); + context.Set(this, source, registration, operationId); + + Task returnedTask = source.Task; + try + { + returnedTask = RegisterForConnectionCloseNotification(returnedTask); + + // @TODO: Replace with native async implementation, make Begin/End implementation rely on async implementation + Task.Factory.FromAsync( + beginMethod: static (callback, stateObject) => + { + // @TODO: With C# 10/net6 add [StackTraceHidden] + return ((ExecuteXmlReaderAsyncCallContext)stateObject).Command.BeginExecuteXmlReaderAsync( + callback, + stateObject); + }, + endMethod: static asyncResult => + { + // @TODO: With C# 10/net6 add [StackTraceHidden] + return ((ExecuteXmlReaderAsyncCallContext)asyncResult.AsyncState).Command.EndExecuteXmlReaderAsync( + asyncResult); + }, + state: context + ).ContinueWith(static task => + { + // @TODO: With C#/net6 add [StackTraceHidden] + ExecuteXmlReaderAsyncCallContext context = (ExecuteXmlReaderAsyncCallContext)task.AsyncState; + SqlCommand command = context.Command; + Guid operationId = context.OperationId; + TaskCompletionSource source = context.TaskCompletionSource; + + context.Dispose(); + + command.CleanupAfterExecuteXmlReaderAsync(task, source, operationId); + }, + scheduler: TaskScheduler.Default); + } + catch (Exception e) + { + #if NET + s_diagnosticListener.WriteCommandError(operationId, this, _transaction, e); + #endif + + source.SetException(e); + } + + return returnedTask; + } + + private Task InternalExecuteXmlReaderWithRetryAsync(CancellationToken cancellationToken) => + RetryLogicProvider.ExecuteAsync( + sender: this, + () => InternalExecuteXmlReaderAsync(cancellationToken), + cancellationToken); + + private void SetCachedCommandExecuteXmlReaderContext(ExecuteXmlReaderAsyncCallContext instance) + { + if (_activeConnection?.InnerConnection is SqlInternalConnection sqlInternalConnection) + { + // @TODO: Move this compare exchange into the SqlInternalConnection class (or better yet, do away with this context) + Interlocked.CompareExchange( + ref sqlInternalConnection.CachedCommandExecuteXmlReaderAsyncContext, + instance, + comparand: null); + } + } + + #endregion + + internal sealed class ExecuteXmlReaderAsyncCallContext + : AAsyncCallContext + { + public SqlCommand Command => _owner; + + public Guid OperationId { get; set; } + + public TaskCompletionSource TaskCompletionSource => _source; + + public void Set( + SqlCommand command, + TaskCompletionSource source, + CancellationTokenRegistration disposable, + Guid operationId) + { + base.Set(command,source, disposable); + OperationId = operationId; + } + + protected override void AfterCleared(SqlCommand owner) + { + owner?.SetCachedCommandExecuteXmlReaderContext(this); + } + + protected override void Clear() + { + OperationId = Guid.Empty; + } + } + } +}