Skip to content

Commit ac8244a

Browse files
Wraith2Wraith2
authored
Perf: add allocation free path for close notifications when task state is setup correctly (#1198)
1 parent a6e7e0d commit ac8244a

File tree

5 files changed

+104
-25
lines changed

5 files changed

+104
-25
lines changed

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2550,7 +2550,10 @@ private Task<int> InternalExecuteNonQueryAsync(CancellationToken cancellationTok
25502550
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);
25512551
Guid operationId = s_diagnosticListener.WriteCommandBefore(this, _transaction);
25522552

2553-
TaskCompletionSource<int> source = new TaskCompletionSource<int>();
2553+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
2554+
// to avoid an allocation so use it as the state value if possible but it can be changed if
2555+
// you need it for a more important piece of data that justifies the tuple allocation later
2556+
TaskCompletionSource<int> source = new TaskCompletionSource<int>(_activeConnection);
25542557

25552558
CancellationTokenRegistration registration = new CancellationTokenRegistration();
25562559
if (cancellationToken.CanBeCanceled)
@@ -2674,7 +2677,10 @@ private Task<SqlDataReader> InternalExecuteReaderAsync(CommandBehavior behavior,
26742677
operationId = s_diagnosticListener.WriteCommandBefore(this, _transaction);
26752678
}
26762679

2677-
TaskCompletionSource<SqlDataReader> source = new TaskCompletionSource<SqlDataReader>();
2680+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
2681+
// to avoid an allocation so use it as the state value if possible but it can be changed if
2682+
// you need it for a more important piece of data that justifies the tuple allocation later
2683+
TaskCompletionSource<SqlDataReader> source = new TaskCompletionSource<SqlDataReader>(_activeConnection);
26782684

26792685
CancellationTokenRegistration registration = new CancellationTokenRegistration();
26802686
if (cancellationToken.CanBeCanceled)
@@ -2846,7 +2852,10 @@ private Task<XmlReader> InternalExecuteXmlReaderAsync(CancellationToken cancella
28462852
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);
28472853
Guid operationId = s_diagnosticListener.WriteCommandBefore(this, _transaction);
28482854

2849-
TaskCompletionSource<XmlReader> source = new TaskCompletionSource<XmlReader>();
2855+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
2856+
// to avoid an allocation so use it as the state value if possible but it can be changed if
2857+
// you need it for a more important piece of data that justifies the tuple allocation later
2858+
TaskCompletionSource<XmlReader> source = new TaskCompletionSource<XmlReader>(_activeConnection);
28502859

28512860
CancellationTokenRegistration registration = new CancellationTokenRegistration();
28522861
if (cancellationToken.CanBeCanceled)

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

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2231,17 +2231,44 @@ private static void ChangePassword(string connectionString, SqlConnectionString
22312231
internal Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outerTask, object value, int tag)
22322232
{
22332233
// Connection exists, schedule removal, will be added to ref collection after calling ValidateAndReconnect
2234+
2235+
object state = null;
2236+
if (outerTask.AsyncState == this)
2237+
{
2238+
// if the caller created the TaskCompletionSource for outerTask with this connection
2239+
// as the state parameter (which is immutable) we can use task.AsyncState and state
2240+
// to carry the two pieces of state that we need into the continuation avoiding the
2241+
// allocation of a new state object to carry them
2242+
state = value;
2243+
}
2244+
else
2245+
{
2246+
// otherwise we need to create a Tuple to carry the two pieces of state
2247+
state = Tuple.Create(this, value);
2248+
}
2249+
22342250
return outerTask.ContinueWith(
22352251
continuationFunction: static (task, state) =>
22362252
{
2237-
Tuple<SqlConnection, object> parameters = (Tuple<SqlConnection, object>)state;
2238-
SqlConnection connection = parameters.Item1;
2239-
object obj = parameters.Item2;
2253+
SqlConnection connection = null;
2254+
object obj = null;
2255+
if (state is Tuple<SqlConnection, object> tuple)
2256+
{
2257+
// special state tuple, unpack it
2258+
connection = tuple.Item1;
2259+
obj = tuple.Item2;
2260+
}
2261+
else
2262+
{
2263+
// use state on task and state object
2264+
connection = (SqlConnection)task.AsyncState;
2265+
obj = state;
2266+
}
22402267

22412268
connection.RemoveWeakReference(obj);
22422269
return task;
22432270
},
2244-
state: Tuple.Create(this, value),
2271+
state: state,
22452272
scheduler: TaskScheduler.Default
22462273
).Unwrap();
22472274
}

‎src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2310,7 +2310,7 @@ private Task ReadWriteColumnValueAsync(int col)
23102310
return writeTask;
23112311
}
23122312

2313-
private void RegisterForConnectionCloseNotification<T>(ref Task<T> outterTask)
2313+
private Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outterTask)
23142314
{
23152315
SqlConnection connection = _connection;
23162316
if (connection == null)
@@ -2319,7 +2319,7 @@ private void RegisterForConnectionCloseNotification<T>(ref Task<T> outterTask)
23192319
throw ADP.ClosedConnectionError();
23202320
}
23212321

2322-
connection.RegisterForConnectionCloseNotification<T>(ref outterTask, this, SqlReferenceCollection.BulkCopyTag);
2322+
return connection.RegisterForConnectionCloseNotification(outterTask, this, SqlReferenceCollection.BulkCopyTag);
23232323
}
23242324

23252325
// Runs a loop to copy all columns of a single row.
@@ -3139,9 +3139,7 @@ private Task WriteToServerInternalAsync(CancellationToken ctoken)
31393139
if (_isAsyncBulkCopy)
31403140
{
31413141
source = new TaskCompletionSource<object>(); // Creating the completion source/Task that we pass to application
3142-
resultTask = source.Task;
3143-
3144-
RegisterForConnectionCloseNotification(ref resultTask);
3142+
resultTask = RegisterForConnectionCloseNotification(source.Task);
31453143
}
31463144

31473145
if (_destinationTableName == null)

‎src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2964,7 +2964,10 @@ private Task<int> InternalExecuteNonQueryAsync(CancellationToken cancellationTok
29642964
SqlClientEventSource.Log.TryCorrelationTraceEvent("<sc.SqlCommand.ExecuteNonQueryAsync|API|Correlation> ObjectID {0}, ActivityID {1}", ObjectID, ActivityCorrelator.Current);
29652965
SqlConnection.ExecutePermission.Demand();
29662966

2967-
TaskCompletionSource<int> source = new TaskCompletionSource<int>();
2967+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
2968+
// to avoid an allocation so use it as the state value if possible but it can be changed if
2969+
// you need it for a more important piece of data that justifies the tuple allocation later
2970+
TaskCompletionSource<int> source = new TaskCompletionSource<int>(_activeConnection);
29682971

29692972
CancellationTokenRegistration registration = new CancellationTokenRegistration();
29702973
if (cancellationToken.CanBeCanceled)
@@ -2980,7 +2983,7 @@ private Task<int> InternalExecuteNonQueryAsync(CancellationToken cancellationTok
29802983
Task<int> returnedTask = source.Task;
29812984
try
29822985
{
2983-
RegisterForConnectionCloseNotification(ref returnedTask);
2986+
returnedTask = RegisterForConnectionCloseNotification(returnedTask);
29842987

29852988
Task<int>.Factory.FromAsync(BeginExecuteNonQueryAsync, EndExecuteNonQueryAsync, null).ContinueWith((t) =>
29862989
{
@@ -3061,7 +3064,10 @@ private Task<SqlDataReader> InternalExecuteReaderAsync(CommandBehavior behavior,
30613064
SqlClientEventSource.Log.TryCorrelationTraceEvent("<sc.SqlCommand.ExecuteReaderAsync|API|Correlation> ObjectID {0}, behavior={1}, ActivityID {2}", ObjectID, (int)behavior, ActivityCorrelator.Current);
30623065
SqlConnection.ExecutePermission.Demand();
30633066

3064-
TaskCompletionSource<SqlDataReader> source = new TaskCompletionSource<SqlDataReader>();
3067+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
3068+
// to avoid an allocation so use it as the state value if possible but it can be changed if
3069+
// you need it for a more important piece of data that justifies the tuple allocation later
3070+
TaskCompletionSource<SqlDataReader> source = new TaskCompletionSource<SqlDataReader>(_activeConnection);
30653071

30663072
CancellationTokenRegistration registration = new CancellationTokenRegistration();
30673073
if (cancellationToken.CanBeCanceled)
@@ -3077,7 +3083,7 @@ private Task<SqlDataReader> InternalExecuteReaderAsync(CommandBehavior behavior,
30773083
Task<SqlDataReader> returnedTask = source.Task;
30783084
try
30793085
{
3080-
RegisterForConnectionCloseNotification(ref returnedTask);
3086+
returnedTask = RegisterForConnectionCloseNotification(returnedTask);
30813087

30823088
Task<SqlDataReader>.Factory.FromAsync(BeginExecuteReaderAsync, EndExecuteReaderAsync, behavior, null).ContinueWith((t) =>
30833089
{
@@ -3207,7 +3213,10 @@ private Task<XmlReader> InternalExecuteXmlReaderAsync(CancellationToken cancella
32073213
SqlClientEventSource.Log.TryCorrelationTraceEvent("<sc.SqlCommand.ExecuteXmlReaderAsync|API|Correlation> ObjectID {0}, ActivityID {1}", ObjectID, ActivityCorrelator.Current);
32083214
SqlConnection.ExecutePermission.Demand();
32093215

3210-
TaskCompletionSource<XmlReader> source = new TaskCompletionSource<XmlReader>();
3216+
// connection can be used as state in RegisterForConnectionCloseNotification continuation
3217+
// to avoid an allocation so use it as the state value if possible but it can be changed if
3218+
// you need it for a more important piece of data that justifies the tuple allocation later
3219+
TaskCompletionSource<XmlReader> source = new TaskCompletionSource<XmlReader>(_activeConnection);
32113220

32123221
CancellationTokenRegistration registration = new CancellationTokenRegistration();
32133222
if (cancellationToken.CanBeCanceled)
@@ -3223,7 +3232,7 @@ private Task<XmlReader> InternalExecuteXmlReaderAsync(CancellationToken cancella
32233232
Task<XmlReader> returnedTask = source.Task;
32243233
try
32253234
{
3226-
RegisterForConnectionCloseNotification(ref returnedTask);
3235+
returnedTask = RegisterForConnectionCloseNotification(returnedTask);
32273236

32283237
Task<XmlReader>.Factory.FromAsync(BeginExecuteXmlReaderAsync, EndExecuteXmlReaderAsync, null).ContinueWith((t) =>
32293238
{
@@ -5814,7 +5823,7 @@ object ICloneable.Clone()
58145823
return Clone();
58155824
}
58165825

5817-
private void RegisterForConnectionCloseNotification<T>(ref Task<T> outterTask)
5826+
private Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outterTask)
58185827
{
58195828
SqlConnection connection = _activeConnection;
58205829
if (connection == null)
@@ -5823,7 +5832,7 @@ private void RegisterForConnectionCloseNotification<T>(ref Task<T> outterTask)
58235832
throw ADP.ClosedConnectionError();
58245833
}
58255834

5826-
connection.RegisterForConnectionCloseNotification<T>(ref outterTask, this, SqlReferenceCollection.CommandTag);
5835+
return connection.RegisterForConnectionCloseNotification(outterTask, this, SqlReferenceCollection.CommandTag);
58275836
}
58285837

58295838
// validates that a command has commandText and a non-busy open connection

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

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2805,16 +2805,52 @@ private static void ChangePassword(string connectionString, SqlConnectionString
28052805
SqlConnectionFactory.SingletonInstance.ClearPool(key);
28062806
}
28072807

2808-
internal void RegisterForConnectionCloseNotification<T>(ref Task<T> outerTask, object value, int tag)
2808+
internal Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outerTask, object value, int tag)
28092809
{
28102810
// Connection exists, schedule removal, will be added to ref collection after calling ValidateAndReconnect
2811-
outerTask = outerTask.ContinueWith(task =>
2811+
2812+
object state = null;
2813+
if (outerTask.AsyncState == this)
2814+
{
2815+
// if the caller created the TaskCompletionSource for outerTask with this connection
2816+
// as the state parameter (which is immutable) we can use task.AsyncState and state
2817+
// to carry the two pieces of state that we need into the continuation avoiding the
2818+
// allocation of a new state object to carry them
2819+
state = value;
2820+
}
2821+
else
28122822
{
2813-
RemoveWeakReference(value);
2814-
return task;
2815-
}, TaskScheduler.Default).Unwrap();
2823+
// otherwise we need to create a Tuple to carry the two pieces of state
2824+
state = Tuple.Create(this, value);
2825+
}
2826+
2827+
return outerTask.ContinueWith(
2828+
continuationFunction: static (task, state) =>
2829+
{
2830+
SqlConnection connection = null;
2831+
object obj = null;
2832+
if (state is Tuple<SqlConnection, object> tuple)
2833+
{
2834+
// special state tuple, unpack it
2835+
connection = tuple.Item1;
2836+
obj = tuple.Item2;
2837+
}
2838+
else
2839+
{
2840+
// use state on task and state object
2841+
connection = (SqlConnection)task.AsyncState;
2842+
obj = state;
2843+
}
2844+
2845+
connection.RemoveWeakReference(obj);
2846+
return task;
2847+
},
2848+
state: state,
2849+
scheduler: TaskScheduler.Default
2850+
).Unwrap();
28162851
}
28172852

2853+
28182854
// updates our context with any changes made to the memory-mapped data by an external process
28192855
static private void RefreshMemoryMappedData(SqlDebugContext sdc)
28202856
{

0 commit comments

Comments
 (0)