Skip to content

Commit 7b82e07

Browse files
Wraith2Wraith2
authored
Perf: Rework ExecuteReaderAsync to minimize allocations (#528)
1 parent 415262d commit 7b82e07

File tree

8 files changed

+362
-108
lines changed

8 files changed

+362
-108
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -706,6 +706,7 @@
706706
<Reference Include="System.Memory" />
707707
</ItemGroup>
708708
<ItemGroup>
709+
<Compile Include="Microsoft\Data\SqlClient\AAsyncCallContext.cs" />
709710
<Compile Include="Resources\SR.Designer.cs">
710711
<DesignTime>True</DesignTime>
711712
<AutoGen>True</AutoGen>

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/ProviderBase/DbConnectionFactory.cs‎

Lines changed: 65 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// The .NET Foundation licenses this file to you under the MIT license.
33
// See the LICENSE file in the project root for more information.
44

5+
using System;
56
using System.Data;
67
using System.Data.Common;
78
using System.Diagnostics;
@@ -15,6 +16,7 @@ namespace Microsoft.Data.ProviderBase
1516
{
1617
internal abstract partial class DbConnectionFactory
1718
{
19+
private static readonly Action<Task<DbConnectionInternal>, object> s_tryGetConnectionCompletedContinuation = TryGetConnectionCompletedContinuation;
1820

1921
internal bool TryGetConnection(DbConnection owningConnection, TaskCompletionSource<DbConnectionInternal> retry, DbConnectionOptions userOptions, DbConnectionInternal oldConnection, out DbConnectionInternal connection)
2022
{
@@ -82,25 +84,7 @@ internal bool TryGetConnection(DbConnection owningConnection, TaskCompletionSour
8284

8385
// now that we have an antecedent task, schedule our work when it is completed.
8486
// If it is a new slot or a completed task, this continuation will start right away.
85-
newTask = s_pendingOpenNonPooled[idx].ContinueWith((_) =>
86-
{
87-
Transaction originalTransaction = ADP.GetCurrentTransaction();
88-
try
89-
{
90-
ADP.SetCurrentTransaction(retry.Task.AsyncState as Transaction);
91-
var newConnection = CreateNonPooledConnection(owningConnection, poolGroup, userOptions);
92-
if ((oldConnection != null) && (oldConnection.State == ConnectionState.Open))
93-
{
94-
oldConnection.PrepareForReplaceConnection();
95-
oldConnection.Dispose();
96-
}
97-
return newConnection;
98-
}
99-
finally
100-
{
101-
ADP.SetCurrentTransaction(originalTransaction);
102-
}
103-
}, cancellationTokenSource.Token, TaskContinuationOptions.LongRunning, TaskScheduler.Default);
87+
newTask = CreateReplaceConnectionContinuation(s_pendingOpenNonPooled[idx], owningConnection, retry, userOptions, oldConnection, poolGroup, cancellationTokenSource);
10488

10589
// Place this new task in the slot so any future work will be queued behind it
10690
s_pendingOpenNonPooled[idx] = newTask;
@@ -114,29 +98,11 @@ internal bool TryGetConnection(DbConnection owningConnection, TaskCompletionSour
11498
}
11599

116100
// once the task is done, propagate the final results to the original caller
117-
newTask.ContinueWith((task) =>
118-
{
119-
cancellationTokenSource.Dispose();
120-
if (task.IsCanceled)
121-
{
122-
retry.TrySetException(ADP.ExceptionWithStackTrace(ADP.NonPooledOpenTimeout()));
123-
}
124-
else if (task.IsFaulted)
125-
{
126-
retry.TrySetException(task.Exception.InnerException);
127-
}
128-
else
129-
{
130-
if (!retry.TrySetResult(task.Result))
131-
{
132-
// The outer TaskCompletionSource was already completed
133-
// Which means that we don't know if someone has messed with the outer connection in the middle of creation
134-
// So the best thing to do now is to destroy the newly created connection
135-
task.Result.DoomThisConnection();
136-
task.Result.Dispose();
137-
}
138-
}
139-
}, TaskScheduler.Default);
101+
newTask.ContinueWith(
102+
continuationAction: s_tryGetConnectionCompletedContinuation,
103+
state: Tuple.Create(cancellationTokenSource, retry),
104+
scheduler: TaskScheduler.Default
105+
);
140106

141107
return false;
142108
}
@@ -188,5 +154,62 @@ internal bool TryGetConnection(DbConnection owningConnection, TaskCompletionSour
188154

189155
return true;
190156
}
157+
158+
private Task<DbConnectionInternal> CreateReplaceConnectionContinuation(Task<DbConnectionInternal> task, DbConnection owningConnection, TaskCompletionSource<DbConnectionInternal> retry, DbConnectionOptions userOptions, DbConnectionInternal oldConnection, DbConnectionPoolGroup poolGroup, CancellationTokenSource cancellationTokenSource)
159+
{
160+
return task.ContinueWith(
161+
(_) =>
162+
{
163+
Transaction originalTransaction = ADP.GetCurrentTransaction();
164+
try
165+
{
166+
ADP.SetCurrentTransaction(retry.Task.AsyncState as Transaction);
167+
var newConnection = CreateNonPooledConnection(owningConnection, poolGroup, userOptions);
168+
if ((oldConnection != null) && (oldConnection.State == ConnectionState.Open))
169+
{
170+
oldConnection.PrepareForReplaceConnection();
171+
oldConnection.Dispose();
172+
}
173+
return newConnection;
174+
}
175+
finally
176+
{
177+
ADP.SetCurrentTransaction(originalTransaction);
178+
}
179+
},
180+
cancellationTokenSource.Token,
181+
TaskContinuationOptions.LongRunning,
182+
TaskScheduler.Default
183+
);
184+
}
185+
186+
private static void TryGetConnectionCompletedContinuation(Task<DbConnectionInternal> task, object state)
187+
{
188+
Tuple<CancellationTokenSource, TaskCompletionSource<DbConnectionInternal>> parameters = (Tuple<CancellationTokenSource, TaskCompletionSource<DbConnectionInternal>>)state;
189+
CancellationTokenSource source = parameters.Item1;
190+
source.Dispose();
191+
192+
TaskCompletionSource<DbConnectionInternal> retryCompletionSource = parameters.Item2;
193+
194+
if (task.IsCanceled)
195+
{
196+
retryCompletionSource.TrySetException(ADP.ExceptionWithStackTrace(ADP.NonPooledOpenTimeout()));
197+
}
198+
else if (task.IsFaulted)
199+
{
200+
retryCompletionSource.TrySetException(task.Exception.InnerException);
201+
}
202+
else
203+
{
204+
if (!retryCompletionSource.TrySetResult(task.Result))
205+
{
206+
// The outer TaskCompletionSource was already completed
207+
// Which means that we don't know if someone has messed with the outer connection in the middle of creation
208+
// So the best thing to do now is to destroy the newly created connection
209+
task.Result.DoomThisConnection();
210+
task.Result.Dispose();
211+
}
212+
}
213+
}
191214
}
192215
}

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System;
6+
using System.Threading.Tasks;
7+
8+
namespace Microsoft.Data.SqlClient
9+
{
10+
// this class is a base class for creating derived objects that will store state for async operations
11+
// avoiding the use of closures and allowing caching/reuse of the instances for frequently used async
12+
// calls
13+
//
14+
// DO derive from this and seal your class
15+
// DO add additional fields or properties needed for the async operation and then override Clear to zero them
16+
// DO override AfterClear and use the owner parameter to return the object to a cache location if you have one, this is the purpose of the method
17+
// CONSIDER creating your own Set method that calls the base Set rather than providing a parameterized ctor, it is friendlier to caching
18+
// DO NOT use this class' state after Dispose has been called. It will not throw ObjectDisposedException but it will be a cleared object
19+
20+
internal abstract class AAsyncCallContext<TOwner, TTask> : IDisposable
21+
where TOwner : class
22+
{
23+
protected TOwner _owner;
24+
protected TaskCompletionSource<TTask> _source;
25+
protected IDisposable _disposable;
26+
27+
protected AAsyncCallContext()
28+
{
29+
}
30+
31+
protected AAsyncCallContext(TOwner owner, TaskCompletionSource<TTask> source, IDisposable disposable = null)
32+
{
33+
Set(owner, source, disposable);
34+
}
35+
36+
protected void Set(TOwner owner, TaskCompletionSource<TTask> source, IDisposable disposable = null)
37+
{
38+
_owner = owner ?? throw new ArgumentNullException(nameof(owner));
39+
_source = source ?? throw new ArgumentNullException(nameof(source));
40+
_disposable = disposable;
41+
}
42+
43+
protected void ClearCore()
44+
{
45+
_source = null;
46+
_owner = default;
47+
IDisposable copyDisposable = _disposable;
48+
_disposable = null;
49+
copyDisposable?.Dispose();
50+
}
51+
52+
/// <summary>
53+
/// override this method to cleanup instance data before ClearCore is called which will blank the base data
54+
/// </summary>
55+
protected virtual void Clear()
56+
{
57+
}
58+
59+
/// <summary>
60+
/// override this method to do work after the instance has been totally blanked, intended for cache return etc
61+
/// </summary>
62+
protected virtual void AfterCleared(TOwner owner)
63+
{
64+
}
65+
66+
public void Dispose()
67+
{
68+
TOwner owner = _owner;
69+
try
70+
{
71+
Clear();
72+
}
73+
finally
74+
{
75+
ClearCore();
76+
}
77+
AfterCleared(owner);
78+
}
79+
}
80+
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2265,7 +2265,7 @@ private Task ReadWriteColumnValueAsync(int col)
22652265
return writeTask;
22662266
}
22672267

2268-
private void RegisterForConnectionCloseNotification<T>(ref Task<T> outerTask)
2268+
private Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outerTask)
22692269
{
22702270
SqlConnection connection = _connection;
22712271
if (connection == null)
@@ -2274,7 +2274,7 @@ private void RegisterForConnectionCloseNotification<T>(ref Task<T> outerTask)
22742274
throw ADP.ClosedConnectionError();
22752275
}
22762276

2277-
connection.RegisterForConnectionCloseNotification<T>(ref outerTask, this, SqlReferenceCollection.BulkCopyTag);
2277+
return connection.RegisterForConnectionCloseNotification(outerTask, this, SqlReferenceCollection.BulkCopyTag);
22782278
}
22792279

22802280
// Runs a loop to copy all columns of a single row.
@@ -3057,7 +3057,7 @@ private Task WriteToServerInternalAsync(CancellationToken ctoken)
30573057
source = new TaskCompletionSource<object>(); // Creating the completion source/Task that we pass to application
30583058
resultTask = source.Task;
30593059

3060-
RegisterForConnectionCloseNotification(ref resultTask);
3060+
resultTask = RegisterForConnectionCloseNotification(resultTask);
30613061
}
30623062

30633063
if (_destinationTableName == null)

0 commit comments

Comments
 (0)