Skip to content

Commit 6431209

Browse files
Implement pool shutdown for ChannelDbConnectionPool and harden WaitHandleDbConnectionPool shutdown (AB#44828)
ChannelDbConnectionPool: - Shutdown() now transitions State to ShuttingDown (FR-001), completes the idle channel writer to unblock async waiters (FR-002, FR-007), and drains and destroys buffered idle connections (FR-003). Implementation is idempotent via Interlocked.CompareExchange (FR-006). - Startup() emits a trace and is a structural counterpart to Shutdown. - ReturnInternalConnection no longer asserts on TryWrite; if the channel was completed by a concurrent shutdown, the returning connection is destroyed (FR-004 race-window guard). - IdleConnectionChannel.Complete() exposes channel completion to the pool. WaitHandleDbConnectionPool: - Shutdown() is now idempotent (FR-006), disposes both _cleanupTimer and _errorTimer (FR-005), wakes threads parked in WaitHandle.WaitAny by releasing the pool semaphore (FR-007), and drains _stackNew/_stackOld (FR-003). - CleanupCallback and ErrorCallback short-circuit when State == ShuttingDown so an in-flight callback racing with Shutdown does not re-arm work. - The private TryGetConnection wait loop re-checks State after WaitHandle.WaitAny and bails out, so waiters cannot spin against a drained pool. Tests: - ChannelDbConnectionPoolShutdownTest (7 tests) covers state transition, drain, return-after-shutdown, idempotency, async waiter unblock, post-shutdown return, and Startup-after-Shutdown. - WaitHandleDbConnectionPoolShutdownTest (9 tests) covers state transition, cleanup/error timer disposal, drain, idempotency, callback no-op after shutdown, sync TryGetConnection short-circuit, and sync waiter unblock. Verified by running targeted suite: 340 tests passing across net8.0, net9.0, net10.0, and net462. Spec: specs/004-pool-shutdown.
1 parent 0759bee commit 6431209

5 files changed

Lines changed: 540 additions & 10 deletions

File tree

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,13 @@ internal sealed class ChannelDbConnectionPool : IDbConnectionPool
9292
/// Must be updated using <see cref="Interlocked"/> operations to ensure thread safety.
9393
/// </summary>
9494
private volatile int _isClearing;
95+
96+
/// <summary>
97+
/// Tracks whether <see cref="Shutdown"/> has already initiated the shutdown sequence so that
98+
/// repeated calls are observed as no-ops (FR-006). Updated atomically via
99+
/// <see cref="Interlocked.CompareExchange(ref int, int, int)"/>.
100+
/// </summary>
101+
private int _shutdownInitiated;
95102
#endregion
96103

97104
/// <summary>
@@ -254,21 +261,57 @@ public void ReturnInternalConnection(DbConnectionInternal connection, DbConnecti
254261
}
255262
else
256263
{
257-
var written = _idleChannel.TryWrite(connection);
258-
Debug.Assert(written, "Failed to write returning connection to the idle channel.");
264+
if (!_idleChannel.TryWrite(connection))
265+
{
266+
// The channel has been completed (pool is shutting down). Race window
267+
// between the State check above and TryWrite: destroy instead of pooling.
268+
RemoveConnection(connection);
269+
}
259270
}
260271
}
261272

262273
/// <inheritdoc />
263274
public void Shutdown()
264275
{
265-
// No-op for now, warmup will be implemented later.
276+
// FR-006: idempotent. Compare-and-exchange ensures only one caller performs shutdown work.
277+
if (Interlocked.CompareExchange(ref _shutdownInitiated, 1, 0) != 0)
278+
{
279+
return;
280+
}
281+
282+
SqlClientEventSource.Log.TryPoolerTraceEvent(
283+
"<prov.DbConnectionPool.Shutdown|RES|INFO|CPOOL> {0}", Id);
284+
285+
// FR-001: transition to ShuttingDown. After this point, ReturnInternalConnection
286+
// routes returning connections to RemoveConnection.
287+
State = ShuttingDown;
288+
289+
// FR-002 + FR-007: complete the channel writer so:
290+
// - no further idle connections can be enqueued (TryWrite returns false), and
291+
// - in-flight / future async waiters on ReadAsync fault with ChannelClosedException.
292+
_idleChannel.Complete();
293+
294+
// FR-003: drain remaining buffered idle connections and destroy them. The channel is
295+
// unbounded so all already-enqueued items can be drained synchronously.
296+
while (_idleChannel.TryRead(out DbConnectionInternal? connection))
297+
{
298+
if (connection is not null)
299+
{
300+
RemoveConnection(connection);
301+
}
302+
// null sentinels are wake-up signals only; nothing to destroy.
303+
}
266304
}
267305

268306
/// <inheritdoc />
269307
public void Startup()
270308
{
271-
// No-op for now, warmup will be implemented later.
309+
// This pool has no background timers today (idle timeout is enforced lazily in
310+
// IsLiveConnection on retrieval; pruning is not implemented). State is set to Running
311+
// in the constructor, so this is currently the symmetrical counterpart of Shutdown.
312+
// Background work (warmup, pruning timers) will be added here when introduced.
313+
SqlClientEventSource.Log.TryPoolerTraceEvent(
314+
"<prov.DbConnectionPool.Startup|RES|INFO|CPOOL> {0}", Id);
272315
}
273316

274317
/// <inheritdoc />

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IdleConnectionChannel.cs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,19 @@ internal IdleConnectionChannel()
2626
{
2727
var channel = Channel.CreateUnbounded<DbConnectionInternal?>();
2828
_reader = channel.Reader;
29-
//TODO: the channel should be completed on pool shutdown
3029
_writer = channel.Writer;
3130
}
3231

32+
/// <summary>
33+
/// Marks the channel writer as complete. After completion, <see cref="TryWrite"/>
34+
/// returns <see langword="false"/> for any future writes, and any in-flight or future
35+
/// <see cref="ReadAsync"/> waiters will fault with <see cref="System.Threading.Channels.ChannelClosedException"/>
36+
/// once the channel is drained. Used by the connection pool to signal shutdown.
37+
/// </summary>
38+
/// <returns><see langword="true"/> if this call completed the channel; otherwise <see langword="false"/>
39+
/// (channel was already completed).</returns>
40+
internal bool Complete() => _writer.TryComplete();
41+
3342
/// <summary>
3443
/// The number of non-null connections currently in the channel.
3544
/// </summary>

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs

Lines changed: 61 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,14 @@ public bool IsRunning
329329

330330
private void CleanupCallback(object state)
331331
{
332+
// FR-005: if the pool is shutting down, skip work. Shutdown disposes the timer, but
333+
// a callback may already be in-flight when Shutdown runs; this guard ensures it does
334+
// not perform pruning or re-arm pool create requests.
335+
if (State == ShuttingDown)
336+
{
337+
return;
338+
}
339+
332340
// Called when the cleanup-timer ticks over.
333341

334342
// This is the automatic pruning method. Every period, we will
@@ -767,6 +775,13 @@ private void DestroyObject(DbConnectionInternal obj)
767775

768776
private void ErrorCallback(object state)
769777
{
778+
// FR-005: skip work if the pool is shutting down. The shutdown path disposes the
779+
// timer; this guard handles the in-flight-callback race.
780+
if (State == ShuttingDown)
781+
{
782+
return;
783+
}
784+
770785
SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.ErrorCallback|RES|CPOOL> {0}, Resetting Error handling.", Id);
771786
_errorOccurred = false;
772787
_waitHandles.ErrorEvent.Reset();
@@ -956,6 +971,16 @@ private bool TryGetConnection(DbConnection owningObject, uint waitForMultipleObj
956971
{
957972
waitResult = WaitHandle.WaitAny(_waitHandles.GetHandles(allowCreate), unchecked((int)waitForMultipleObjectsTimeout));
958973

974+
// FR-007: after waking, observe shutdown state and bail out so waiters
975+
// do not spin against a drained pool.
976+
if (State != Running)
977+
{
978+
SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.GetConnection|RES|CPOOL> {0}, Pool is shutting down; abandoning wait.", Id);
979+
Interlocked.Decrement(ref _waitCount);
980+
connection = null;
981+
return false;
982+
}
983+
959984
// From the WaitAny docs: "If more than one object became signaled during
960985
// the call, this is the array index of the signaled object with the
961986
// smallest index value of all the signaled objects." This is important
@@ -1481,14 +1506,45 @@ public void Startup()
14811506
public void Shutdown()
14821507
{
14831508
SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.Shutdown|RES|INFO|CPOOL> {0}", Id);
1509+
1510+
// FR-006: idempotent. Subsequent calls observe ShuttingDown and bail.
1511+
if (State == ShuttingDown)
1512+
{
1513+
return;
1514+
}
14841515
State = ShuttingDown;
14851516

1486-
// deactivate timer callbacks
1487-
Timer t = _cleanupTimer;
1488-
_cleanupTimer = null;
1489-
if (t != null)
1517+
// FR-005: dispose all background timers so they no longer schedule new work.
1518+
// Note that any timer callback already in flight may still observe State == ShuttingDown
1519+
// and short-circuit (see CleanupCallback / ErrorCallback).
1520+
Timer cleanup = Interlocked.Exchange(ref _cleanupTimer, null);
1521+
cleanup?.Dispose();
1522+
Timer error = Interlocked.Exchange(ref _errorTimer, null);
1523+
error?.Dispose();
1524+
1525+
// FR-007: wake any threads parked in WaitHandle.WaitAny by releasing the pool semaphore
1526+
// up to its max count. Waiters check State == Running after wake-up and bail.
1527+
// We may release more than necessary; SemaphoreFullException is harmless when the
1528+
// semaphore is already saturated.
1529+
try
1530+
{
1531+
_waitHandles.PoolSemaphore.Release(MaxPoolSize);
1532+
}
1533+
catch (SemaphoreFullException)
1534+
{
1535+
// Already at max count - all slots free. Nothing to do.
1536+
}
1537+
1538+
// FR-003: drain idle stacks and destroy contained connections. Active checked-out
1539+
// connections are destroyed when they are returned (see DeactivateObject's
1540+
// State is ShuttingDown branch).
1541+
while (_stackNew.TryPop(out DbConnectionInternal newObj))
1542+
{
1543+
DestroyObject(newObj);
1544+
}
1545+
while (_stackOld.TryPop(out DbConnectionInternal oldObj))
14901546
{
1491-
t.Dispose();
1547+
DestroyObject(oldObj);
14921548
}
14931549
}
14941550

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
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.Collections.Generic;
7+
using System.Data.Common;
8+
using System.Threading.Tasks;
9+
using Microsoft.Data.Common.ConnectionString;
10+
using Microsoft.Data.ProviderBase;
11+
using Microsoft.Data.SqlClient.ConnectionPool;
12+
using Xunit;
13+
14+
namespace Microsoft.Data.SqlClient.UnitTests.ConnectionPool
15+
{
16+
/// <summary>
17+
/// Deterministic tests for <see cref="ChannelDbConnectionPool"/> shutdown behavior.
18+
/// Verifies the contract defined by spec 004-pool-shutdown FR-001 .. FR-007 (P1 scope).
19+
/// </summary>
20+
public class ChannelDbConnectionPoolShutdownTest
21+
{
22+
private static ChannelDbConnectionPool ConstructPool(int maxPoolSize = 5)
23+
{
24+
var poolGroupOptions = new DbConnectionPoolGroupOptions(
25+
poolByIdentity: false,
26+
minPoolSize: 0,
27+
maxPoolSize: maxPoolSize,
28+
creationTimeout: 15,
29+
loadBalanceTimeout: 0,
30+
hasTransactionAffinity: true);
31+
32+
var dbConnectionPoolGroup = new DbConnectionPoolGroup(
33+
new SqlConnectionOptions("Data Source=localhost;"),
34+
new ConnectionPoolKey("TestDataSource", credential: null, accessToken: null, accessTokenCallback: null, sspiContextProvider: null),
35+
poolGroupOptions);
36+
37+
return new ChannelDbConnectionPool(
38+
new ChannelDbConnectionPoolTest.SuccessfulSqlConnectionFactory(),
39+
dbConnectionPoolGroup,
40+
DbConnectionPoolIdentity.NoIdentity,
41+
new DbConnectionPoolProviderInfo());
42+
}
43+
44+
// FR-001
45+
[Fact]
46+
public void Shutdown_TransitionsState_ToShuttingDown()
47+
{
48+
var pool = ConstructPool();
49+
Assert.True(pool.IsRunning);
50+
Assert.Equal(DbConnectionPoolState.Running, pool.State);
51+
52+
pool.Shutdown();
53+
54+
Assert.False(pool.IsRunning);
55+
Assert.Equal(DbConnectionPoolState.ShuttingDown, pool.State);
56+
}
57+
58+
// FR-003
59+
[Fact]
60+
public void Shutdown_DrainsIdleConnections()
61+
{
62+
var pool = ConstructPool();
63+
64+
// Vend and return three connections so they sit idle in the channel.
65+
var owners = new List<SqlConnection>();
66+
var conns = new List<DbConnectionInternal>();
67+
for (int i = 0; i < 3; i++)
68+
{
69+
var owner = new SqlConnection();
70+
owners.Add(owner);
71+
Assert.True(pool.TryGetConnection(owner, taskCompletionSource: null, out DbConnectionInternal? c));
72+
Assert.NotNull(c);
73+
conns.Add(c!);
74+
}
75+
for (int i = 0; i < conns.Count; i++)
76+
{
77+
pool.ReturnInternalConnection(conns[i], owners[i]);
78+
}
79+
Assert.Equal(3, pool.IdleCount);
80+
81+
pool.Shutdown();
82+
83+
Assert.Equal(0, pool.IdleCount);
84+
Assert.Equal(0, pool.Count);
85+
}
86+
87+
// FR-004 - returned connection while shutting down is destroyed, not pooled.
88+
[Fact]
89+
public void Shutdown_ReturnedConnection_IsDestroyedNotPooled()
90+
{
91+
var pool = ConstructPool();
92+
var owner = new SqlConnection();
93+
Assert.True(pool.TryGetConnection(owner, taskCompletionSource: null, out DbConnectionInternal? conn));
94+
Assert.NotNull(conn);
95+
Assert.Equal(1, pool.Count);
96+
Assert.Equal(0, pool.IdleCount);
97+
98+
pool.Shutdown();
99+
100+
// Connection is still checked out; return it now.
101+
pool.ReturnInternalConnection(conn!, owner);
102+
103+
Assert.Equal(0, pool.IdleCount);
104+
Assert.Equal(0, pool.Count);
105+
}
106+
107+
// FR-006
108+
[Fact]
109+
public void Shutdown_IsIdempotent()
110+
{
111+
var pool = ConstructPool();
112+
pool.Shutdown();
113+
// Second call must not throw and must leave state intact.
114+
pool.Shutdown();
115+
pool.Shutdown();
116+
Assert.Equal(DbConnectionPoolState.ShuttingDown, pool.State);
117+
}
118+
119+
// FR-007 - async waiter is unblocked when the pool shuts down.
120+
[Fact]
121+
public async Task Shutdown_UnblocksAsyncWaiter()
122+
{
123+
var pool = ConstructPool(maxPoolSize: 1);
124+
125+
// Saturate the pool.
126+
Assert.True(pool.TryGetConnection(new SqlConnection(), taskCompletionSource: null, out DbConnectionInternal? blocking));
127+
Assert.NotNull(blocking);
128+
129+
// Park an async waiter.
130+
var tcs = new TaskCompletionSource<DbConnectionInternal>();
131+
bool completed = pool.TryGetConnection(new SqlConnection(), tcs, out DbConnectionInternal? waiter);
132+
Assert.False(completed);
133+
Assert.Null(waiter);
134+
Assert.False(tcs.Task.IsCompleted);
135+
136+
// Shut down the pool.
137+
pool.Shutdown();
138+
139+
// Waiter must complete (faulted or with a null connection) within a bounded window.
140+
var winner = await Task.WhenAny(tcs.Task, Task.Delay(TimeSpan.FromSeconds(5)));
141+
Assert.Same(tcs.Task, winner);
142+
// Either an exception was set (channel closed) or the result is null - both are acceptable
143+
// shutdown signals. What matters is the waiter does NOT block forever.
144+
if (tcs.Task.IsFaulted)
145+
{
146+
// Expected path: ChannelClosedException or a wrapped exception.
147+
Assert.NotNull(tcs.Task.Exception);
148+
}
149+
else
150+
{
151+
// Permitted: completed with null result.
152+
Assert.Null(tcs.Task.Result);
153+
}
154+
}
155+
156+
// Story 3 acceptance #3 - sync get fails fast after shutdown.
157+
// The factory-level retry guard checks IsRunning, but the pool itself must not vend
158+
// new connections after Shutdown. We verify by exhausting the pool first then
159+
// checking that returned connections are destroyed (Count goes back to 0).
160+
[Fact]
161+
public void Shutdown_AfterShutdown_NewReturnsAreDestroyed()
162+
{
163+
var pool = ConstructPool();
164+
var owner = new SqlConnection();
165+
Assert.True(pool.TryGetConnection(owner, taskCompletionSource: null, out DbConnectionInternal? c));
166+
Assert.NotNull(c);
167+
168+
pool.Shutdown();
169+
pool.ReturnInternalConnection(c!, owner);
170+
171+
Assert.False(pool.IsRunning);
172+
Assert.Equal(0, pool.Count);
173+
Assert.Equal(0, pool.IdleCount);
174+
}
175+
176+
// Sanity: Startup is currently a no-op for this pool but must not throw or change
177+
// shutdown state if invoked after Shutdown.
178+
[Fact]
179+
public void Startup_AfterShutdown_DoesNotResurrectPool()
180+
{
181+
var pool = ConstructPool();
182+
pool.Shutdown();
183+
184+
pool.Startup();
185+
186+
Assert.Equal(DbConnectionPoolState.ShuttingDown, pool.State);
187+
Assert.False(pool.IsRunning);
188+
}
189+
}
190+
}

0 commit comments

Comments
 (0)