diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs index 16d707f3fd..e7f43765f1 100644 --- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs +++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs @@ -2460,7 +2460,10 @@ private Task InternalExecuteNonQueryAsync(CancellationToken cancellationTok SqlClientEventSource.Log.TryCorrelationTraceEvent("SqlCommand.InternalExecuteNonQueryAsync | API | Correlation | Object Id {0}, Activity Id {1}, Client Connection Id {2}, Command Text '{3}'", ObjectID, ActivityCorrelator.Current, Connection?.ClientConnectionId, CommandText); Guid operationId = _diagnosticListener.WriteCommandBefore(this, _transaction); - TaskCompletionSource source = new TaskCompletionSource(); + // 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) @@ -2565,7 +2568,10 @@ private Task InternalExecuteReaderAsync(CommandBehavior behavior, operationId = _diagnosticListener.WriteCommandBefore(this, _transaction); } - TaskCompletionSource source = new TaskCompletionSource(); + // 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) @@ -2737,7 +2743,10 @@ private Task InternalExecuteXmlReaderAsync(CancellationToken cancella 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 = _diagnosticListener.WriteCommandBefore(this, _transaction); - TaskCompletionSource source = new TaskCompletionSource(); + // 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) diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs index 4879a3111c..19f72d14b6 100644 --- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs +++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs @@ -2151,17 +2151,44 @@ private static void ChangePassword(string connectionString, SqlConnectionString internal Task RegisterForConnectionCloseNotification(Task outerTask, object value, int tag) { // Connection exists, schedule removal, will be added to ref collection after calling ValidateAndReconnect + + object state = null; + if (outerTask.AsyncState == this) + { + // if the caller created the TaskCompletionSource for outerTask with this connection + // as the state parameter (which is immutable) we can use task.AsyncState and state + // to carry the two pieces of state that we need into the continuation avoiding the + // allocation of a new state object to carry them + state = value; + } + else + { + // otherwise we need to create a Tuple to carry the two pieces of state + state = Tuple.Create(this, value); + } + return outerTask.ContinueWith( continuationFunction: static (task, state) => { - Tuple parameters = (Tuple)state; - SqlConnection connection = parameters.Item1; - object obj = parameters.Item2; + SqlConnection connection = null; + object obj = null; + if (state is Tuple tuple) + { + // special state tuple, unpack it + connection = tuple.Item1; + obj = tuple.Item2; + } + else + { + // use state on task and state object + connection = (SqlConnection)task.AsyncState; + obj = state; + } connection.RemoveWeakReference(obj); return task; }, - state: Tuple.Create(this, value), + state: state, scheduler: TaskScheduler.Default ).Unwrap(); } diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs index 3a2f91e7e7..1dc63f61bc 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs @@ -2310,7 +2310,7 @@ private Task ReadWriteColumnValueAsync(int col) return writeTask; } - private void RegisterForConnectionCloseNotification(ref Task outterTask) + private Task RegisterForConnectionCloseNotification(Task outterTask) { SqlConnection connection = _connection; if (connection == null) @@ -2319,7 +2319,7 @@ private void RegisterForConnectionCloseNotification(ref Task outterTask) throw ADP.ClosedConnectionError(); } - connection.RegisterForConnectionCloseNotification(ref outterTask, this, SqlReferenceCollection.BulkCopyTag); + return connection.RegisterForConnectionCloseNotification(outterTask, this, SqlReferenceCollection.BulkCopyTag); } // Runs a loop to copy all columns of a single row. @@ -3139,9 +3139,7 @@ private Task WriteToServerInternalAsync(CancellationToken ctoken) if (_isAsyncBulkCopy) { source = new TaskCompletionSource(); // Creating the completion source/Task that we pass to application - resultTask = source.Task; - - RegisterForConnectionCloseNotification(ref resultTask); + resultTask = RegisterForConnectionCloseNotification(source.Task); } if (_destinationTableName == null) diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs index a4de7ead2b..8df8db168e 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs @@ -2935,7 +2935,10 @@ private Task InternalExecuteNonQueryAsync(CancellationToken cancellationTok SqlClientEventSource.Log.TryCorrelationTraceEvent(" ObjectID {0}, ActivityID {1}", ObjectID, ActivityCorrelator.Current); SqlConnection.ExecutePermission.Demand(); - TaskCompletionSource source = new TaskCompletionSource(); + // 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) @@ -2951,7 +2954,7 @@ private Task InternalExecuteNonQueryAsync(CancellationToken cancellationTok Task returnedTask = source.Task; try { - RegisterForConnectionCloseNotification(ref returnedTask); + returnedTask = RegisterForConnectionCloseNotification(returnedTask); Task.Factory.FromAsync(BeginExecuteNonQueryAsync, EndExecuteNonQueryAsync, null).ContinueWith((t) => { @@ -3032,7 +3035,10 @@ private Task InternalExecuteReaderAsync(CommandBehavior behavior, SqlClientEventSource.Log.TryCorrelationTraceEvent(" ObjectID {0}, behavior={1}, ActivityID {2}", ObjectID, (int)behavior, ActivityCorrelator.Current); SqlConnection.ExecutePermission.Demand(); - TaskCompletionSource source = new TaskCompletionSource(); + // 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) @@ -3048,7 +3054,7 @@ private Task InternalExecuteReaderAsync(CommandBehavior behavior, Task returnedTask = source.Task; try { - RegisterForConnectionCloseNotification(ref returnedTask); + returnedTask = RegisterForConnectionCloseNotification(returnedTask); Task.Factory.FromAsync(BeginExecuteReaderAsync, EndExecuteReaderAsync, behavior, null).ContinueWith((t) => { @@ -3182,7 +3188,10 @@ private Task InternalExecuteXmlReaderAsync(CancellationToken cancella SqlClientEventSource.Log.TryCorrelationTraceEvent(" ObjectID {0}, ActivityID {1}", ObjectID, ActivityCorrelator.Current); SqlConnection.ExecutePermission.Demand(); - TaskCompletionSource source = new TaskCompletionSource(); + // 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) @@ -3198,7 +3207,7 @@ private Task InternalExecuteXmlReaderAsync(CancellationToken cancella Task returnedTask = source.Task; try { - RegisterForConnectionCloseNotification(ref returnedTask); + returnedTask = RegisterForConnectionCloseNotification(returnedTask); Task.Factory.FromAsync(BeginExecuteXmlReaderAsync, EndExecuteXmlReaderAsync, null).ContinueWith((t) => { @@ -5771,7 +5780,7 @@ object ICloneable.Clone() return Clone(); } - private void RegisterForConnectionCloseNotification(ref Task outterTask) + private Task RegisterForConnectionCloseNotification(Task outterTask) { SqlConnection connection = _activeConnection; if (connection == null) @@ -5780,7 +5789,7 @@ private void RegisterForConnectionCloseNotification(ref Task outterTask) throw ADP.ClosedConnectionError(); } - connection.RegisterForConnectionCloseNotification(ref outterTask, this, SqlReferenceCollection.CommandTag); + return connection.RegisterForConnectionCloseNotification(outterTask, this, SqlReferenceCollection.CommandTag); } // validates that a command has commandText and a non-busy open connection diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs index cf77f152d6..655f74a2ad 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs @@ -2819,16 +2819,52 @@ private static void ChangePassword(string connectionString, SqlConnectionString SqlConnectionFactory.SingletonInstance.ClearPool(key); } - internal void RegisterForConnectionCloseNotification(ref Task outerTask, object value, int tag) + internal Task RegisterForConnectionCloseNotification(Task outerTask, object value, int tag) { // Connection exists, schedule removal, will be added to ref collection after calling ValidateAndReconnect - outerTask = outerTask.ContinueWith(task => + + object state = null; + if (outerTask.AsyncState == this) + { + // if the caller created the TaskCompletionSource for outerTask with this connection + // as the state parameter (which is immutable) we can use task.AsyncState and state + // to carry the two pieces of state that we need into the continuation avoiding the + // allocation of a new state object to carry them + state = value; + } + else { - RemoveWeakReference(value); - return task; - }, TaskScheduler.Default).Unwrap(); + // otherwise we need to create a Tuple to carry the two pieces of state + state = Tuple.Create(this, value); + } + + return outerTask.ContinueWith( + continuationFunction: static (task, state) => + { + SqlConnection connection = null; + object obj = null; + if (state is Tuple tuple) + { + // special state tuple, unpack it + connection = tuple.Item1; + obj = tuple.Item2; + } + else + { + // use state on task and state object + connection = (SqlConnection)task.AsyncState; + obj = state; + } + + connection.RemoveWeakReference(obj); + return task; + }, + state: state, + scheduler: TaskScheduler.Default + ).Unwrap(); } + // updates our context with any changes made to the memory-mapped data by an external process static private void RefreshMemoryMappedData(SqlDebugContext sdc) {