Skip to content

Commit 5c66f1b

Browse files
Wraith2Wraith2
committed
rework ExecuteReaderAsync to minimize allocations
1 parent 46ab2ec commit 5c66f1b

File tree

8 files changed

+366
-110
lines changed

8 files changed

+366
-110
lines changed

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
<Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientEventSource.cs">
3434
<Link>Microsoft\Data\SqlClient\SqlClientEventSource.cs</Link>
3535
</Compile>
36-
<Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientLogger.cs" >
36+
<Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientLogger.cs">
3737
<Link>Microsoft\Data\SqlClient\SqlClientLogger.cs</Link>
3838
</Compile>
3939
<Compile Include="..\..\src\Microsoft\Data\SqlClient\DataClassification\SensitivityClassification.cs">
@@ -57,7 +57,7 @@
5757
<Compile Include="..\..\src\Microsoft\Data\SqlClient\ColumnEncryptionKeyInfo.cs">
5858
<Link>Microsoft\Data\SqlClient\ColumnEncryptionKeyInfo.cs</Link>
5959
</Compile>
60-
<Compile Include="..\..\src\Microsoft\Data\SqlClient\OnChangedEventHandler.cs" >
60+
<Compile Include="..\..\src\Microsoft\Data\SqlClient\OnChangedEventHandler.cs">
6161
<Link>Microsoft\Data\SqlClient\OnChangedEventHandler.cs</Link>
6262
</Compile>
6363
<Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlAeadAes256CbcHmac256Algorithm.cs">
@@ -635,6 +635,7 @@
635635
<Reference Include="System.Memory" />
636636
</ItemGroup>
637637
<ItemGroup>
638+
<Compile Include="Microsoft\Data\SqlClient\AAsyncCallContext.cs" />
638639
<Compile Include="Resources\SR.Designer.cs">
639640
<DesignTime>True</DesignTime>
640641
<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> retryCompletionSrouce = parameters.Item2;
193+
194+
if (task.IsCanceled)
195+
{
196+
retryCompletionSrouce.TrySetException(ADP.ExceptionWithStackTrace(ADP.NonPooledOpenTimeout()));
197+
}
198+
else if (task.IsFaulted)
199+
{
200+
retryCompletionSrouce.TrySetException(task.Exception.InnerException);
201+
}
202+
else
203+
{
204+
if (!retryCompletionSrouce.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: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
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 and 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 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+
/// overrride 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+
67+
public void Dispose()
68+
{
69+
TOwner owner = _owner;
70+
try
71+
{
72+
Clear();
73+
}
74+
finally
75+
{
76+
ClearCore();
77+
}
78+
AfterCleared(owner);
79+
}
80+
}
81+
}

‎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
@@ -2224,7 +2224,7 @@ private Task ReadWriteColumnValueAsync(int col)
22242224
return writeTask;
22252225
}
22262226

2227-
private void RegisterForConnectionCloseNotification<T>(ref Task<T> outerTask)
2227+
private Task<T> RegisterForConnectionCloseNotification<T>(Task<T> outerTask)
22282228
{
22292229
SqlConnection connection = _connection;
22302230
if (connection == null)
@@ -2233,7 +2233,7 @@ private void RegisterForConnectionCloseNotification<T>(ref Task<T> outerTask)
22332233
throw ADP.ClosedConnectionError();
22342234
}
22352235

2236-
connection.RegisterForConnectionCloseNotification<T>(ref outerTask, this, SqlReferenceCollection.BulkCopyTag);
2236+
return connection.RegisterForConnectionCloseNotification(outerTask, this, SqlReferenceCollection.BulkCopyTag);
22372237
}
22382238

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

3019-
RegisterForConnectionCloseNotification(ref resultTask);
3019+
resultTask = RegisterForConnectionCloseNotification(resultTask);
30203020
}
30213021

30223022
if (_destinationTableName == null)

0 commit comments

Comments
 (0)