Skip to content

Commit efc0c28

Browse files
benrr101benrr101
committed
Move TryGetConnection
1 parent 5aa2d13 commit efc0c28

File tree

2 files changed

+147
-143
lines changed

2 files changed

+147
-143
lines changed

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

Lines changed: 4 additions & 143 deletions
Original file line numberDiff line numberDiff line change
@@ -41,152 +41,13 @@ protected DbConnectionOptions FindConnectionOptions(DbConnectionPoolKey key)
4141
return null;
4242
}
4343

44-
private static Task<DbConnectionInternal> GetCompletedTask()
44+
protected static Task<DbConnectionInternal> GetCompletedTask()
4545
{
4646
Debug.Assert(Monitor.IsEntered(s_pendingOpenNonPooled), $"Expected {nameof(s_pendingOpenNonPooled)} lock to be held.");
4747
return s_completedTask ?? (s_completedTask = Task.FromResult<DbConnectionInternal>(null));
4848
}
4949

50-
internal bool TryGetConnection(DbConnection owningConnection, TaskCompletionSource<DbConnectionInternal> retry, DbConnectionOptions userOptions, DbConnectionInternal oldConnection, out DbConnectionInternal connection)
51-
{
52-
Debug.Assert(owningConnection != null, "null owningConnection?");
53-
54-
DbConnectionPoolGroup poolGroup;
55-
IDbConnectionPool connectionPool;
56-
connection = null;
57-
58-
// Work around race condition with clearing the pool between GetConnectionPool obtaining pool
59-
// and GetConnection on the pool checking the pool state. Clearing the pool in this window
60-
// will switch the pool into the ShuttingDown state, and GetConnection will return null.
61-
// There is probably a better solution involving locking the pool/group, but that entails a major
62-
// re-design of the connection pooling synchronization, so is postponed for now.
63-
64-
// Use retriesLeft to prevent CPU spikes with incremental sleep
65-
// start with one msec, double the time every retry
66-
// max time is: 1 + 2 + 4 + ... + 2^(retries-1) == 2^retries -1 == 1023ms (for 10 retries)
67-
int retriesLeft = 10;
68-
int timeBetweenRetriesMilliseconds = 1;
69-
70-
do
71-
{
72-
poolGroup = GetConnectionPoolGroup(owningConnection);
73-
// Doing this on the callers thread is important because it looks up the WindowsIdentity from the thread.
74-
connectionPool = GetConnectionPool(owningConnection, poolGroup);
75-
if (connectionPool == null)
76-
{
77-
// If GetConnectionPool returns null, we can be certain that
78-
// this connection should not be pooled via DbConnectionPool
79-
// or have a disabled pool entry.
80-
poolGroup = GetConnectionPoolGroup(owningConnection); // previous entry have been disabled
81-
82-
if (retry != null)
83-
{
84-
Task<DbConnectionInternal> newTask;
85-
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
86-
lock (s_pendingOpenNonPooled)
87-
{
88-
// look for an available task slot (completed or empty)
89-
int idx;
90-
for (idx = 0; idx < s_pendingOpenNonPooled.Length; idx++)
91-
{
92-
Task task = s_pendingOpenNonPooled[idx];
93-
if (task == null)
94-
{
95-
s_pendingOpenNonPooled[idx] = GetCompletedTask();
96-
break;
97-
}
98-
else if (task.IsCompleted)
99-
{
100-
break;
101-
}
102-
}
103-
104-
// if didn't find one, pick the next one in round-robin fashion
105-
if (idx == s_pendingOpenNonPooled.Length)
106-
{
107-
idx = (int)(s_pendingOpenNonPooledNext % s_pendingOpenNonPooled.Length);
108-
unchecked
109-
{
110-
s_pendingOpenNonPooledNext++;
111-
}
112-
}
113-
114-
// now that we have an antecedent task, schedule our work when it is completed.
115-
// If it is a new slot or a completed task, this continuation will start right away.
116-
newTask = CreateReplaceConnectionContinuation(s_pendingOpenNonPooled[idx], owningConnection, retry, userOptions, oldConnection, poolGroup, cancellationTokenSource);
117-
118-
// Place this new task in the slot so any future work will be queued behind it
119-
s_pendingOpenNonPooled[idx] = newTask;
120-
}
121-
122-
// Set up the timeout (if needed)
123-
if (owningConnection.ConnectionTimeout > 0)
124-
{
125-
int connectionTimeoutMilliseconds = owningConnection.ConnectionTimeout * 1000;
126-
cancellationTokenSource.CancelAfter(connectionTimeoutMilliseconds);
127-
}
128-
129-
// once the task is done, propagate the final results to the original caller
130-
newTask.ContinueWith(
131-
continuationAction: TryGetConnectionCompletedContinuation,
132-
state: Tuple.Create(cancellationTokenSource, retry),
133-
scheduler: TaskScheduler.Default
134-
);
135-
136-
return false;
137-
}
138-
139-
connection = CreateNonPooledConnection(owningConnection, poolGroup, userOptions);
140-
141-
SqlClientEventSource.Metrics.EnterNonPooledConnection();
142-
}
143-
else
144-
{
145-
if (((SqlClient.SqlConnection)owningConnection).ForceNewConnection)
146-
{
147-
Debug.Assert(!(oldConnection is DbConnectionClosed), "Force new connection, but there is no old connection");
148-
connection = connectionPool.ReplaceConnection(owningConnection, userOptions, oldConnection);
149-
}
150-
else
151-
{
152-
if (!connectionPool.TryGetConnection(owningConnection, retry, userOptions, out connection))
153-
{
154-
return false;
155-
}
156-
}
157-
158-
if (connection == null)
159-
{
160-
// connection creation failed on semaphore waiting or if max pool reached
161-
if (connectionPool.IsRunning)
162-
{
163-
SqlClientEventSource.Log.TryTraceEvent("<prov.DbConnectionFactory.GetConnection|RES|CPOOL> {0}, GetConnection failed because a pool timeout occurred.", ObjectID);
164-
// If GetConnection failed while the pool is running, the pool timeout occurred.
165-
throw ADP.PooledOpenTimeout();
166-
}
167-
else
168-
{
169-
// We've hit the race condition, where the pool was shut down after we got it from the group.
170-
// Yield time slice to allow shut down activities to complete and a new, running pool to be instantiated
171-
// before retrying.
172-
System.Threading.Thread.Sleep(timeBetweenRetriesMilliseconds);
173-
timeBetweenRetriesMilliseconds *= 2; // double the wait time for next iteration
174-
}
175-
}
176-
}
177-
} while (connection == null && retriesLeft-- > 0);
178-
179-
if (connection == null)
180-
{
181-
SqlClientEventSource.Log.TryTraceEvent("<prov.DbConnectionFactory.GetConnection|RES|CPOOL> {0}, GetConnection failed because a pool timeout occurred and all retries were exhausted.", ObjectID);
182-
// exhausted all retries or timed out - give up
183-
throw ADP.PooledOpenTimeout();
184-
}
185-
186-
return true;
187-
}
188-
189-
private Task<DbConnectionInternal> CreateReplaceConnectionContinuation(Task<DbConnectionInternal> task, DbConnection owningConnection, TaskCompletionSource<DbConnectionInternal> retry, DbConnectionOptions userOptions, DbConnectionInternal oldConnection, DbConnectionPoolGroup poolGroup, CancellationTokenSource cancellationTokenSource)
50+
protected Task<DbConnectionInternal> CreateReplaceConnectionContinuation(Task<DbConnectionInternal> task, DbConnection owningConnection, TaskCompletionSource<DbConnectionInternal> retry, DbConnectionOptions userOptions, DbConnectionInternal oldConnection, DbConnectionPoolGroup poolGroup, CancellationTokenSource cancellationTokenSource)
19051
{
19152
return task.ContinueWith(
19253
(_) =>
@@ -214,7 +75,7 @@ private Task<DbConnectionInternal> CreateReplaceConnectionContinuation(Task<DbCo
21475
);
21576
}
21677

217-
private void TryGetConnectionCompletedContinuation(Task<DbConnectionInternal> task, object state)
78+
protected void TryGetConnectionCompletedContinuation(Task<DbConnectionInternal> task, object state)
21879
{
21980
Tuple<CancellationTokenSource, TaskCompletionSource<DbConnectionInternal>> parameters = (Tuple<CancellationTokenSource, TaskCompletionSource<DbConnectionInternal>>)state;
22081
CancellationTokenSource source = parameters.Item1;
@@ -247,7 +108,7 @@ private void TryGetConnectionCompletedContinuation(Task<DbConnectionInternal> ta
247108
}
248109
}
249110

250-
private IDbConnectionPool GetConnectionPool(DbConnection owningObject, DbConnectionPoolGroup connectionPoolGroup)
111+
protected IDbConnectionPool GetConnectionPool(DbConnection owningObject, DbConnectionPoolGroup connectionPoolGroup)
251112
{
252113
// if poolgroup is disabled, it will be replaced with a new entry
253114

‎src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnectionFactory.cs

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,149 @@ internal void QueuePoolGroupForRelease(DbConnectionPoolGroup poolGroup)
197197
SqlClientEventSource.Metrics.EnterInactiveConnectionPoolGroup();
198198
SqlClientEventSource.Metrics.ExitActiveConnectionPoolGroup();
199199
}
200+
201+
internal bool TryGetConnection(
202+
DbConnection owningConnection,
203+
TaskCompletionSource<DbConnectionInternal> retry,
204+
DbConnectionOptions userOptions,
205+
DbConnectionInternal oldConnection,
206+
out DbConnectionInternal connection)
207+
{
208+
Debug.Assert(owningConnection is not null, "null owningConnection?");
209+
210+
connection = null;
211+
212+
// Work around race condition with clearing the pool between GetConnectionPool obtaining pool
213+
// and GetConnection on the pool checking the pool state. Clearing the pool in this window
214+
// will switch the pool into the ShuttingDown state, and GetConnection will return null.
215+
// There is probably a better solution involving locking the pool/group, but that entails a major
216+
// re-design of the connection pooling synchronization, so is postponed for now.
217+
218+
// Use retriesLeft to prevent CPU spikes with incremental sleep
219+
// start with one msec, double the time every retry
220+
// max time is: 1 + 2 + 4 + ... + 2^(retries-1) == 2^retries -1 == 1023ms (for 10 retries)
221+
int retriesLeft = 10;
222+
int timeBetweenRetriesMilliseconds = 1;
223+
224+
do
225+
{
226+
DbConnectionPoolGroup poolGroup = GetConnectionPoolGroup(owningConnection);
227+
228+
// Doing this on the callers thread is important because it looks up the WindowsIdentity from the thread.
229+
IDbConnectionPool connectionPool = GetConnectionPool(owningConnection, poolGroup);
230+
if (connectionPool == null)
231+
{
232+
// If GetConnectionPool returns null, we can be certain that this connection
233+
// should not be pooled via DbConnectionPool or have a disabled pool entry.
234+
poolGroup = GetConnectionPoolGroup(owningConnection); // previous entry have been disabled
235+
236+
if (retry is not null)
237+
{
238+
Task<DbConnectionInternal> newTask;
239+
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
240+
lock (s_pendingOpenNonPooled)
241+
{
242+
// look for an available task slot (completed or empty)
243+
int idx;
244+
for (idx = 0; idx < s_pendingOpenNonPooled.Length; idx++)
245+
{
246+
Task task = s_pendingOpenNonPooled[idx];
247+
if (task is null)
248+
{
249+
s_pendingOpenNonPooled[idx] = GetCompletedTask();
250+
break;
251+
}
252+
253+
if (task.IsCompleted)
254+
{
255+
break;
256+
}
257+
}
258+
259+
// if didn't find one, pick the next one in round-robin fashion
260+
if (idx == s_pendingOpenNonPooled.Length)
261+
{
262+
idx = (int)(s_pendingOpenNonPooledNext % s_pendingOpenNonPooled.Length);
263+
unchecked
264+
{
265+
s_pendingOpenNonPooledNext++;
266+
}
267+
}
268+
269+
// now that we have an antecedent task, schedule our work when it is completed.
270+
// If it is a new slot or a completed task, this continuation will start right away.
271+
newTask = CreateReplaceConnectionContinuation(s_pendingOpenNonPooled[idx], owningConnection, retry, userOptions, oldConnection, poolGroup, cancellationTokenSource);
272+
273+
// Place this new task in the slot so any future work will be queued behind it
274+
s_pendingOpenNonPooled[idx] = newTask;
275+
}
276+
277+
// Set up the timeout (if needed)
278+
if (owningConnection.ConnectionTimeout > 0)
279+
{
280+
int connectionTimeoutMilliseconds = owningConnection.ConnectionTimeout * 1000;
281+
cancellationTokenSource.CancelAfter(connectionTimeoutMilliseconds);
282+
}
283+
284+
// once the task is done, propagate the final results to the original caller
285+
newTask.ContinueWith(
286+
continuationAction: TryGetConnectionCompletedContinuation,
287+
state: Tuple.Create(cancellationTokenSource, retry),
288+
scheduler: TaskScheduler.Default
289+
);
290+
291+
return false;
292+
}
293+
294+
connection = CreateNonPooledConnection(owningConnection, poolGroup, userOptions);
295+
296+
SqlClientEventSource.Metrics.EnterNonPooledConnection();
297+
}
298+
else
299+
{
300+
if (((SqlConnection)owningConnection).ForceNewConnection)
301+
{
302+
Debug.Assert(oldConnection is not DbConnectionClosed, "Force new connection, but there is no old connection");
303+
304+
connection = connectionPool.ReplaceConnection(owningConnection, userOptions, oldConnection);
305+
}
306+
else
307+
{
308+
if (!connectionPool.TryGetConnection(owningConnection, retry, userOptions, out connection))
309+
{
310+
return false;
311+
}
312+
}
313+
314+
if (connection is null)
315+
{
316+
// connection creation failed on semaphore waiting or if max pool reached
317+
if (connectionPool.IsRunning)
318+
{
319+
SqlClientEventSource.Log.TryTraceEvent("<prov.SqlConnectionFactory.GetConnection|RES|CPOOL> {0}, GetConnection failed because a pool timeout occurred.", ObjectId);
320+
// If GetConnection failed while the pool is running, the pool timeout occurred.
321+
throw ADP.PooledOpenTimeout();
322+
}
323+
324+
// We've hit the race condition, where the pool was shut down after we
325+
// got it from the group. Yield time slice to allow shutdown activities
326+
// to complete and a new, running pool to be instantiated before
327+
// retrying.
328+
Thread.Sleep(timeBetweenRetriesMilliseconds);
329+
timeBetweenRetriesMilliseconds *= 2; // double the wait time for next iteration
330+
}
331+
}
332+
} while (connection == null && retriesLeft-- > 0);
333+
334+
if (connection == null)
335+
{
336+
SqlClientEventSource.Log.TryTraceEvent("<prov.SqlConnectionFactory.GetConnection|RES|CPOOL> {0}, GetConnection failed because a pool timeout occurred and all retries were exhausted.", ObjectId);
337+
// exhausted all retries or timed out - give up
338+
throw ADP.PooledOpenTimeout();
339+
}
340+
341+
return true;
342+
}
200343

201344
#endregion
202345

0 commit comments

Comments
 (0)