diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 8ccee11a7..dcea07d75 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -47,7 +47,10 @@ static SqlHandlePtr getEnvHandle() { // transaction control, and autocommit configuration. //------------------------------------------------------------------------------------------------- Connection::Connection(const std::u16string& conn_str, bool use_pool) - : _connStr(conn_str), _autocommit(false), _fromPool(use_pool) { + : _connStr(conn_str), + _autocommit(false), + _fromPool(use_pool), + _lastUsed(std::chrono::steady_clock::now()) { PERF_TIMER("Connection::Connection"); allocateDbcHandle(); } @@ -76,6 +79,10 @@ void Connection::allocateDbcHandle() { void Connection::connect(const py::dict& attrs_before) { PERF_TIMER("Connection::connect"); + if (_isMock) { + updateLastUsed(); + return; + } LOG("Connecting to database"); // Apply access token before connect if (!attrs_before.is_none() && py::len(attrs_before) > 0) { @@ -103,6 +110,10 @@ void Connection::connect(const py::dict& attrs_before) { void Connection::disconnect() { PERF_TIMER("Connection::disconnect"); + if (_isMock) { + _dbcHandle.reset(); + return; + } // Determine GIL state once, up front. disconnect() runs both from // pybind11-bound methods (GIL held) and from GIL-less destructor / shutdown // paths: Connection::~Connection() dropping the last shared_ptr, or teardown @@ -507,6 +518,9 @@ void Connection::applyAttrsBefore(const py::dict& attrs) { } bool Connection::isAlive() const { + if (_isMock) { + return true; + } if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } @@ -517,6 +531,9 @@ bool Connection::isAlive() const { } bool Connection::reset() { + if (_isMock) { + return true; + } if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } @@ -642,6 +659,9 @@ ConnectionHandle::ConnectionHandle(const std::u16string& connStr, bool usePool, } if (!_usePool) { _conn = std::make_shared(_connStr, false); + if (ConnectionPoolManager::getInstance().mock_mode()) { + _conn->setMock(true); + } // Non-pooled connect still honors the lazy token factory: a // token is materialized only when a physical connection is opened. The // factory may also carry the token expiry, but a non-pooled diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index d5300aff4..2273a31c4 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -55,6 +55,17 @@ class Connection { bool reset(); void updateLastUsed(); std::chrono::steady_clock::time_point lastUsed() const; + void setMock(bool mock) { _isMock = mock; } + bool isMock() const { return _isMock; } + void setPoolOrigin(uint64_t pool_id, uint64_t generation) { + _originPoolId = pool_id; + _originGeneration = generation; + } + bool matchesPoolOrigin(uint64_t pool_id, uint64_t generation) const { + return _originPoolId == pool_id && _originGeneration == generation; + } + uint64_t originPoolId() const { return _originPoolId; } + uint64_t originGeneration() const { return _originGeneration; } // Materialize connect-attrs from a Python token-factory callback. // The factory may return either a bare attrs dict (legacy) or a @@ -100,6 +111,9 @@ class Connection { std::u16string _connStr; bool _fromPool = false; bool _autocommit = true; + bool _isMock = false; + uint64_t _originPoolId = 0; + uint64_t _originGeneration = 0; SqlHandlePtr _dbcHandle; std::chrono::steady_clock::time_point _lastUsed; // POSIX-epoch expiry (seconds) of the access token this connection last @@ -150,6 +164,9 @@ class ConnectionHandle { // Get information about the driver and data source py::object getInfo(SQLUSMALLINT infoType) const; + uint64_t originGeneration() const { return _conn ? _conn->originGeneration() : 0; } + uint64_t originPoolId() const { return _conn ? _conn->originPoolId() : 0; } + private: std::shared_ptr _conn; bool _usePool; diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 831a01db2..f8ecad5f3 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -51,266 +51,414 @@ static std::string extractAccessToken(const py::dict& attrs) { return std::string(); } +// Process-wide monotonic counter for pool IDs to prevent ABA address-reuse (#746) +static std::atomic s_next_pool_id{1}; + ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs) : _max_size(max_size), _idle_timeout_secs(idle_timeout_secs), - _current_size(0) {} + _current_size(0), + _checked_out(0), + _in_flight(0), + _pool_id(s_next_pool_id.fetch_add(1)) {} + +void ConnectionPool::invokeDisconnectHook() { + std::shared_ptr> hook; + { + std::lock_guard lock(_mutex); + hook = _on_disconnect_hook; + } + if (hook) { + if (*hook) { + (*hook)(); + } + py::gil_scoped_acquire gil; + hook.reset(); + } +} + +void ConnectionPool::drainDisconnectList(std::vector>& list) { + for (auto& conn : list) { + if (!conn) { + continue; + } + invokeDisconnectHook(); + try { + conn->disconnect(); + } catch (const std::exception& ex) { + LOG("ConnectionPool::drainDisconnectList: disconnect failed: %s", ex.what()); + } + { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } + } + list.clear(); +} std::shared_ptr ConnectionPool::acquire(const std::u16string& connStr, const py::dict& attrs_before, const py::object& token_factory) { PERF_TIMER("ConnectionPool::acquire"); std::vector> to_disconnect; + struct DisconnectGuard { + ConnectionPool& pool; + std::vector>& list; + ~DisconnectGuard() { + pool.drainDisconnectList(list); + } + } guard{*this, to_disconnect}; std::shared_ptr valid_conn = nullptr; - bool needs_connect = false; - // Phase 1: Prune stale connections (under mutex — no ODBC calls). - { - std::lock_guard lock(_mutex); - auto now = std::chrono::steady_clock::now(); - size_t before = _pool.size(); - - _pool.erase(std::remove_if(_pool.begin(), _pool.end(), - [&](const std::shared_ptr& conn) { - auto idle_time = - std::chrono::duration_cast( - now - conn->lastUsed()) - .count(); - if (idle_time > _idle_timeout_secs) { - to_disconnect.push_back(conn); - return true; - } - return false; - }), - _pool.end()); - - size_t pruned = before - _pool.size(); - // Decrement _current_size eagerly so new slots can be reserved while - // stale connections are being disconnected (Phase 4). This means - // _current_size tracks *reserved capacity* (pooled + checked-out + - // in-flight new), not necessarily live ODBC handles. - _current_size = (_current_size >= pruned) ? (_current_size - pruned) : 0; - } + while (valid_conn == nullptr) { + bool needs_connect = false; + py::dict pending_attrs; + long long pending_expiry = 0; + bool have_pending_token = false; + uint64_t reservation_generation = 0; - // Phase 2: Pop one candidate at a time and validate it outside the - // mutex. isAlive() and reset() perform ODBC calls that release the - // GIL; calling them while holding the mutex would create a mutex/GIL - // lock-ordering deadlock when multiple threads acquire concurrently. - // - // Expiry-aware checkout may capture a freshly minted token here so a - // rotated-token pool can be reopened without invoking the factory twice. - py::dict pending_attrs; - long long pending_expiry = 0; - bool have_pending_token = false; - while (true) { - std::shared_ptr candidate; + // Phase 1: Prune stale connections (under mutex — no ODBC calls). { - std::unique_lock lock(_mutex); - if (_pool.empty()) { - // No more candidates — try to reserve a slot for a new connection. - if (_current_size < _max_size) { - // Reserve the slot here but construct the Connection outside - // _mutex (Phase 3): the Connection constructor allocates ODBC - // handles and emits log records that acquire the GIL, and - // holding _mutex across a GIL acquisition deadlocks a thread - // that holds the GIL and is waiting on _mutex (#671). - ++_current_size; - needs_connect = true; - break; + std::lock_guard lock(_mutex); + auto now = std::chrono::steady_clock::now(); + size_t before = _pool.size(); + + _pool.erase(std::remove_if(_pool.begin(), _pool.end(), + [&](const std::shared_ptr& conn) { + auto idle_time = + std::chrono::duration_cast( + now - conn->lastUsed()) + .count(); + if (idle_time > _idle_timeout_secs) { + to_disconnect.push_back(conn); + return true; + } + return false; + }), + _pool.end()); + + size_t pruned = before - _pool.size(); + // Retain capacity for pruned stale connections by accounting for them + // in _in_flight until Phase 4 disconnects them outside the mutex (#746). + _in_flight += pruned; + } + + // Disconnect pruned stale connections outside lock BEFORE attempting + // candidate validation or slot reservation in Phase 2/3. As each disconnect + // finishes, drainDisconnectList decrements _in_flight and _current_size. + drainDisconnectList(to_disconnect); + + // Phase 2: Pop one candidate at a time and validate it outside the + // mutex. isAlive() and reset() perform ODBC calls that release the + // GIL; calling them while holding the mutex would create a mutex/GIL + // lock-ordering deadlock when multiple threads acquire concurrently. + // + // Expiry-aware checkout may capture a freshly minted token here so a + // rotated-token pool can be reopened without invoking the factory twice. + while (true) { + std::shared_ptr candidate; + uint64_t candidate_generation = 0; + { + std::unique_lock lock(_mutex); + if (_pool.empty()) { + // No more candidates — try to reserve a slot for a new connection. + if (_current_size < _max_size) { + // Reserve the slot here but construct the Connection outside + // _mutex (Phase 3): the Connection constructor allocates ODBC + // handles and emits log records that acquire the GIL, and + // holding _mutex across a GIL acquisition deadlocks a thread + // that holds the GIL and is waiting on _mutex (#671). + ++_current_size; + ++_in_flight; + reservation_generation = _generation; + needs_connect = true; + break; + } + // Pool is full — throw immediately. Another thread may be + // validating a popped candidate outside the mutex right now, so + // a transient "pool full" is an acceptable trade-off that + // callers can retry. + throw std::runtime_error( + "ConnectionPool::acquire: pool size limit reached"); } - // Pool is full — throw immediately. Another thread may be - // validating a popped candidate outside the mutex right now, so - // a transient "pool full" is an acceptable trade-off that - // callers can retry. - throw std::runtime_error( - "ConnectionPool::acquire: pool size limit reached"); + candidate = _pool.front(); + _pool.pop_front(); + ++_in_flight; + candidate_generation = _generation; } - candidate = _pool.front(); - _pool.pop_front(); - } - // Validate the candidate outside the mutex. - bool reuse_candidate = false; - try { - if (token_factory && !token_factory.is_none() && - candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) { - // Expiry-aware checkout with token compare: the pooled token is - // at/near expiry, so mint a fresh one and compare. If the - // provider returns the SAME token (its cache is still valid), - // the connection is healthy — refresh the recorded expiry and - // reuse it rather than needlessly churning. Only a DIFFERENT - // (rotated) token forces discard-and-reopen, and we carry the - // fresh attrs forward so the reopen below does not invoke the - // factory a second time. - long long fresh_expiry = 0; - py::dict fresh_attrs = - Connection::invokeTokenFactory(token_factory, fresh_expiry); - std::string fresh_token = extractAccessToken(fresh_attrs); - if (!fresh_token.empty() && - fresh_token == candidate->currentAccessToken() && - tokenExpirySafelyBeyond(fresh_expiry, TOKEN_EXPIRY_THRESHOLD_SECS)) { - // Same token AND its refreshed expiry is safely beyond the - // threshold: the provider's cache is still valid and the - // connection is healthy, so refresh the recorded expiry and - // reuse. We deliberately do NOT reuse when the returned - // expiry is unknown (<=0) or still inside the threshold — - // extending the recorded expiry and handing the connection - // back would defeat the very refresh this checkout intended - // (the token could expire mid-query). Those cases fall - // through to discard-and-reopen below. - // - // Narrow edge: a MISBEHAVING provider that repeatedly hands - // back the same token still inside the threshold makes every - // checkout discard + reopen (and get the same near-expiry - // token) — pure churn, no benefit. This is acceptable: a - // well-behaved azure-identity credential refreshes - // proactively (returning a token with a fresh, far-out - // expiry) before the threshold, so the safe-reuse path above - // is taken in practice. We favor never handing out a token - // that may expire mid-query over avoiding the churn. - candidate->setTokenExpiry(fresh_expiry); - reuse_candidate = candidate->isAlive() && candidate->reset(); - if (!reuse_candidate) { - // The token is still valid but the socket is dead - // (isAlive()/reset() failed). We already minted the - // fresh attrs, so carry them forward and let Phase 3 - // reopen with them instead of invoking the factory a - // second time. No sibling drain: siblings hold the same - // still-valid token and remain reusable. - pending_attrs = fresh_attrs; - pending_expiry = fresh_expiry; - have_pending_token = true; + // Validate the candidate outside the mutex. + bool reuse_candidate = false; + try { + if (token_factory && !token_factory.is_none() && + candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) { + // Expiry-aware checkout with token compare: the pooled token is + // at/near expiry, so mint a fresh one and compare. If the + // provider returns the SAME token (its cache is still valid), + // the connection is healthy — refresh the recorded expiry and + // reuse it rather than needlessly churning. Only a DIFFERENT + // (rotated) token forces discard-and-reopen, and we carry the + // fresh attrs forward so the reopen below does not invoke the + // factory a second time. + long long fresh_expiry = 0; + py::dict fresh_attrs = + Connection::invokeTokenFactory(token_factory, fresh_expiry); + if (candidate->isMock()) { + candidate->setTokenExpiry(fresh_expiry); + reuse_candidate = true; + } else { + std::string fresh_token = extractAccessToken(fresh_attrs); + if (!fresh_token.empty() && + fresh_token == candidate->currentAccessToken() && + tokenExpirySafelyBeyond(fresh_expiry, TOKEN_EXPIRY_THRESHOLD_SECS)) { + candidate->setTokenExpiry(fresh_expiry); + reuse_candidate = candidate->isAlive() && candidate->reset(); + if (!reuse_candidate) { + pending_attrs = fresh_attrs; + pending_expiry = fresh_expiry; + have_pending_token = true; + } + } else { + pending_attrs = fresh_attrs; + pending_expiry = fresh_expiry; + have_pending_token = true; + const std::string stale_token = candidate->currentAccessToken(); + if (!stale_token.empty()) { + std::lock_guard lock(_mutex); + if (_generation == candidate_generation) { + _pool.erase( + std::remove_if( + _pool.begin(), _pool.end(), + [&](const std::shared_ptr& sibling) { + if (sibling->currentAccessToken() == stale_token) { + to_disconnect.push_back(sibling); + ++_in_flight; + return true; + } + return false; + }), + _pool.end()); + } + } + } } } else { - // Token rotated, or the "fresh" token is still at/near - // expiry (or has an unknown expiry): discard and reopen with - // the fresh attrs. Remember the fresh token to reopen with, - // and eagerly drain the sibling idle connections that still - // hold the now-stale token. They were all minted from the - // same provider before the rotation, so they are equally - // stale; discarding them together here avoids rediscovering - // each one (and paying another factory compare) on later - // checkouts. No ODBC calls under the mutex — the actual - // disconnects happen in Phase 4, outside the lock. - pending_attrs = fresh_attrs; - pending_expiry = fresh_expiry; - have_pending_token = true; - const std::string stale_token = candidate->currentAccessToken(); - if (!stale_token.empty()) { - std::lock_guard lock(_mutex); - _pool.erase( - std::remove_if( - _pool.begin(), _pool.end(), - [&](const std::shared_ptr& sibling) { - if (sibling->currentAccessToken() == stale_token) { - to_disconnect.push_back(sibling); - if (_current_size > 0) --_current_size; - return true; - } - return false; - }), - _pool.end()); - } + reuse_candidate = candidate->isAlive() && candidate->reset(); } - } else { - reuse_candidate = candidate->isAlive() && candidate->reset(); + } catch (const std::exception& ex) { + LOG("Candidate connection validation failed: %s", ex.what()); } - } catch (const std::exception& ex) { - LOG("Candidate connection validation failed: %s", ex.what()); - } - if (reuse_candidate) { - valid_conn = candidate; - break; - } - - // Candidate is dead, reset failed, or its token rotated — mark for - // disconnect and decrement the pool size. - to_disconnect.push_back(candidate); - { - std::lock_guard lock(_mutex); - if (_current_size > 0) --_current_size; - } - - // If a rotated token was captured, reserve a slot and reopen with it - // immediately instead of churning through the remaining candidates - // (which hold the same stale token and would all be discarded anyway). - if (have_pending_token) { - std::lock_guard lock(_mutex); - if (_current_size < _max_size) { - // Reserve the slot here but construct the Connection outside - // _mutex (Phase 3): the constructor emits GIL-acquiring log - // records, and holding _mutex across a GIL acquisition - // deadlocks a thread that holds the GIL and waits on _mutex (#671). - ++_current_size; - needs_connect = true; - break; + if (reuse_candidate) { + bool gen_valid = false; + { + std::lock_guard lock(_mutex); + if (_generation == candidate_generation) { + candidate->updateLastUsed(); + candidate->setPoolOrigin(_pool_id, _generation); + valid_conn = candidate; + ++_checked_out; + if (_in_flight > 0) { + --_in_flight; + } + gen_valid = true; + } + } + if (gen_valid) { + break; + } + // Pool was closed while validating candidate (#746); discard stale + // candidate and release in-flight reservation. + try { + candidate->disconnect(); + } catch (const std::exception& ex) { + LOG("Disconnect candidate failed: %s", ex.what()); + } + { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } + continue; } - // Pool momentarily full; fall through and retry the loop. On the - // retry another near-expiry candidate may re-invoke the factory and - // overwrite pending_attrs/pending_expiry with a newer token. That - // needs a full pool AND a simultaneous rotation, is rare, and is - // harmless: we simply reopen with the most recently minted token. - } - } - // Phase 3: Construct and connect the new connection outside the mutex. - if (needs_connect) { - try { - // Construct the Connection outside _mutex (#671): the constructor - // allocates ODBC handles and emits log records that acquire the GIL, - // so it must not run while _mutex is held. - valid_conn = std::make_shared(connStr, true); - if (have_pending_token) { - // Reopen with the fresh token captured during expiry-aware - // checkout (the previous connection's token had rotated). - valid_conn->connect(pending_attrs); - valid_conn->setTokenExpiry(pending_expiry); - } else if (token_factory && !token_factory.is_none()) { - // Lazy token acquisition: only now, when a physical - // connection is actually being opened, do we materialize the - // token. On a pool reuse this whole branch is skipped, so a - // same-identity hit never acquires a token. The GIL is held here - // (connect() releases it only around the ODBC call itself), so - // invoking the Python callback is safe. - long long expiry = 0; - py::dict connect_attrs = Connection::invokeTokenFactory(token_factory, expiry); - valid_conn->connect(connect_attrs); - // Record the token expiry so a later checkout can refresh this - // connection before the token lapses. - valid_conn->setTokenExpiry(expiry); - } else { - valid_conn->connect(attrs_before); + // Candidate is dead, reset failed, or its token rotated — disconnect and + // release the in-flight reservation (#746). + try { + candidate->disconnect(); + } catch (const std::exception& ex) { + LOG("Disconnect candidate failed: %s", ex.what()); } - } catch (...) { - // Construct/connect failed — release the reserved slot { std::lock_guard lock(_mutex); - if (_current_size > 0) --_current_size; + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } + + // If a rotated token was captured, reserve a slot and reopen with it + // immediately instead of churning through the remaining candidates + // (which hold the same stale token and would all be discarded anyway). + if (have_pending_token) { + // Drain any siblings placed into to_disconnect before reserving a new slot + drainDisconnectList(to_disconnect); + std::lock_guard lock(_mutex); + if (_current_size < _max_size) { + // Reserve the slot here but construct the Connection outside + // _mutex (Phase 3): the constructor emits GIL-acquiring log + // records, and holding _mutex across a GIL acquisition + // deadlocks a thread that holds the GIL and waits on _mutex (#671). + ++_current_size; + ++_in_flight; + reservation_generation = _generation; + needs_connect = true; + break; + } + // Pool momentarily full; fall through and retry the loop. On the + // retry another near-expiry candidate may re-invoke the factory and + // overwrite pending_attrs/pending_expiry with a newer token. That + // needs a full pool AND a simultaneous rotation, is rare, and is + // harmless: we simply reopen with the most recently minted token. } - throw; } - } - // Phase 4: Disconnect expired/bad connections outside lock. - for (auto& conn : to_disconnect) { - try { - conn->disconnect(); - } catch (const std::exception& ex) { - LOG("Disconnect bad/expired connections failed: %s", ex.what()); + if (valid_conn != nullptr) { + break; + } + + // Phase 3: Construct and connect the new connection outside the mutex. + if (needs_connect) { + try { + // Construct the Connection outside _mutex (#671): the constructor + // allocates ODBC handles and emits log records that acquire the GIL, + // so it must not run while _mutex is held. + auto new_conn = std::make_shared(connStr, true); + if (_mock_mode) { + new_conn->setMock(true); + } + if (have_pending_token) { + // Reopen with the fresh token captured during expiry-aware + // checkout (the previous connection's token had rotated). + new_conn->connect(pending_attrs); + new_conn->setTokenExpiry(pending_expiry); + } else if (token_factory && !token_factory.is_none()) { + // Lazy token acquisition: only now, when a physical + // connection is actually being opened, do we materialize the + // token. On a pool reuse this whole branch is skipped, so a + // same-identity hit never acquires a token. The GIL is held here + // (connect() releases it only around the ODBC call itself), so + // invoking the Python callback is safe. + long long expiry = 0; + py::dict connect_attrs = Connection::invokeTokenFactory(token_factory, expiry); + new_conn->connect(connect_attrs); + // Record the token expiry so a later checkout can refresh this + // connection before the token lapses. + new_conn->setTokenExpiry(expiry); + } else { + new_conn->connect(attrs_before); + } + + // Verify that pool was not closed while connecting outside the mutex (#746). + bool gen_valid = false; + { + std::lock_guard lock(_mutex); + if (_generation == reservation_generation) { + new_conn->updateLastUsed(); + new_conn->setPoolOrigin(_pool_id, _generation); + valid_conn = new_conn; + ++_checked_out; + if (_in_flight > 0) { + --_in_flight; + } + gen_valid = true; + } + } + if (gen_valid) { + break; + } + // Pool was closed while connecting. Disconnect the stale connection + // immediately BEFORE relinquishing the in-flight reservation, so that + // the stale physical connection and any newly reserved connections do + // not co-exist and exceed max_size (#746). + invokeDisconnectHook(); + try { + new_conn->disconnect(); + } catch (const std::exception& ex) { + LOG("Disconnect stale connection failed: %s", ex.what()); + } + { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } + continue; + } catch (...) { + // Construct/connect failed — release the reserved slot and in-flight count. + { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } + throw; + } } } + + // Phase 4: Disconnect expired/bad connections outside lock and decrement in-flight capacity. + drainDisconnectList(to_disconnect); return valid_conn; } void ConnectionPool::release(std::shared_ptr conn) { PERF_TIMER("ConnectionPool::release"); + if (!conn) { + return; + } bool should_disconnect = false; + bool decrement_in_flight = false; { std::lock_guard lock(_mutex); - if (_pool.size() < _max_size) { - conn->updateLastUsed(); - _pool.push_back(conn); + if (conn->originPoolId() == _pool_id) { + bool generation_matches = (conn->originGeneration() == _generation); + if (generation_matches && _pool.size() < _max_size) { + conn->updateLastUsed(); + _pool.push_back(conn); + if (_checked_out > 0) { + --_checked_out; + } + conn->setPoolOrigin(0, 0); + } else { + should_disconnect = true; + if (_checked_out > 0) { + --_checked_out; + // Keep this connection accounted for as in-flight until disconnect + // completes outside the mutex, so a concurrent acquire cannot reserve + // and open a new physical handle while this handle is still live (#746). + ++_in_flight; + decrement_in_flight = true; + } + conn->setPoolOrigin(0, 0); + } } else { should_disconnect = true; } @@ -318,14 +466,21 @@ void ConnectionPool::release(std::shared_ptr conn) { // Disconnect outside the mutex to avoid holding it during the // blocking ODBC call (which releases the GIL). if (should_disconnect) { + invokeDisconnectHook(); try { conn->disconnect(); } catch (const std::exception& ex) { LOG("ConnectionPool::release: disconnect failed: %s", ex.what()); } - std::lock_guard lock(_mutex); - if (_current_size > 0) - --_current_size; + if (decrement_in_flight) { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } } } @@ -334,20 +489,16 @@ bool ConnectionPool::canEvict() { // Never evict while any connection is checked out or in-flight. Reserved // capacity (_current_size) beyond what is sitting idle in _pool means a // caller still holds one, so the pool must stay. - size_t checked_out = (_current_size > _pool.size()) ? (_current_size - _pool.size()) : 0; - if (checked_out > 0) { + size_t in_flight_or_checked_out = + (_current_size > _pool.size()) ? (_current_size - _pool.size()) : 0; + if (in_flight_or_checked_out > 0 || _checked_out > 0 || _in_flight > 0) { return false; } // Nothing checked out and the pool is empty: safe to drop immediately. if (_pool.empty()) { return true; } - // Nothing checked out but idle connections remain. Evict the whole pool - // only once EVERY pooled connection has been idle longer than the idle - // timeout. This is what reclaims pools for rotating / single-use identities - // (e.g. per-request Entra users keyed by token hash): such a pool is never - // acquired again, so its idle connection is never pruned by acquire() and - // _current_size would otherwise stay > 0 forever. Evaluating the idle + // Empty pools past idle timeout can be evicted. Checking the idle // timeout here lets the next acquireConnection() on any key sweep it away. auto now = std::chrono::steady_clock::now(); for (const auto& conn : _pool) { @@ -369,14 +520,29 @@ void ConnectionPool::close() { to_close.push_back(_pool.front()); _pool.pop_front(); } - _current_size = 0; + // Account for closing idle connections in _in_flight so a concurrent + // acquire cannot reserve and open a new physical handle while these + // old handles are still connected outside the mutex (#746). + _in_flight += to_close.size(); + _current_size = _checked_out + _in_flight; + ++_generation; } for (auto& conn : to_close) { + invokeDisconnectHook(); try { conn->disconnect(); } catch (const std::exception& ex) { LOG("ConnectionPool::close: disconnect failed: %s", ex.what()); } + { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } } } @@ -385,71 +551,167 @@ ConnectionPoolManager& ConnectionPoolManager::getInstance() { return manager; } -std::shared_ptr ConnectionPoolManager::acquireConnection(const std::u16string& connStr, - const py::dict& attrs_before, - const std::u16string& pool_key, - const py::object& token_factory) { +std::shared_ptr ConnectionPoolManager::acquireConnection( + const std::u16string& connStr, + const py::dict& attrs_before, + const std::u16string& pool_key, + const py::object& token_factory) { PERF_TIMER("ConnectionPoolManager::acquireConnection"); // Key the pool by pool_key when provided (identity-aware), // else fall back to the connection string (legacy behavior). const std::u16string& key = pool_key.empty() ? connStr : pool_key; std::shared_ptr pool; + std::shared_ptr old_pool_to_close; bool created = false; - std::vector> evicted; + std::vector>> evicted; + + // RAII guard ensuring any key placed in _closing_keys is removed and + // _manager_cv notified even if an exception or early return occurs (#746). + struct ClosingGuard { + ConnectionPoolManager& mgr; + std::vector keys; + ~ClosingGuard() { + if (!keys.empty()) { + std::lock_guard lock(mgr._manager_mutex); + for (const auto& k : keys) { + mgr._closing_keys.erase(k); + } + mgr._manager_cv.notify_all(); + } + } + void remove(const std::u16string& k) { + mgr._closing_keys.erase(k); + keys.erase(std::remove(keys.begin(), keys.end(), k), keys.end()); + mgr._manager_cv.notify_all(); + } + } closing_guard{*this}; + { - std::lock_guard lock(_manager_mutex); - // Pooling disabled (a concurrent disable_pooling() disarmed us): decline - // to create or hand out a pool. Because this check and the pool creation - // below share _manager_mutex with the setAccepting(false) in - // disable_pooling(), the decision is atomic — a connect either creates - // its pool before the disable (and closePools() then reaps it) or sees - // _accepting == false here and never creates one. The caller - // (ConnectionHandle) falls back to a non-pooled connection. - if (!_accepting) { - return nullptr; + py::gil_scoped_release release_gil; + { + std::unique_lock lock(_manager_mutex); + // Wait if this key is currently undergoing close/replacement by another thread, + // or until pooling is disabled. Serializes replacement creation with old-pool teardown (#746). + _manager_cv.wait(lock, [this, &key]() { + return !_accepting || _closing_keys.find(key) == _closing_keys.end(); + }); + + // Pooling disabled (a concurrent disable_pooling() disarmed us): decline + // to create or hand out a pool. Because this check and the pool creation + // below share _manager_mutex with the setAccepting(false) in + // disable_pooling(), the decision is atomic — a connect either creates + // its pool before the disable (and closePools() then reaps it) or sees + // _accepting == false here and never creates one. The caller + // (ConnectionHandle) falls back to a non-pooled connection. + if (!_accepting) { + return nullptr; + } + + // Lazy eviction: drop pools whose connections are all idle past the + // idle timeout (and none checked out) so distinct short-lived + // identities (e.g. per-request Entra users keyed by token hash) do not + // accumulate pools forever. canEvict() only inspects state (no ODBC + // calls), so it is safe under _manager_mutex; the actual disconnects + // happen via close() below, outside the lock. The pool we are about to + // use is skipped so it is never evicted from under us. + // + // The sweep is O(pools × idle-conns) under the global mutex, so it is + // throttled: a pool can only become evictable after its connections + // sit idle past the idle timeout, so sweeping more often than that + // window is pure overhead. Between sweeps we skip straight to the pool + // lookup, keeping the hot path cheap under a many-identity connect load. + auto now = std::chrono::steady_clock::now(); + auto sweep_interval = std::chrono::seconds(std::max(1, _default_idle_secs)); + if (now - _last_sweep >= sweep_interval) { + _last_sweep = now; + for (auto it = _pools.begin(); it != _pools.end();) { + // Only evict a pool that no one else is holding: use_count == 1 + // means the map is the sole owner. An in-flight acquirer copies + // its pool shared_ptr while holding _manager_mutex (same section + // as this sweep) and keeps that copy across the unlocked + // acquire(); returnConnection() likewise takes a ref under the + // mutex before releasing. Either bumps use_count above 1 for the + // whole window, so this guard prevents evicting — and then + // closing (disconnecting) — a pool a peer thread has already + // selected but not yet finished using. + if (it->first != key && it->second && it->second.use_count() == 1 && + it->second->canEvict()) { + _closing_keys.insert(it->first); + closing_guard.keys.push_back(it->first); + evicted.push_back({it->first, it->second}); + it = _pools.erase(it); + } else { + ++it; + } + } + } + // Defer replacement-pool creation if the existing pool still has live + // work. If the existing pool has finished all live work (canEvict() == true) + // and is not held by concurrent acquirers (use_count() == 1), evict it + // and serialize its close BEFORE publishing a new replacement pool (#746). + auto it = _pools.find(key); + if (it != _pools.end() && it->second && it->second.use_count() == 1 && + it->second->canEvict()) { + old_pool_to_close = it->second; + _pools.erase(it); + _closing_keys.insert(key); + closing_guard.keys.push_back(key); + } + if (!old_pool_to_close) { + auto& pool_ref = _pools[key]; + if (!pool_ref) { + pool_ref = std::make_shared(_default_max_size, _default_idle_secs); + if (_mock_mode) { + pool_ref->set_mock_mode(true); + } + created = true; + } + pool = pool_ref; + } } - // Lazy eviction: drop pools whose connections are all idle past the - // idle timeout (and none checked out) so distinct short-lived - // identities (e.g. per-request Entra users keyed by token hash) do not - // accumulate pools forever. canEvict() only inspects state (no ODBC - // calls), so it is safe under _manager_mutex; the actual disconnects - // happen via close() below, outside the lock. The pool we are about to - // use is skipped so it is never evicted from under us. - // - // The sweep is O(pools × idle-conns) under the global mutex, so it is - // throttled: a pool can only become evictable after its connections - // sit idle past the idle timeout, so sweeping more often than that - // window is pure overhead. Between sweeps we skip straight to the pool - // lookup, keeping the hot path cheap under a many-identity connect load. - auto now = std::chrono::steady_clock::now(); - auto sweep_interval = std::chrono::seconds(std::max(1, _default_idle_secs)); - if (now - _last_sweep >= sweep_interval) { - _last_sweep = now; - for (auto it = _pools.begin(); it != _pools.end();) { - // Only evict a pool that no one else is holding: use_count == 1 - // means the map is the sole owner. An in-flight acquirer copies - // its pool shared_ptr while holding _manager_mutex (same section - // as this sweep) and keeps that copy across the unlocked - // acquire(); returnConnection() likewise takes a ref under the - // mutex before releasing. Either bumps use_count above 1 for the - // whole window, so this guard prevents evicting — and then - // closing (disconnecting) — a pool a peer thread has already - // selected but not yet finished using. - if (it->first != key && it->second && it->second.use_count() == 1 && - it->second->canEvict()) { - evicted.push_back(it->second); - it = _pools.erase(it); - } else { - ++it; + if (old_pool_to_close) { + // Close the old pool completely BEFORE creating and publishing the replacement, + // ensuring its physical handles are disconnected before new ones can be opened (#746). + try { + old_pool_to_close->close(); + } catch (const std::exception& ex) { + LOG("ConnectionPoolManager: closing evicted pool failed: %s", ex.what()); + } + old_pool_to_close.reset(); + { + std::lock_guard lock(_manager_mutex); + if (_accepting) { + auto& pool_ref = _pools[key]; + if (!pool_ref) { + pool_ref = std::make_shared(_default_max_size, _default_idle_secs); + if (_mock_mode) { + pool_ref->set_mock_mode(true); + } + created = true; + } + pool = pool_ref; } + closing_guard.remove(key); + } + if (!_accepting) { + return nullptr; } } - auto& pool_ref = _pools[key]; - if (!pool_ref) { - pool_ref = std::make_shared(_default_max_size, _default_idle_secs); - created = true; + // Close evicted pools outside _manager_mutex: close() disconnects ODBC + // handles (releasing the GIL), which must never run while holding + // _manager_mutex or we risk a mutex/GIL lock-ordering deadlock. + for (auto& [evicted_key, evicted_pool] : evicted) { + try { + evicted_pool->close(); + } catch (const std::exception& ex) { + LOG("ConnectionPoolManager: closing evicted pool failed: %s", ex.what()); + } + evicted_pool.reset(); + { + std::lock_guard lock(_manager_mutex); + closing_guard.remove(evicted_key); + } } - pool = pool_ref; } // Log after releasing _manager_mutex (#671): LOG() acquires the GIL, and // holding a native mutex across a GIL acquisition deadlocks a thread that @@ -457,16 +719,6 @@ std::shared_ptr ConnectionPoolManager::acquireConnection(const std:: if (created) { LOG("Creating new connection pool"); } - // Close evicted pools outside _manager_mutex: close() disconnects ODBC - // handles (releasing the GIL), which must never run while holding - // _manager_mutex or we risk a mutex/GIL lock-ordering deadlock. - for (auto& evicted_pool : evicted) { - try { - evicted_pool->close(); - } catch (const std::exception& ex) { - LOG("ConnectionPoolManager: closing evicted pool failed: %s", ex.what()); - } - } // Call acquire() outside _manager_mutex. acquire() may release the GIL // during the ODBC connect call; holding _manager_mutex across that would // create a mutex/GIL lock-ordering deadlock. connStr (not key) is used to @@ -516,26 +768,22 @@ void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) { } void ConnectionPoolManager::closePools() { - // Mirror the eviction-sweep pattern: under the mutex, move every pool into - // a local vector and clear the map, then release the mutex before closing. - // close() disconnects ODBC handles (releasing the GIL), which must never - // run while holding _manager_mutex or we risk a mutex/GIL lock-ordering - // deadlock with a concurrent acquireConnection()/returnConnection(). + py::gil_scoped_release release_gil; + // Under _manager_mutex, snapshot all pools to close their idle connections. + // Wait for any in-flight same-key pool replacements to finish closing first (#746). std::vector> to_close; { - std::lock_guard lock(_manager_mutex); + std::unique_lock lock(_manager_mutex); + _manager_cv.wait(lock, [this]() { return _closing_keys.empty(); }); to_close.reserve(_pools.size()); for (auto& [conn_str, pool] : _pools) { if (pool) { to_close.push_back(pool); } } - _pools.clear(); - // Nothing left to sweep; reset the throttle so a fresh pool set after - // this is swept on its next acquireConnection(). - _last_sweep = std::chrono::steady_clock::time_point{}; } - // Close each pool outside _manager_mutex. + // Close each pool outside _manager_mutex: close() drains idle connections, + // bumps _generation, and sets _current_size = _checked_out + _in_flight (#746). for (auto& pool : to_close) { try { pool->close(); @@ -543,9 +791,33 @@ void ConnectionPoolManager::closePools() { LOG("ConnectionPoolManager::closePools: closing pool failed: %s", ex.what()); } } + to_close.clear(); + { + std::lock_guard lock(_manager_mutex); + // Only evict pools that have no live work left (canEvict() == true) and + // are not held by any concurrent thread (use_count() == 1). + // If an old pool still has checked-out connections or in-flight opens, + // retain it in _pools so that: + // 1. Creation of a replacement pool is deferred until the old pool has + // no live work, preventing capacity overflow across recreation (#746). + // 2. Any subsequent acquireConnection() respects live capacity. + // 3. returnConnection() continues to route to this pool to decrement + // _checked_out and _current_size as connections are released. + for (auto it = _pools.begin(); it != _pools.end();) { + if (!it->second || (it->second.use_count() == 1 && it->second->canEvict())) { + it = _pools.erase(it); + } else { + ++it; + } + } + // Reset the sweep throttle so a fresh pool set after this is swept + // on its next acquireConnection(). + _last_sweep = std::chrono::steady_clock::time_point{}; + } } void ConnectionPoolManager::setAccepting(bool accepting) { std::lock_guard lock(_manager_mutex); _accepting = accepting; + _manager_cv.notify_all(); } diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index edc87c865..5bbaebb22 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -6,12 +6,16 @@ #pragma once #include "connection/connection.h" +#include #include +#include #include +#include #include #include #include #include +#include // Manages a fixed-size pool of reusable database connections for a // single connection string @@ -37,14 +41,69 @@ class ConnectionPool { // Closes all connections in the pool, releasing resources void close(); + // Drains and disconnects connections outside the lock, decrementing in-flight capacity + void drainDisconnectList(std::vector>& list); + // True when the pool holds no live or in-flight connections and can be // dropped by the manager to reclaim memory (lazy eviction). bool canEvict(); + // Test accessors for pool generation, checked-out count, and current size + size_t current_size() const { + std::lock_guard lock(const_cast(_mutex)); + return _current_size; + } + size_t checked_out() const { + std::lock_guard lock(const_cast(_mutex)); + return _checked_out; + } + size_t in_flight() const { + std::lock_guard lock(const_cast(_mutex)); + return _in_flight; + } + uint64_t generation() const { + std::lock_guard lock(const_cast(_mutex)); + return _generation; + } + uint64_t pool_id() const { + return _pool_id; + } + + // Test helper to inject a candidate connection for race testing + void inject_candidate(std::shared_ptr conn) { + std::lock_guard lock(_mutex); + _pool.push_back(conn); + ++_current_size; + } + + // Test hooks for deterministic race testing (#746) + void set_mock_mode(bool enable) { + _mock_mode = enable; + } + bool mock_mode() const { + return _mock_mode; + } + void set_on_disconnect_hook(std::shared_ptr> hook) { + std::shared_ptr> old_hook; + { + std::lock_guard lock(_mutex); + old_hook = std::move(_on_disconnect_hook); + _on_disconnect_hook = std::move(hook); + } + } + private: + void invokeDisconnectHook(); + size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale size_t _current_size = 0; + size_t _checked_out = 0; // Live connections currently checked out by callers (#746) + size_t _in_flight = 0; // Connects or validations currently in flight (#746) + uint64_t _generation = 0; // Pool reset generation for reservation attribution (#746) + uint64_t _pool_id = 0; // Monotonic process-wide pool ID to avoid ABA reuse (#746) + std::atomic _mock_mode{false}; + std::shared_ptr> _on_disconnect_hook; std::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access }; @@ -85,6 +144,28 @@ class ConnectionPoolManager { // Closes all pools and their connections void closePools(); + // Test hooks for mock mode + void set_mock_mode(bool enable) { + std::lock_guard lock(_manager_mutex); + _mock_mode = enable; + for (auto& [_, pool] : _pools) { + if (pool) { + pool->set_mock_mode(enable); + } + } + } + bool mock_mode() const { + std::lock_guard lock(const_cast(_manager_mutex)); + return _mock_mode; + } + + // Test accessor to look up an existing pool for deterministic testing (#746) + std::shared_ptr getPool(const std::u16string& key) { + std::lock_guard lock(_manager_mutex); + auto it = _pools.find(key); + return (it != _pools.end()) ? it->second : nullptr; + } + private: ConnectionPoolManager() = default; ~ConnectionPoolManager() = default; @@ -92,6 +173,11 @@ class ConnectionPoolManager { // Map from connection string to connection pool std::unordered_map> _pools; + // Keys whose pools are currently being closed outside _manager_mutex (#746). + // Serializes same-key replacement creation with old-pool teardown. + std::unordered_set _closing_keys; + std::condition_variable _manager_cv; + // Protects access to the _pools map std::mutex _manager_mutex; size_t _default_max_size = 10; @@ -104,6 +190,7 @@ class ConnectionPoolManager { // explicit enable_pooling() call; only disable_pooling() disarms it, and // enable_pooling() re-arms it. bool _accepting = true; + bool _mock_mode = false; // Throttle for the lazy-eviction sweep in acquireConnection(). The sweep // iterates every pool (and every idle connection within each) under diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 35a3b6da4..0e92f8076 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -20,6 +20,7 @@ #include // For std::memcpy #include #include +#include #include // std::forward #include // CPython datetime API (PyDateTime_IMPORT, PyDateTime_GET_*, etc.) @@ -6108,9 +6109,17 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("set_attr", &ConnectionHandle::setAttr, py::arg("attribute"), py::arg("value"), "Set connection attribute") .def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle) - .def("get_info", &ConnectionHandle::getInfo, py::arg("info_type")); + .def("get_info", &ConnectionHandle::getInfo, py::arg("info_type")) + .def_property_readonly("origin_generation", &ConnectionHandle::originGeneration) + .def_property_readonly("origin_pool_id", &ConnectionHandle::originPoolId); m.def("enable_pooling", &enable_pooling, "Enable global connection pooling"); m.def("close_pooling", []() { ConnectionPoolManager::getInstance().closePools(); }); + m.def("_set_pool_manager_mock_mode", [](bool enable) { + ConnectionPoolManager::getInstance().set_mock_mode(enable); + }, py::arg("enable") = true); + m.def("_get_pool_for_key", [](const std::u16string& key) { + return ConnectionPoolManager::getInstance().getPool(key); + }, py::arg("key"), "Get internal pool instance for testing (#746)"); m.def("disable_pooling", []() { // Disarm new-pool creation *before* closing so a connect racing this // disable cannot resurrect a pool after the map is cleared: any @@ -6121,6 +6130,88 @@ PYBIND11_MODULE(ddbc_bindings, m) { manager.setAccepting(false); manager.closePools(); }, "Disable global connection pooling and close all pools"); + // Internal test seam: allows deterministic unit testing of ConnectionPool + // concurrency and generation tracking (#746). + py::class_>(m, "_TestPooledConnection") + .def_property_readonly("origin_generation", &Connection::originGeneration) + .def_property_readonly("origin_pool_id", &Connection::originPoolId); + struct PyObjectHolder { + PyObject* ptr = nullptr; + explicit PyObjectHolder(py::object obj) : ptr(obj.release().ptr()) {} + ~PyObjectHolder() { + if (ptr) { + py::gil_scoped_acquire gil; + Py_XDECREF(ptr); + ptr = nullptr; + } + } + PyObjectHolder(const PyObjectHolder&) = delete; + PyObjectHolder& operator=(const PyObjectHolder&) = delete; + PyObjectHolder(PyObjectHolder&& other) noexcept : ptr(other.ptr) { + other.ptr = nullptr; + } + PyObjectHolder& operator=(PyObjectHolder&& other) noexcept { + if (this != &other) { + if (ptr) { + py::gil_scoped_acquire gil; + Py_XDECREF(ptr); + } + ptr = other.ptr; + other.ptr = nullptr; + } + return *this; + } + }; + py::class_>(m, "_TestConnectionPool") + .def(py::init(), py::arg("max_size") = 1, py::arg("idle_timeout_secs") = 600) + .def( + "acquire", + [](ConnectionPool& pool, const std::u16string& connStr, + const py::object& token_factory) { + return pool.acquire(connStr, py::dict(), token_factory); + }, + py::arg("conn_str"), py::arg("token_factory") = py::none()) + .def("release", &ConnectionPool::release, py::call_guard(), py::arg("conn")) + .def("close", &ConnectionPool::close, py::call_guard()) + .def("set_mock_mode", &ConnectionPool::set_mock_mode, py::arg("enable") = true) + .def( + "set_on_disconnect_hook", + [](ConnectionPool& pool, py::object hook) { + if (hook.is_none()) { + pool.set_on_disconnect_hook(nullptr); + } else { + auto holder = std::make_shared(std::move(hook)); + auto fn = std::make_shared>([holder]() { + py::gil_scoped_acquire gil; + if (holder && holder->ptr) { + py::handle h(holder->ptr); + h(); + } + }); + pool.set_on_disconnect_hook(fn); + } + }, + py::arg("hook")) + .def_property_readonly("current_size", &ConnectionPool::current_size) + .def_property_readonly("checked_out", &ConnectionPool::checked_out) + .def_property_readonly("in_flight", &ConnectionPool::in_flight) + .def_property_readonly("generation", &ConnectionPool::generation) + .def_property_readonly("pool_id", &ConnectionPool::pool_id) + .def( + "inject_candidate", + [](ConnectionPool& pool, const std::u16string& connStr, long long expiry, bool mock) { + auto conn = std::make_shared(connStr, true); + if (mock || pool.mock_mode()) { + conn->setMock(true); + } + if (expiry > 0) { + conn->setTokenExpiry(expiry); + } + conn->updateLastUsed(); + conn->setPoolOrigin(pool.pool_id(), pool.generation()); + pool.inject_candidate(conn); + }, + py::arg("conn_str"), py::arg("expiry") = 0, py::arg("mock") = false); m.def("DDBCSQLExecDirect", &SQLExecDirect_wrap, "Execute a SQL query directly"); m.def("DDBCSQLExecute", &SQLExecute_wrap, "DetectParamTypes + BindParameters + SQLExecute all in C++", diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index e33f71030..0225411eb 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -39,7 +39,10 @@ def _run_in_subprocess(body: str, conn_str: str) -> None: is fine). """ env = os.environ.copy() - env["DB_CONNECTION_STRING"] = conn_str + if conn_str: + env["DB_CONNECTION_STRING"] = conn_str + else: + env.pop("DB_CONNECTION_STRING", None) proc = subprocess.run( [sys.executable, "-c", textwrap.dedent(body)], env=env, @@ -981,6 +984,1056 @@ def test_pooling_state_consistency(conn_str): print("Pooling state consistency verified") +def test_pool_size_accounting_race_on_close_interleave(conn_str): + """Regression test for GH-746: connection pool size accounting drift on close race. + + When a connection-open failure races pool close(), in-flight reservations + and checked-out connections are retained in _current_size across close(). + Coupled with pool generation tracking, Thread A's cleanup on failure only + decrements capacity if its generation still matches, preventing drift and + ensuring max_size is never exceeded across close() interleavings. + + Uses _TestConnectionPool to ensure Thread A, Thread B, and Thread C all + operate deterministically against the exact same pool instance. + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(2, 600) + assert pool.current_size == 0 + assert pool.generation == 0 + + in_factory_a = threading.Event() + release_factory_a = threading.Event() + + def factory_a(): + in_factory_a.set() + assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release factory A" + raise RuntimeError("simulated open failure A") + + t_a_error = [] + + def run_a(): + try: + pool.acquire("SERVER=dummy_test_746;", factory_a) + except Exception as exc: + t_a_error.append(exc) + + t_a = threading.Thread(target=run_a) + t_a.start() + assert in_factory_a.wait(timeout=5.0), "Timed out waiting for Thread A to enter factory" + assert pool.current_size == 1 + assert pool.in_flight == 1 + + # Thread A has reserved slot 1. Pool is closed while Thread A is in-flight. + # Reserved capacity is retained for in-flight opens so max_size is never exceeded (#746). + pool.close() + assert pool.current_size == 1 + assert pool.in_flight == 1 + assert pool.generation == 1 + + # Thread B initiates acquire and reserves the 2nd slot under the new generation. + in_factory_b = threading.Event() + release_factory_b = threading.Event() + + def factory_b(): + in_factory_b.set() + assert release_factory_b.wait(timeout=5.0), "Timed out waiting to release factory B" + raise RuntimeError("simulated open failure B") + + t_b_error = [] + + def run_b(): + try: + pool.acquire("SERVER=dummy_test_746;", factory_b) + except Exception as exc: + t_b_error.append(exc) + + t_b = threading.Thread(target=run_b) + t_b.start() + assert in_factory_b.wait(timeout=5.0), "Timed out waiting for Thread B to enter factory" + assert pool.current_size == 2 + assert pool.in_flight == 2 + + # Thread A now raises its error. Its cleanup decrements in_flight and current_size + # for Thread A (2 -> 1). Thread B's reservation under generation 1 is preserved. + release_factory_a.set() + t_a.join(timeout=5.0) + assert len(t_a_error) == 1 and "simulated open failure A" in str(t_a_error[0]) + + # Under the generation fix, current_size MUST still be 1 (Thread B's reservation is preserved). + assert pool.current_size == 1, ( + f"Expected pool.current_size to be 1, but got {pool.current_size} (drift occurred!)" + ) + assert pool.in_flight == 1 + + # Clean up Thread B + release_factory_b.set() + t_b.join(timeout=5.0) + assert len(t_b_error) == 1 and "simulated open failure B" in str(t_b_error[0]) + assert pool.current_size == 0 + assert pool.in_flight == 0 + pool.close() + """, + conn_str, + ) + + +def test_pool_size_accounting_race_on_candidate_validation_close_interleave(conn_str): + """Regression test for GH-746: candidate validation failure racing pool close(). + + When a candidate popped from the pool fails validation (e.g. dead socket or + token rotation failure) while racing a pool close(), the pool retains + capacity for the in-flight validation across close(). Coupled with generation + guarding, Thread A's cleanup on validation failure decrements only its own + reservation without corrupting any newer generation's reservation. + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(2, 600) + # Inject an expired candidate into the idle pool + pool.inject_candidate("SERVER=dummy_test_746;", 1) + assert pool.current_size == 1 + assert pool.generation == 0 + + in_factory_a = threading.Event() + release_factory_a = threading.Event() + + def factory_a(): + in_factory_a.set() + assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release factory A" + raise RuntimeError("simulated token rotation validation failure") + + t_a_error = [] + + def run_a(): + try: + pool.acquire("SERVER=dummy_test_746;", factory_a) + except Exception as exc: + t_a_error.append(exc) + + t_a = threading.Thread(target=run_a) + t_a.start() + assert in_factory_a.wait(timeout=5.0), "Timed out waiting for Thread A to enter factory" + assert pool.in_flight == 1 + + # Thread A popped the candidate (generation 0) and is validating it in factory_a. + # Pool close retains capacity for in-flight validation: current_size stays 1. + pool.close() + assert pool.current_size == 1 + assert pool.in_flight == 1 + assert pool.generation == 1 + + # Thread B reserves the 2nd slot under generation 1. + in_factory_b = threading.Event() + release_factory_b = threading.Event() + + def factory_b(): + in_factory_b.set() + assert release_factory_b.wait(timeout=5.0), "Timed out waiting to release factory B" + raise RuntimeError("simulated open failure B") + + t_b_error = [] + + def run_b(): + try: + pool.acquire("SERVER=dummy_test_746;", factory_b) + except Exception as exc: + t_b_error.append(exc) + + t_b = threading.Thread(target=run_b) + t_b.start() + assert in_factory_b.wait(timeout=5.0), "Timed out waiting for Thread B to enter factory" + assert pool.current_size == 2 + assert pool.in_flight == 2 + + # Thread A finishes factory_a (validation fails). + # Thread A releases its in-flight reservation (2 -> 1). + # Thread B's reservation under generation 1 is preserved. + release_factory_a.set() + t_a.join(timeout=5.0) + + # Verify Thread B's reservation was not cancelled + assert pool.current_size == 1, ( + f"Expected pool.current_size to be 1, but got {pool.current_size} (drift occurred!)" + ) + assert pool.in_flight == 1 + + release_factory_b.set() + t_b.join(timeout=5.0) + assert pool.current_size == 0 + assert pool.in_flight == 0 + pool.close() + """, + conn_str, + ) + + +def test_pool_size_accounting_race_on_successful_candidate_reuse_close_interleave(conn_str): + """Regression test for GH-746: candidate reuse success racing pool close(). + + When a candidate popped under generation 0 succeeds validation while racing + a pool close(), the candidate must NOT be returned as a valid connection + under the new generation. Returning it without an active reservation in the + new generation would allow another thread to reserve up to max_size, causing + the pool to exceed max_size. Instead, the stale candidate is discarded and + acquire retries under the new generation (or fails if the pool is full). + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(2, 600) + pool.set_mock_mode(True) + # Inject candidate with near-expiry so factory_a runs to check it + pool.inject_candidate("SERVER=dummy_test_746;", 1, True) + assert pool.current_size == 1 + assert pool.generation == 0 + + in_factory_a = threading.Event() + release_factory_a = threading.Event() + factory_a_calls = [0] + + def factory_a(): + factory_a_calls[0] += 1 + if factory_a_calls[0] == 1: + in_factory_a.set() + assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release factory A" + return {}, 9999999999 + + t_a_conn = [] + t_a_error = [] + + def run_a(): + try: + conn = pool.acquire("SERVER=dummy_test_746;", factory_a) + t_a_conn.append(conn) + except Exception as exc: + t_a_error.append(exc) + + t_a = threading.Thread(target=run_a) + t_a.start() + assert in_factory_a.wait(timeout=5.0), "Timed out waiting for Thread A to enter factory" + assert pool.in_flight == 1 + + # Thread A popped candidate (generation 0). Pool close retains in-flight validation. + pool.close() + assert pool.current_size == 1 + assert pool.in_flight == 1 + assert pool.generation == 1 + + # Thread B reserves slot 2 under generation 1 + in_factory_b = threading.Event() + release_factory_b = threading.Event() + + def factory_b(): + in_factory_b.set() + assert release_factory_b.wait(timeout=5.0), "Timed out waiting to release factory B" + return {} + + t_b_conn = [] + t_b_error = [] + + def run_b(): + try: + conn = pool.acquire("SERVER=dummy_test_746;", factory_b) + t_b_conn.append(conn) + except Exception as exc: + t_b_error.append(exc) + + t_b = threading.Thread(target=run_b) + t_b.start() + assert in_factory_b.wait(timeout=5.0), "Timed out waiting for Thread B to enter factory" + assert pool.current_size == 2 + assert pool.in_flight == 2 + + # 1. Thread B finishes connecting first and publishes under generation 1 + release_factory_b.set() + t_b.join(timeout=5.0) + assert len(t_b_error) == 0 + assert len(t_b_conn) == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 1 + assert pool.current_size == 2 + + # 2. Thread A finishes validation second. Under the generation fix, Thread A detects + # generation mismatch (0 != 1), discards stale candidate, decrements in_flight and current_size (2 -> 1). + # Thread A retries acquire under generation 1 and successfully acquires slot 2. + release_factory_a.set() + t_a.join(timeout=5.0) + assert len(t_a_error) == 0 + assert len(t_a_conn) == 1 + # Verify candidate was discarded and A performed a fresh open under generation 1 (#746) + assert factory_a_calls[0] == 2, f"Expected 2 factory invocations (validation + retry open), got {factory_a_calls[0]}" + assert t_a_conn[0].origin_generation == 1, f"Expected generation 1, got {t_a_conn[0].origin_generation}" + assert pool.checked_out == 2 + assert pool.in_flight == 0 + assert pool.current_size == 2 + + # Return both connections and clean up + pool.release(t_a_conn[0]) + pool.release(t_b_conn[0]) + assert pool.checked_out == 0 + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + +def test_pool_size_accounting_race_on_successful_open_close_interleave(conn_str): + """Regression test for GH-746: connection open success racing pool close(). + + When a connection open succeeds after racing a pool close(), the newly opened + connection must NOT be returned into the pool under the new generation without + validating the generation counter. If Thread A connected under generation 0, + close() reset the pool, and Thread B reserved generation 1, returning Thread A's + connection would result in 2 live connections when max_size=1. The generation + check ensures Thread A discards the orphaned connection and retries under the + new generation (failing if Thread B has claimed the capacity). + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(2, 600) + pool.set_mock_mode(True) + assert pool.current_size == 0 + assert pool.generation == 0 + + in_factory_a = threading.Event() + release_factory_a = threading.Event() + factory_a_calls = [0] + + def factory_a(): + factory_a_calls[0] += 1 + if factory_a_calls[0] == 1: + in_factory_a.set() + assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release factory A" + return {} + + t_a_conn = [] + t_a_error = [] + + def run_a(): + try: + conn = pool.acquire("SERVER=dummy_test_746;", factory_a) + t_a_conn.append(conn) + except Exception as exc: + t_a_error.append(exc) + + t_a = threading.Thread(target=run_a) + t_a.start() + assert in_factory_a.wait(timeout=5.0), "Timed out waiting for Thread A to enter factory" + assert pool.current_size == 1 + assert pool.in_flight == 1 + + # Thread A reserved slot under generation 0. Pool is closed while Thread A is in-flight. + # In-flight capacity is retained so max_size is not exceeded (#746). + pool.close() + assert pool.current_size == 1 + assert pool.in_flight == 1 + assert pool.generation == 1 + + # Thread B begins acquiring under generation 1 (reserves slot 2 of max_size=2) + in_factory_b = threading.Event() + release_factory_b = threading.Event() + + def factory_b(): + in_factory_b.set() + assert release_factory_b.wait(timeout=5.0), "Timed out waiting to release factory B" + return {} + + t_b_conn = [] + t_b_error = [] + + def run_b(): + try: + conn = pool.acquire("SERVER=dummy_test_746;", factory_b) + t_b_conn.append(conn) + except Exception as exc: + t_b_error.append(exc) + + t_b = threading.Thread(target=run_b) + t_b.start() + assert in_factory_b.wait(timeout=5.0), "Timed out waiting for Thread B to enter factory" + assert pool.current_size == 2 + assert pool.in_flight == 2 + + # 1. NEW-GENERATION OPEN COMPLETES FIRST: + # Thread B finishes connecting and publishes valid_conn under generation 1 + release_factory_b.set() + t_b.join(timeout=5.0) + assert len(t_b_error) == 0 + assert len(t_b_conn) == 1 + assert t_b_conn[0].origin_generation == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 1 + assert pool.current_size == 2 + + # 2. OLD-GENERATION OPEN COMPLETES SECOND: + # Thread A finishes connecting outside the lock. + # Under generation fix, Thread A detects reservation generation mismatch (0 != 1), + # disconnects its stale connection, decrements in_flight (1 -> 0) and current_size (2 -> 1). + # Thread B's connection is intact and valid! + # Thread A retries acquire under generation 1 and successfully acquires the freed slot. + release_factory_a.set() + t_a.join(timeout=5.0) + assert len(t_a_error) == 0 + assert len(t_a_conn) == 1 + # Verify stale open was discarded and A performed a fresh open under generation 1 (#746) + assert factory_a_calls[0] == 2, f"Expected 2 factory invocations (initial open + retry open), got {factory_a_calls[0]}" + assert t_a_conn[0].origin_generation == 1, f"Expected generation 1, got {t_a_conn[0].origin_generation}" + assert pool.checked_out == 2 + assert pool.in_flight == 0 + assert pool.current_size == 2 + + # Clean up both connections + pool.release(t_a_conn[0]) + pool.release(t_b_conn[0]) + assert pool.checked_out == 0 + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + +def test_pool_in_flight_open_blocks_acquire_exceeding_max_size_1(conn_str): + """When max_size=1 and an open is in flight across close(), new acquire is blocked. + + Ensures that an in-flight open retains reserved capacity across close(), + preventing a new-generation thread from opening another physical connection + while the stale open is still establishing its socket (#746). + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(1, 600) + pool.set_mock_mode(True) + + in_factory_a = threading.Event() + release_factory_a = threading.Event() + + def factory_a(): + in_factory_a.set() + assert release_factory_a.wait(timeout=5.0) + return {} + + t_a_conn = [] + + def run_a(): + conn = pool.acquire("SERVER=dummy_test_746;", factory_a) + t_a_conn.append(conn) + + t_a = threading.Thread(target=run_a) + t_a.start() + assert in_factory_a.wait(timeout=5.0) + assert pool.current_size == 1 + assert pool.in_flight == 1 + + # Close pool while Thread A is in-flight + pool.close() + # In-flight capacity is retained: current_size stays 1! + assert pool.current_size == 1 + assert pool.in_flight == 1 + assert pool.generation == 1 + + # Thread B tries to acquire under generation 1: REJECTED because capacity is held! + rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", lambda: {}) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "Thread B must be rejected while Thread A's open is still in-flight" + + # Thread A completes and disconnects stale connection, freeing the slot + release_factory_a.set() + t_a.join(timeout=5.0) + + # Thread A's retry acquired the freed slot under generation 1 + assert len(t_a_conn) == 1 + assert pool.current_size == 1 + assert pool.checked_out == 1 + + pool.release(t_a_conn[0]) + assert pool.checked_out == 0 + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + +def test_pool_release_from_stale_generation_does_not_pollute_pool(conn_str): + """Regression test for GH-746: releasing a connection from an invalidated pool/generation. + + When a connection checked out from an earlier pool generation is released after + pool.close() has advanced the generation, release() must NOT push that stale connection + back into the active pool nor decrement current_size of the new generation. Instead, it + must cleanly disconnect the stale connection, ensuring the new pool generation remains + uncorrupted and never exceeds max_size. + """ + _run_in_subprocess( + """ + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(1, 600) + pool.set_mock_mode(True) + + # 1. Acquire conn_1 under generation 0 + conn_1 = pool.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_1 is not None + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.generation == 0 + + # 2. Pool is closed while conn_1 is still checked out. + # Reserved capacity is retained for checked-out connections so the + # max_size cap is not exceeded while conn_1 is live (#746). + pool.close() + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.generation == 1 + + # 3. An acquire under generation 1 must be rejected while conn_1 is still checked out, + # preserving the max_size=1 invariant across pool.close() (#746). + rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", lambda: {}) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "A new acquire must be rejected while conn_1 is still checked out" + + # 4. Release conn_1 (from generation 0). + # It belongs to this pool but its generation is stale. It is disconnected, + # and the retained checked-out capacity is released: current_size drops to 0. + pool.release(conn_1) + assert pool.current_size == 0 + assert pool.checked_out == 0 + + # 5. Now that capacity has freed up, acquire conn_2 under generation 1 succeeds + conn_2 = pool.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_2 is not None + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.generation == 1 + + # 6. Release conn_2 (matches generation 1). It returns to the pool idle deque. + pool.release(conn_2) + assert pool.current_size == 1 + assert pool.checked_out == 0 + + # 7. Next acquire reuses conn_2 from the pool + conn_3 = pool.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_3 is not None + assert pool.current_size == 1 + assert pool.checked_out == 1 + + pool.release(conn_3) + assert pool.checked_out == 0 + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + +def test_pool_release_after_pool_recreation(conn_str): + """Releasing a connection to a newly recreated pool must not corrupt the new pool's size. + + Verifies that monotonic pool IDs prevent address-reuse (ABA) corruption: + even if pool_2 were to be allocated at the same memory address as pool_1, + releasing conn_1 (from pool_1) to pool_2 will not match pool_2's monotonic pool ID. + Therefore, pool_2's current_size is not erroneously decremented or corrupted (#746). + """ + _run_in_subprocess( + """ + from mssql_python import ddbc_bindings + + pool_1 = ddbc_bindings._TestConnectionPool(1, 600) + pool_1.set_mock_mode(True) + assert pool_1.pool_id > 0 + + # Acquire conn_1 from pool_1 + conn_1 = pool_1.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_1 is not None + assert pool_1.current_size == 1 + assert pool_1.checked_out == 1 + + # Create pool_2 with its own distinct monotonic pool_id + pool_2 = ddbc_bindings._TestConnectionPool(1, 600) + pool_2.set_mock_mode(True) + assert pool_2.pool_id > pool_1.pool_id + assert pool_2.current_size == 0 + assert pool_2.checked_out == 0 + + # Release conn_1 into pool_2 (wrong pool ID) + pool_2.release(conn_1) + # pool_2 must not adopt conn_1 or decrement its size: stays 0 + assert pool_2.current_size == 0 + assert pool_2.checked_out == 0 + + # pool_2 can acquire normally + conn_2 = pool_2.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_2 is not None + assert pool_2.current_size == 1 + assert pool_2.checked_out == 1 + + # Releasing conn_1 again to pool_2 does not corrupt pool_2's active connection + pool_2.release(conn_1) + assert pool_2.current_size == 1 + assert pool_2.checked_out == 1 + + # Cleanly release conn_2 + pool_2.release(conn_2) + assert pool_2.checked_out == 0 + pool_2.close() + assert pool_2.current_size == 0 + """, + conn_str, + ) + + +def test_pool_manager_defers_replacement_while_connection_checked_out(conn_str): + """Across a disable/enable cycle, acquires respect capacity of checked-out connections. + + When ConnectionPoolManager::closePools() runs while a connection is checked out, + the pool is retained in _pools, deferring replacement pool creation until the old pool + has no live work. Subsequent acquires under the re-enabled pool manager must not + allow exceeding max_size while the old connection remains checked out (#746). + """ + _run_in_subprocess( + """ + from mssql_python import ddbc_bindings + + ddbc_bindings._set_pool_manager_mock_mode(True) + + pool_key = "SERVER=dummy_test_746;test_replace" + conn_str = "SERVER=dummy_test_746;" + + # 1. Enable pooling with max_size=2 + ddbc_bindings.enable_pooling(2, 600) + + # 2. Acquire conn_1 from pool (generation 0) + conn_1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + assert conn_1.origin_generation == 0 + + # 3. Disable pooling (invokes closePools()) + # The pool has checked_out=1, so it is retained in _pools with generation bumped to 1. + ddbc_bindings.disable_pooling() + + # 4. Re-enable pooling with max_size=2 + ddbc_bindings.enable_pooling(2, 600) + + # 5. Acquire conn_2 from the manager while conn_1 is still checked out. + # Replacement pool creation is deferred because conn_1 is still live. + # conn_2 is acquired under generation 1 (1 + 1 = 2 connections active). + conn_2 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + assert conn_2.origin_generation == 1 + assert conn_2.origin_pool_id == conn_1.origin_pool_id + + # 6. Attempting to acquire a 3rd connection must fail because max_size=2 is reached! + rejected = False + try: + ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "Must reject acquire when max_size=2 is reached across recreation!" + + # 7. Close conn_1: releases the old-generation connection and frees capacity. + conn_1.close() + + # 8. Now acquire conn_3: capacity is freed, so acquire succeeds! + conn_3 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + assert conn_3.origin_generation == 1 + + # Clean up remaining connections + conn_2.close() + conn_3.close() + ddbc_bindings.close_pooling() + """, + conn_str, + ) + + +def test_pool_release_disconnect_keeps_in_flight_until_disconnected(conn_str): + """Regression test for GH-746: stale/overflow release disconnect retains in-flight accounting. + + When an expired, generation-mismatched, or overflow connection is released, + it must be transitioned from _checked_out to _in_flight during the + unlocked conn->disconnect() call, and only decremented after disconnect finishes. + This guarantees that concurrent callers cannot reserve a slot or open a new + physical connection before the old physical connection has finished closing. + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(1, 600) + pool.set_mock_mode(True) + + # 1. Acquire connection (generation 0, checked_out=1, current_size=1) + conn = pool.acquire("SERVER=dummy_test_746;", None) + assert conn is not None + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 0 + + # 2. Close the pool to bump the generation so releasing conn triggers a disconnect. + # Since conn was checked out, close() keeps current_size=1, checked_out=1, in_flight=0. + pool.close() + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 0 + assert pool.generation == 1 + + in_disconnect = threading.Event() + release_disconnect = threading.Event() + hook_observed = {} + + def on_disconnect(): + # At this point, release() has moved conn from _checked_out to _in_flight, + # but has NOT yet decremented current_size. + hook_observed["current_size"] = pool.current_size + hook_observed["in_flight"] = pool.in_flight + hook_observed["checked_out"] = pool.checked_out + in_disconnect.set() + assert release_disconnect.wait(timeout=5.0), "Timed out waiting to release disconnect" + + pool.set_on_disconnect_hook(on_disconnect) + + t_err = [] + def run_release(): + try: + pool.release(conn) + except Exception as exc: + t_err.append(exc) + + t = threading.Thread(target=run_release) + t.start() + assert in_disconnect.wait(timeout=5.0), "Timed out waiting for disconnect hook" + + # Verify hook observed state: + assert hook_observed["current_size"] == 1 + assert hook_observed["in_flight"] == 1 + assert hook_observed["checked_out"] == 0 + + # While disconnect is in progress, any concurrent acquire is blocked from + # allocating because max_size=1 is still fully occupied by in_flight teardown! + rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", None) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "Concurrent acquire must be rejected while teardown disconnect is in-flight!" + + # Let the disconnect complete + release_disconnect.set() + t.join(timeout=5.0) + assert not t_err, f"Release thread error: {t_err}" + + # Now that disconnect is finished, counters are decremented to 0 + assert pool.current_size == 0 + assert pool.in_flight == 0 + assert pool.checked_out == 0 + + # Now acquire succeeds + pool.set_on_disconnect_hook(None) + conn_new = pool.acquire("SERVER=dummy_test_746;", None) + assert conn_new is not None + assert pool.current_size == 1 + assert pool.checked_out == 1 + pool.release(conn_new) + pool.close() + """, + conn_str, + ) + + +def test_pool_close_disconnect_keeps_in_flight_until_disconnected(conn_str): + """Regression test for GH-746: idle connection disconnect in close() retains in-flight accounting. + + When close() drains idle connections from the pool, they must be added to + _in_flight and only decremented from _in_flight and _current_size after each + physical disconnect finishes, preventing concurrent acquires from observing freed + slots while physical sockets are still closing. + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(1, 600) + pool.set_mock_mode(True) + + # Inject an idle candidate into the pool + pool.inject_candidate("SERVER=dummy_test_746;", 0, True) + assert pool.current_size == 1 + assert pool.checked_out == 0 + assert pool.in_flight == 0 + + in_disconnect = threading.Event() + release_disconnect = threading.Event() + hook_observed = {} + + def on_disconnect(): + hook_observed["current_size"] = pool.current_size + hook_observed["in_flight"] = pool.in_flight + hook_observed["checked_out"] = pool.checked_out + in_disconnect.set() + assert release_disconnect.wait(timeout=5.0), "Timed out waiting to release disconnect" + + pool.set_on_disconnect_hook(on_disconnect) + + t_err = [] + def run_close(): + try: + pool.close() + except Exception as exc: + t_err.append(exc) + + t = threading.Thread(target=run_close) + t.start() + assert in_disconnect.wait(timeout=5.0), "Timed out waiting for disconnect hook in close" + + # Verify hook observed state: + assert hook_observed["current_size"] == 1 + assert hook_observed["in_flight"] == 1 + assert hook_observed["checked_out"] == 0 + + # Concurrent acquire cannot exceed max_size while idle connection is disconnecting + rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", None) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "Concurrent acquire must be rejected while close disconnect is in-flight!" + + # Let close complete + release_disconnect.set() + t.join(timeout=5.0) + assert not t_err, f"Close thread error: {t_err}" + + # Verify clean post-close state + assert pool.current_size == 0 + assert pool.in_flight == 0 + assert pool.checked_out == 0 + """, + conn_str, + ) + + +def test_pool_prune_stale_disconnect_keeps_in_flight_until_disconnected(conn_str): + """Regression test for GH-746: Phase 1 stale idle pruning retains in-flight capacity. + + When Phase 1 prunes stale idle connections past idle_timeout, they are moved to + _in_flight and only decremented after physical disconnect in Phase 4 finishes, + preventing concurrent callers from allocating into slots of disconnecting handles. + """ + _run_in_subprocess( + """ + import time + import threading + from mssql_python import ddbc_bindings + + # Pool with idle timeout of 0 seconds and max_size=1 + pool = ddbc_bindings._TestConnectionPool(1, 0) + pool.set_mock_mode(True) + + pool.inject_candidate("SERVER=dummy_test_746;", 0, True) + assert pool.current_size == 1 + assert pool.checked_out == 0 + assert pool.in_flight == 0 + + # Wait for candidate to exceed idle timeout + time.sleep(1.1) + + in_disconnect = threading.Event() + release_disconnect = threading.Event() + hook_observed = {} + + def on_disconnect(): + hook_observed["current_size"] = pool.current_size + hook_observed["in_flight"] = pool.in_flight + hook_observed["checked_out"] = pool.checked_out + in_disconnect.set() + assert release_disconnect.wait(timeout=5.0), "Timed out waiting to release disconnect" + + pool.set_on_disconnect_hook(on_disconnect) + + t_err = [] + t_conn = [] + + def run_acquire(): + try: + conn = pool.acquire("SERVER=dummy_test_746;", None) + t_conn.append(conn) + except Exception as exc: + t_err.append(exc) + + t = threading.Thread(target=run_acquire) + t.start() + assert in_disconnect.wait(timeout=5.0), "Timed out waiting for disconnect hook in Phase 4" + + # Verify hook observed state: + assert hook_observed["current_size"] == 1 + assert hook_observed["in_flight"] == 1 + assert hook_observed["checked_out"] == 0 + + # Concurrent acquire cannot exceed max_size while pruned connection is disconnecting + rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", None) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + rejected = True + assert rejected, "Concurrent acquire must be rejected while pruned disconnect is in-flight!" + + # Let disconnect finish + release_disconnect.set() + t.join(timeout=5.0) + assert not t_err, f"Acquire thread error: {t_err}" + assert len(t_conn) == 1 + + # Now newly acquired connection is checked out + assert pool.current_size == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 0 + + pool.set_on_disconnect_hook(None) + pool.release(t_conn[0]) + pool.close() + """, + conn_str, + ) + + +def test_pool_manager_serializes_same_key_replacement_while_old_pool_closing(conn_str): + """Regression test for GH-746: serialize replacement-pool creation while old pool is closing. + + When ConnectionPoolManager evicts an evictable pool and begins closing it, concurrent + acquires for the same key must wait for teardown to finish rather than creating and + publishing a competing pool before old idle handles are disconnected. + """ + _run_in_subprocess( + """ + import threading + import time + from mssql_python import ddbc_bindings + + ddbc_bindings._set_pool_manager_mock_mode(True) + + pool_key = "SERVER=dummy_test_746;test_replace_closing" + conn_str = "SERVER=dummy_test_746;" + + # 1. Enable pooling with max_size=1, idle_timeout=0 so idle pools become evictable + ddbc_bindings.enable_pooling(1, 0) + + # 2. Acquire a connection from the manager and return it to make the pool idle + conn_init = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + conn_init.close() + + # Sleep briefly so idle_time > 0 (idle_timeout=0) makes canEvict() return True + time.sleep(1.05) + + # 3. Retrieve the internal pool instance for pool_key and attach a disconnect hook + old_pool = ddbc_bindings._get_pool_for_key(pool_key) + assert old_pool is not None + + hook_entered = threading.Event() + proceed_disconnect = threading.Event() + thread_b_started = threading.Event() + thread_b_finished = threading.Event() + thread_b_result = [] + + def on_disconnect(): + hook_entered.set() + # Wait until Thread B has launched its acquire attempt + thread_b_started.wait(timeout=5.0) + # Sleep a moment to ensure Thread B enters acquireConnection and waits on _manager_cv + time.sleep(0.15) + assert not thread_b_finished.is_set(), "Thread B must be blocked waiting on _closing_keys!" + proceed_disconnect.wait(timeout=5.0) + + old_pool.set_on_disconnect_hook(on_disconnect) + # Drop Python reference so it->second.use_count() == 1 allows eviction + del old_pool + + # 4. Thread A triggers acquireConnection, detecting the evictable pool and calling close() + thread_a_result = [] + def thread_a_worker(): + try: + c = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + thread_a_result.append(c) + except Exception as e: + thread_a_result.append(e) + + def thread_b_worker(): + thread_b_started.set() + try: + c = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + thread_b_result.append(c) + except Exception as e: + thread_b_result.append(e) + finally: + thread_b_finished.set() + + t_a = threading.Thread(target=thread_a_worker) + t_a.start() + + assert hook_entered.wait(timeout=5.0), "Disconnect hook was not reached" + + # 5. Launch Thread B: tries to acquire for the same key while Thread A is closing old pool + t_b = threading.Thread(target=thread_b_worker) + t_b.start() + + # Let Thread A complete the disconnect and pool replacement + proceed_disconnect.set() + + t_a.join(timeout=5.0) + t_b.join(timeout=5.0) + + assert len(thread_a_result) == 1 and not isinstance(thread_a_result[0], Exception) + conn_a = thread_a_result[0] + + # Thread B must have been serialized and checked out from the replacement pool. + # Since Thread A took the only slot on the replacement pool (max_size=1), + # Thread B was rejected with 'pool size limit reached'. + assert len(thread_b_result) == 1 + res_b = thread_b_result[0] + assert isinstance(res_b, RuntimeError) and "pool size limit reached" in str(res_b), ( + f"Expected pool size limit reached on serialized replacement pool, got: {res_b}" + ) + + # Free slot on replacement pool and verify subsequent acquire succeeds + expected_pool_id = conn_a.origin_pool_id + conn_a.close() + conn_after = ddbc_bindings.Connection(conn_str, True, {}, pool_key, lambda: {}) + assert conn_after.origin_pool_id == expected_pool_id + conn_after.close() + + ddbc_bindings.close_pooling() + """, + conn_str, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # =============================================================================