From 5328d5d44f3a40171e4c532cd6a851d084ec7c76 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:12:06 -0400 Subject: [PATCH 01/12] FIX: Prevent connection pool size drift when close races failed open (#746) ### Work Item / Issue Reference > GitHub Issue: #746 ------------------------------------------------------------------- ### Summary This pull request fixes a race condition in `ConnectionPool` where an interleaved pool `close()` (or `close_pooling()`) during a failing connection open could corrupt `_current_size`, causing the pool to undercount active connections and briefly exceed `max_size` under concurrent load (#746). **Root Cause:** When a thread fails to open/construct a connection in Phase 3 of `ConnectionPool::acquire()`, its `catch (...)` handler decrements `_current_size` to release the reserved slot. If `ConnectionPool::close()` executed concurrently outside the lock, `close()` already cleared all idle connections and reset `_current_size = 0`. If another thread then acquired and reserved a slot (`_current_size = 1`), the first thread's subsequent decrement wiped out the new thread's reservation instead of its own, causing the counter to drift lower than the true active connection count. **Key Changes:** * Added `uint64_t _generation` counter to `ConnectionPool`, incremented on every `ConnectionPool::close()`. * In `ConnectionPool::acquire()`, captured `reservation_generation = _generation` alongside slot reservations (`++_current_size`). * In Phase 3 error recovery, guarded slot decrements with `if (_generation == reservation_generation && _current_size > 0)` so that stale reservation cleanups from an earlier generation do not cancel newer reservations. * Added regression test `test_pool_size_accounting_race_on_close_interleave` in `tests/test_009_pooling.py` reproducing the exact interleaving and verifying slot bounds are preserved. --- .../pybind/connection/connection_pool.cpp | 14 ++- .../pybind/connection/connection_pool.h | 1 + tests/test_009_pooling.py | 100 +++++++++++++++++- 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 831a01db2..7cdaca0f5 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -102,6 +102,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt py::dict pending_attrs; long long pending_expiry = 0; bool have_pending_token = false; + uint64_t reservation_generation = 0; while (true) { std::shared_ptr candidate; { @@ -115,6 +116,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // holding _mutex across a GIL acquisition deadlocks a thread // that holds the GIL and is waiting on _mutex (#671). ++_current_size; + reservation_generation = _generation; needs_connect = true; break; } @@ -243,6 +245,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // records, and holding _mutex across a GIL acquisition // deadlocks a thread that holds the GIL and waits on _mutex (#671). ++_current_size; + reservation_generation = _generation; needs_connect = true; break; } @@ -283,10 +286,16 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt valid_conn->connect(attrs_before); } } catch (...) { - // Construct/connect failed — release the reserved slot + // Construct/connect failed — release the reserved slot only if the pool + // has not been reset in the meantime (#746). If close() ran while we were + // connecting outside the lock, close() already set _current_size = 0 and + // bumped _generation; decrementing here would cancel another thread's + // newer reservation instead of our own. { std::lock_guard lock(_mutex); - if (_current_size > 0) --_current_size; + if (_generation == reservation_generation && _current_size > 0) { + --_current_size; + } } throw; } @@ -370,6 +379,7 @@ void ConnectionPool::close() { _pool.pop_front(); } _current_size = 0; + ++_generation; } for (auto& conn : to_close) { try { diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index edc87c865..55ad34669 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -45,6 +45,7 @@ class ConnectionPool { size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale size_t _current_size = 0; + uint64_t _generation = 0; // Pool reset generation for reservation attribution (#746) std::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access }; diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index e33f71030..168f0631b 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,101 @@ 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 close_pooling() (or pool close), + the failed thread's decrement must not cancel a newer generation's + reservation. With max_size=1, if Thread A's open failure wrongly + decrements the counter after Thread B reserved the slot under the new + generation, Thread C would be allowed to connect, exceeding max_size. + The generation counter guarantees that Thread A's cleanup only decrements + if the pool generation still matches its reservation. + + Run in a subprocess so pooling(max_size=1) is the effective configuration. + """ + _run_in_subprocess( + """ + import os, threading + from mssql_python import ddbc_bindings + + ddbc_bindings.enable_pooling(1, 600) + base_conn = os.environ.get("DB_CONNECTION_STRING") or "SERVER=dummy_test_746;" + pool_key = base_conn + "\\x00mssql_test_746_race" + + 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: + ddbc_bindings.Connection(base_conn, True, {}, pool_key, 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" + + # Thread A has reserved the slot. Now close_pooling() resets the pool + # and increments the pool generation counter. + ddbc_bindings.close_pooling() + + # Thread B initiates acquire and reserves the freed 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: + ddbc_bindings.Connection(base_conn, True, {}, pool_key, 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" + + # Thread A now raises its error. With the generation counter fix, its cleanup + # detects that the pool generation changed and does NOT decrement _current_size. + 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]) + + # Thread C now attempts to acquire on the same pool key. + # Since max_size=1 and Thread B is still reserving the slot, Thread C must fail + # with 'pool size limit reached'. + thread_c_rejected = False + try: + ddbc_bindings.Connection(base_conn, True, {}, pool_key, lambda: {}) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + thread_c_rejected = True + + assert thread_c_rejected, "Thread C should have been rejected due to pool capacity limit" + + # 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]) + ddbc_bindings.close_pooling() + """, + conn_str, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From 8153fdc0bfdf435aeada3d2067e797b4ed414ab2 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:22:08 -0400 Subject: [PATCH 02/12] TEST: Add native _TestConnectionPool seam for deterministic race verification (#746) - Expose _TestConnectionPool in ddbc_bindings with current_size and generation accessors. - Update test_pool_size_accounting_race_on_close_interleave to share a single ConnectionPool instance across Thread A, Thread B, and Thread C, verifying that Thread A's failed open cannot cancel Thread B's reservation under a new generation. --- .../pybind/connection/connection_pool.h | 10 ++++ mssql_python/pybind/ddbc_bindings.cpp | 15 ++++++ tests/test_009_pooling.py | 54 +++++++++++-------- 3 files changed, 58 insertions(+), 21 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index 55ad34669..baf5c0abd 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -41,6 +41,16 @@ class ConnectionPool { // dropped by the manager to reclaim memory (lazy eviction). bool canEvict(); + // Test accessors for pool generation and current size + size_t current_size() const { + std::lock_guard lock(const_cast(_mutex)); + return _current_size; + } + uint64_t generation() const { + std::lock_guard lock(const_cast(_mutex)); + return _generation; + } + private: size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 35a3b6da4..0ed652176 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.) @@ -6121,6 +6122,20 @@ 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, "_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) { + pool.acquire(connStr, py::dict(), token_factory); + }, + py::arg("conn_str"), py::arg("token_factory") = py::none()) + .def("close", &ConnectionPool::close) + .def_property_readonly("current_size", &ConnectionPool::current_size) + .def_property_readonly("generation", &ConnectionPool::generation); 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 168f0631b..d255c785c 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -987,24 +987,25 @@ def test_pooling_state_consistency(conn_str): 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 close_pooling() (or pool close), - the failed thread's decrement must not cancel a newer generation's - reservation. With max_size=1, if Thread A's open failure wrongly - decrements the counter after Thread B reserved the slot under the new - generation, Thread C would be allowed to connect, exceeding max_size. - The generation counter guarantees that Thread A's cleanup only decrements - if the pool generation still matches its reservation. - - Run in a subprocess so pooling(max_size=1) is the effective configuration. + When a connection-open failure races pool close(), the failed thread's + decrement must not cancel a newer generation's reservation on that pool. + With max_size=1, if Thread A's open failure wrongly decrements the counter + after Thread B reserved the slot under the new generation, Thread C would + be allowed to connect, exceeding max_size. The generation counter guarantees + that Thread A's cleanup only decrements if the pool generation still matches + its reservation. + + Uses _TestConnectionPool to ensure Thread A, Thread B, and Thread C all + operate deterministically against the exact same pool instance. """ _run_in_subprocess( """ - import os, threading + import threading from mssql_python import ddbc_bindings - ddbc_bindings.enable_pooling(1, 600) - base_conn = os.environ.get("DB_CONNECTION_STRING") or "SERVER=dummy_test_746;" - pool_key = base_conn + "\\x00mssql_test_746_race" + pool = ddbc_bindings._TestConnectionPool(1, 600) + assert pool.current_size == 0 + assert pool.generation == 0 in_factory_a = threading.Event() release_factory_a = threading.Event() @@ -1018,17 +1019,20 @@ def factory_a(): def run_a(): try: - ddbc_bindings.Connection(base_conn, True, {}, pool_key, factory_a) + 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 - # Thread A has reserved the slot. Now close_pooling() resets the pool + # Thread A has reserved the slot. Now pool.close() resets the pool # and increments the pool generation counter. - ddbc_bindings.close_pooling() + pool.close() + assert pool.current_size == 0 + assert pool.generation == 1 # Thread B initiates acquire and reserves the freed slot under the new generation. in_factory_b = threading.Event() @@ -1043,26 +1047,33 @@ def factory_b(): def run_b(): try: - ddbc_bindings.Connection(base_conn, True, {}, pool_key, factory_b) + 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 == 1 # Thread A now raises its error. With the generation counter fix, its cleanup - # detects that the pool generation changed and does NOT decrement _current_size. + # detects that the pool generation changed (0 != 1) and does NOT decrement current_size. + # On buggy code without the fix, Thread A's cleanup would decrement current_size back to 0. 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]) - # Thread C now attempts to acquire on the same pool key. + # 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!)" + ) + + # Thread C now attempts to acquire on the same pool. # Since max_size=1 and Thread B is still reserving the slot, Thread C must fail # with 'pool size limit reached'. thread_c_rejected = False try: - ddbc_bindings.Connection(base_conn, True, {}, pool_key, lambda: {}) + pool.acquire("SERVER=dummy_test_746;", lambda: {}) except RuntimeError as exc: if "pool size limit reached" in str(exc): thread_c_rejected = True @@ -1073,7 +1084,8 @@ def run_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]) - ddbc_bindings.close_pooling() + assert pool.current_size == 0 + pool.close() """, conn_str, ) From e67da28e1e514ca792adcf2446bd0edb4cad7d78 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:27:29 -0400 Subject: [PATCH 03/12] FIX: Guard candidate validation and release decrements with pool generation (#746) - Capture candidate_generation when popping a candidate in ConnectionPool::acquire and guard validation failure cleanup against generation mismatches. - Capture release_generation in ConnectionPool::release to guard overflow disconnect decrements. - Add inject_candidate test helper to _TestConnectionPool. - Add regression test test_pool_size_accounting_race_on_candidate_validation_close_interleave. --- .../pybind/connection/connection_pool.cpp | 18 +++- .../pybind/connection/connection_pool.h | 7 ++ mssql_python/pybind/ddbc_bindings.cpp | 12 ++- tests/test_009_pooling.py | 97 +++++++++++++++++++ 4 files changed, 129 insertions(+), 5 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 7cdaca0f5..2ddcba38b 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -105,6 +105,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt uint64_t reservation_generation = 0; while (true) { std::shared_ptr candidate; + uint64_t candidate_generation = 0; { std::unique_lock lock(_mutex); if (_pool.empty()) { @@ -129,6 +130,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } candidate = _pool.front(); _pool.pop_front(); + candidate_generation = _generation; } // Validate the candidate outside the mutex. @@ -206,7 +208,10 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt [&](const std::shared_ptr& sibling) { if (sibling->currentAccessToken() == stale_token) { to_disconnect.push_back(sibling); - if (_current_size > 0) --_current_size; + if (_generation == candidate_generation && + _current_size > 0) { + --_current_size; + } return true; } return false; @@ -227,11 +232,13 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } // Candidate is dead, reset failed, or its token rotated — mark for - // disconnect and decrement the pool size. + // disconnect and decrement the pool size if the pool generation still matches (#746). to_disconnect.push_back(candidate); { std::lock_guard lock(_mutex); - if (_current_size > 0) --_current_size; + if (_generation == candidate_generation && _current_size > 0) { + --_current_size; + } } // If a rotated token was captured, reserve a slot and reopen with it @@ -315,8 +322,10 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt void ConnectionPool::release(std::shared_ptr conn) { PERF_TIMER("ConnectionPool::release"); bool should_disconnect = false; + uint64_t release_generation = 0; { std::lock_guard lock(_mutex); + release_generation = _generation; if (_pool.size() < _max_size) { conn->updateLastUsed(); _pool.push_back(conn); @@ -333,8 +342,9 @@ void ConnectionPool::release(std::shared_ptr conn) { LOG("ConnectionPool::release: disconnect failed: %s", ex.what()); } std::lock_guard lock(_mutex); - if (_current_size > 0) + if (_generation == release_generation && _current_size > 0) { --_current_size; + } } } diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index baf5c0abd..4a79aacb9 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -51,6 +51,13 @@ class ConnectionPool { return _generation; } + // 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; + } + private: size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 0ed652176..b2c9d279b 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6135,7 +6135,17 @@ PYBIND11_MODULE(ddbc_bindings, m) { py::arg("conn_str"), py::arg("token_factory") = py::none()) .def("close", &ConnectionPool::close) .def_property_readonly("current_size", &ConnectionPool::current_size) - .def_property_readonly("generation", &ConnectionPool::generation); + .def_property_readonly("generation", &ConnectionPool::generation) + .def( + "inject_candidate", + [](ConnectionPool& pool, const std::u16string& connStr, long long expiry) { + auto conn = std::make_shared(connStr, true); + if (expiry > 0) { + conn->setTokenExpiry(expiry); + } + pool.inject_candidate(conn); + }, + py::arg("conn_str"), py::arg("expiry") = 0); 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 d255c785c..9dee25ea1 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1091,6 +1091,103 @@ def run_b(): ) +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 popped candidate's + cleanup must not decrement a reservation created under the new pool generation. + """ + _run_in_subprocess( + """ + import threading + from mssql_python import ddbc_bindings + + pool = ddbc_bindings._TestConnectionPool(1, 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" + + # Thread A popped the candidate (generation 0) and is validating it in factory_a. + # Now pool.close() clears idle connections, resets current_size=0, and bumps generation=1. + pool.close() + assert pool.current_size == 0 + assert pool.generation == 1 + + # Thread B reserves the freed 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 == 1 + + # Thread A finishes factory_a (which raises an error / fails validation). + # Thread A's cleanup detects that the candidate was from generation 0 != 1, + # so it does NOT decrement current_size! + # On unpatched code, Thread A would decrement current_size to 0. + 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!)" + ) + + # Thread C must be rejected because max_size=1 and Thread B holds the slot + thread_c_rejected = False + try: + pool.acquire("SERVER=dummy_test_746;", lambda: {}) + except RuntimeError as exc: + if "pool size limit reached" in str(exc): + thread_c_rejected = True + + assert thread_c_rejected, "Thread C should have been rejected due to pool capacity limit" + + release_factory_b.set() + t_b.join(timeout=5.0) + assert pool.current_size == 0 + pool.close() + """, + conn_str, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From 6f8369c653ebf60d47b7300423c5b09d272502cd Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:31:13 -0400 Subject: [PATCH 04/12] TEST: Retain acquired Connection and expose release on _TestConnectionPool (#746) - Register _TestPooledConnection in ddbc_bindings. - Return acquired Connection from _TestConnectionPool::acquire instead of dropping it. - Expose ConnectionPool::release on _TestConnectionPool so acquired connections can be returned to the pool. --- mssql_python/pybind/ddbc_bindings.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index b2c9d279b..555d8c3d5 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6124,15 +6124,17 @@ PYBIND11_MODULE(ddbc_bindings, m) { }, "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"); 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) { - pool.acquire(connStr, py::dict(), 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::arg("conn")) .def("close", &ConnectionPool::close) .def_property_readonly("current_size", &ConnectionPool::current_size) .def_property_readonly("generation", &ConnectionPool::generation) From 5282437c8c4ed867a2ae795e04d8da6d1db037b3 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:44:18 -0400 Subject: [PATCH 05/12] FIX: Check pool generation on candidate reuse and connect success (#746) - Reject candidate reuse when the pool was closed during validation; discard stale candidate and retry acquire under the current generation. - Check generation validity after Phase 3 connect succeeds; if the pool was closed while connecting, discard the connection and retry acquire under the current generation rather than creating an uncounted connection exceeding max_size. - Add regression tests test_pool_size_accounting_race_on_successful_candidate_reuse_close_interleave and test_pool_size_accounting_race_on_successful_open_close_interleave. --- mssql_python/pybind/connection/connection.cpp | 14 + mssql_python/pybind/connection/connection.h | 3 + .../pybind/connection/connection_pool.cpp | 460 +++++++++--------- .../pybind/connection/connection_pool.h | 10 + mssql_python/pybind/ddbc_bindings.cpp | 8 +- tests/test_009_pooling.py | 206 ++++++++ 6 files changed, 472 insertions(+), 229 deletions(-) diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 8ccee11a7..e0d0b02aa 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -76,6 +76,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 +107,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 +515,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 +528,9 @@ bool Connection::isAlive() const { } bool Connection::reset() { + if (_isMock) { + return true; + } if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index d5300aff4..2b572b4b2 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -55,6 +55,8 @@ 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; } // Materialize connect-attrs from a Python token-factory callback. // The factory may return either a bare attrs dict (legacy) or a @@ -100,6 +102,7 @@ class Connection { std::u16string _connStr; bool _fromPool = false; bool _autocommit = true; + bool _isMock = false; SqlHandlePtr _dbcHandle; std::chrono::steady_clock::time_point _lastUsed; // POSIX-epoch expiry (seconds) of the access token this connection last diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 2ddcba38b..bac39e88b 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -62,249 +62,254 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt PERF_TIMER("ConnectionPool::acquire"); std::vector> 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(); + 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; - _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()); + // 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(); - 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; - } + _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()); - // 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; - uint64_t reservation_generation = 0; - 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; - reservation_generation = _generation; - needs_connect = true; - break; + 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; + } + + // 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; + 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(); + candidate_generation = _generation; } - candidate = _pool.front(); - _pool.pop_front(); - candidate_generation = _generation; - } - // 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); + _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 (_generation == candidate_generation && + _current_size > 0) { + --_current_size; + } + 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 (_generation == candidate_generation && - _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 if the pool generation still matches (#746). - to_disconnect.push_back(candidate); - { - std::lock_guard lock(_mutex); - if (_generation == candidate_generation && _current_size > 0) { - --_current_size; + if (reuse_candidate) { + bool gen_valid = false; + { + std::lock_guard lock(_mutex); + gen_valid = (_generation == candidate_generation); + } + if (gen_valid) { + valid_conn = candidate; + break; + } + // Pool was closed while validating candidate (#746); discard stale + // candidate and retry acquire. + to_disconnect.push_back(candidate); + continue; } - } - // 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; - reservation_generation = _generation; - needs_connect = true; - break; + // Candidate is dead, reset failed, or its token rotated — mark for + // disconnect and decrement the pool size if the pool generation still matches (#746). + to_disconnect.push_back(candidate); + { + std::lock_guard lock(_mutex); + if (_generation == candidate_generation && _current_size > 0) { + --_current_size; + } } - // 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 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) { - // 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); - } - } catch (...) { - // Construct/connect failed — release the reserved slot only if the pool - // has not been reset in the meantime (#746). If close() ran while we were - // connecting outside the lock, close() already set _current_size = 0 and - // bumped _generation; decrementing here would cancel another thread's - // newer reservation instead of our own. - { std::lock_guard lock(_mutex); - if (_generation == reservation_generation && _current_size > 0) { - --_current_size; + 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; + 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. + } + } + + 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); + gen_valid = (_generation == reservation_generation); + } + if (gen_valid) { + valid_conn = new_conn; + break; + } + // Pool was closed while connecting; queue the stale connection + // for disconnect and retry acquire under the new generation. + to_disconnect.push_back(new_conn); + } catch (...) { + // Construct/connect failed — release the reserved slot only if the pool + // has not been reset in the meantime (#746). If close() ran while we were + // connecting outside the lock, close() already set _current_size = 0 and + // bumped _generation; decrementing here would cancel another thread's + // newer reservation instead of our own. + { + std::lock_guard lock(_mutex); + if (_generation == reservation_generation && _current_size > 0) { + --_current_size; + } } + throw; } - throw; } } @@ -405,10 +410,11 @@ 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). diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index 4a79aacb9..6e9912f1b 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -6,6 +6,7 @@ #pragma once #include "connection/connection.h" +#include #include #include #include @@ -58,11 +59,20 @@ class ConnectionPool { ++_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; + } + private: size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale size_t _current_size = 0; uint64_t _generation = 0; // Pool reset generation for reservation attribution (#746) + std::atomic _mock_mode{false}; std::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access }; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 555d8c3d5..3508c150a 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6136,18 +6136,22 @@ PYBIND11_MODULE(ddbc_bindings, m) { py::arg("conn_str"), py::arg("token_factory") = py::none()) .def("release", &ConnectionPool::release, py::arg("conn")) .def("close", &ConnectionPool::close) + .def("set_mock_mode", &ConnectionPool::set_mock_mode, py::arg("enable") = true) .def_property_readonly("current_size", &ConnectionPool::current_size) .def_property_readonly("generation", &ConnectionPool::generation) .def( "inject_candidate", - [](ConnectionPool& pool, const std::u16string& connStr, long long expiry) { + [](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); } pool.inject_candidate(conn); }, - py::arg("conn_str"), py::arg("expiry") = 0); + 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 9dee25ea1..2d2445b85 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1188,6 +1188,212 @@ def run_b(): ) +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(1, 600) + pool.set_mock_mode(True) + # Inject candidate with near-expiry token to trigger token-factory validation + 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" + 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" + + # Thread A popped candidate (generation 0). Now pool.close() wipes the pool. + pool.close() + assert pool.current_size == 0 + assert pool.generation == 1 + + # Thread B reserves the freed 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" + 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 == 1 + + # Thread A's validation succeeds. Under the generation fix, Thread A detects + # generation mismatch (0 != 1), discards the stale candidate, and retries acquire. + # Since max_size=1 and Thread B currently holds the generation 1 reservation, + # Thread A's retry is rejected with 'pool size limit reached' instead of handing out + # an uncounted live connection. + release_factory_a.set() + t_a.join(timeout=5.0) + + assert len(t_a_error) == 1 and "pool size limit reached" in str(t_a_error[0]), ( + f"Expected Thread A to be rejected due to pool limit, got {t_a_error}" + ) + assert len(t_a_conn) == 0 + + # Thread B finishes connecting and successfully checks out its connection + release_factory_b.set() + t_b.join(timeout=5.0) + assert len(t_b_error) == 0 + assert len(t_b_conn) == 1 + assert pool.current_size == 1 + + # Return Thread B's connection and clean up + pool.release(t_b_conn[0]) + assert pool.current_size == 1 + 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(1, 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() + + def factory_a(): + 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 + + # Thread A reserved slot under generation 0. Now pool.close() wipes the pool. + pool.close() + assert pool.current_size == 0 + assert pool.generation == 1 + + # Thread B reserves the freed 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" + 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 == 1 + + # Thread A's open succeeds. Under the generation fix, Thread A detects + # reservation generation mismatch (0 != 1), discards the opened connection, + # and retries acquire under generation 1. + # Since max_size=1 and Thread B holds the generation 1 slot, Thread A's retry + # fails with 'pool size limit reached' rather than creating a 2nd concurrent connection. + release_factory_a.set() + t_a.join(timeout=5.0) + + assert len(t_a_error) == 1 and "pool size limit reached" in str(t_a_error[0]), ( + f"Expected Thread A to be rejected due to pool limit, got {t_a_error}" + ) + assert len(t_a_conn) == 0 + + # Thread B finishes connecting and successfully checks out its connection + release_factory_b.set() + t_b.join(timeout=5.0) + assert len(t_b_error) == 0 + assert len(t_b_conn) == 1 + assert pool.current_size == 1 + + # Return Thread B's connection and clean up + pool.release(t_b_conn[0]) + assert pool.current_size == 1 + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From 2753fb77170fe4d2d7ed1150f036046e45e155f7 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 18:51:58 -0400 Subject: [PATCH 06/12] FIX: Atomically publish acquired connections and track pool origin on release (#746) - Assign valid_conn and stamp pool origin under _mutex when generation matches to prevent check-then-publish races with close(). - Track originating pool and generation on Connection; reject and disconnect stale connections released after pool.close() or pool recreation to prevent uncounted connections in new pool. - Initialize _lastUsed to steady_clock::now() in Connection constructor and inject_candidate to prevent premature pruning during race tests. - Add regression test test_pool_release_from_stale_generation_does_not_pollute_pool. --- mssql_python/pybind/connection/connection.cpp | 5 +- mssql_python/pybind/connection/connection.h | 11 +++ .../pybind/connection/connection_pool.cpp | 34 +++++++--- mssql_python/pybind/ddbc_bindings.cpp | 2 + tests/test_009_pooling.py | 67 +++++++++++++++++++ 5 files changed, 108 insertions(+), 11 deletions(-) diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index e0d0b02aa..ee3f4e242 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(); } diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index 2b572b4b2..2a7eedab4 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -57,6 +57,15 @@ class Connection { std::chrono::steady_clock::time_point lastUsed() const; void setMock(bool mock) { _isMock = mock; } bool isMock() const { return _isMock; } + void setPoolOrigin(void* pool, uint64_t generation) { + _originPool = pool; + _originGeneration = generation; + } + bool matchesPoolOrigin(void* pool, uint64_t generation) const { + return _originPool == pool && _originGeneration == generation; + } + void* originPool() const { return _originPool; } + 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 @@ -103,6 +112,8 @@ class Connection { bool _fromPool = false; bool _autocommit = true; bool _isMock = false; + void* _originPool = nullptr; + uint64_t _originGeneration = 0; SqlHandlePtr _dbcHandle; std::chrono::steady_clock::time_point _lastUsed; // POSIX-epoch expiry (seconds) of the access token this connection last diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index bac39e88b..80a64226f 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -202,10 +202,14 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt bool gen_valid = false; { std::lock_guard lock(_mutex); - gen_valid = (_generation == candidate_generation); + if (_generation == candidate_generation) { + candidate->updateLastUsed(); + candidate->setPoolOrigin(this, _generation); + valid_conn = candidate; + gen_valid = true; + } } if (gen_valid) { - valid_conn = candidate; break; } // Pool was closed while validating candidate (#746); discard stale @@ -287,10 +291,14 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt bool gen_valid = false; { std::lock_guard lock(_mutex); - gen_valid = (_generation == reservation_generation); + if (_generation == reservation_generation) { + new_conn->updateLastUsed(); + new_conn->setPoolOrigin(this, _generation); + valid_conn = new_conn; + gen_valid = true; + } } if (gen_valid) { - valid_conn = new_conn; break; } // Pool was closed while connecting; queue the stale connection @@ -326,12 +334,16 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt void ConnectionPool::release(std::shared_ptr conn) { PERF_TIMER("ConnectionPool::release"); + if (!conn) { + return; + } bool should_disconnect = false; - uint64_t release_generation = 0; + bool gen_matches = false; + uint64_t conn_gen = conn->originGeneration(); { std::lock_guard lock(_mutex); - release_generation = _generation; - if (_pool.size() < _max_size) { + gen_matches = conn->matchesPoolOrigin(this, _generation); + if (gen_matches && _pool.size() < _max_size) { conn->updateLastUsed(); _pool.push_back(conn); } else { @@ -346,9 +358,11 @@ void ConnectionPool::release(std::shared_ptr conn) { } catch (const std::exception& ex) { LOG("ConnectionPool::release: disconnect failed: %s", ex.what()); } - std::lock_guard lock(_mutex); - if (_generation == release_generation && _current_size > 0) { - --_current_size; + if (gen_matches) { + std::lock_guard lock(_mutex); + if (_generation == conn_gen && _current_size > 0) { + --_current_size; + } } } } diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 3508c150a..0700d3cf1 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6149,6 +6149,8 @@ PYBIND11_MODULE(ddbc_bindings, m) { if (expiry > 0) { conn->setTokenExpiry(expiry); } + conn->updateLastUsed(); + conn->setPoolOrigin(&pool, pool.generation()); pool.inject_candidate(conn); }, py::arg("conn_str"), py::arg("expiry") = 0, py::arg("mock") = false); diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 2d2445b85..01a05bc9c 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1394,6 +1394,73 @@ def run_b(): ) +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.generation == 0 + + # 2. Pool is closed while conn_1 is still checked out + pool.close() + assert pool.current_size == 0 + assert pool.generation == 1 + + # 3. Acquire conn_2 under generation 1 (consumes the 1 slot of max_size=1) + conn_2 = pool.acquire("SERVER=dummy_test_746;", lambda: {}) + assert conn_2 is not None + assert pool.current_size == 1 + assert pool.generation == 1 + + # 4. Release conn_1 (from generation 0). + # The pool origin check ensures conn_1 is NOT added to the idle pool + # and current_size of generation 1 is NOT decremented. + pool.release(conn_1) + assert pool.current_size == 1, ( + f"Expected pool.current_size to stay 1, but got {pool.current_size}" + ) + + # 5. Since conn_2 is still checked out and max_size=1, acquire must be rejected + 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 when max_size=1 capacity is occupied" + + # 6. Release conn_2 (matches generation 1). It returns to the pool. + pool.release(conn_2) + assert pool.current_size == 1 + + # 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 + + pool.release(conn_3) + pool.close() + assert pool.current_size == 0 + """, + conn_str, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From c4cdde545e739a41e7b448dbf174dead6c9e372c Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 19:03:54 -0400 Subject: [PATCH 07/12] FIX: Monotonic pool identity and checked-out capacity retention on close (#746) --- mssql_python/pybind/connection/connection.h | 12 +-- .../pybind/connection/connection_pool.cpp | 84 +++++++++------ .../pybind/connection/connection_pool.h | 11 +- mssql_python/pybind/ddbc_bindings.cpp | 4 +- tests/test_009_pooling.py | 100 ++++++++++++++---- 5 files changed, 152 insertions(+), 59 deletions(-) diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index 2a7eedab4..d0bee3832 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -57,14 +57,14 @@ class Connection { std::chrono::steady_clock::time_point lastUsed() const; void setMock(bool mock) { _isMock = mock; } bool isMock() const { return _isMock; } - void setPoolOrigin(void* pool, uint64_t generation) { - _originPool = pool; + void setPoolOrigin(uint64_t pool_id, uint64_t generation) { + _originPoolId = pool_id; _originGeneration = generation; } - bool matchesPoolOrigin(void* pool, uint64_t generation) const { - return _originPool == pool && _originGeneration == generation; + bool matchesPoolOrigin(uint64_t pool_id, uint64_t generation) const { + return _originPoolId == pool_id && _originGeneration == generation; } - void* originPool() const { return _originPool; } + uint64_t originPoolId() const { return _originPoolId; } uint64_t originGeneration() const { return _originGeneration; } // Materialize connect-attrs from a Python token-factory callback. @@ -112,7 +112,7 @@ class Connection { bool _fromPool = false; bool _autocommit = true; bool _isMock = false; - void* _originPool = nullptr; + uint64_t _originPoolId = 0; uint64_t _originGeneration = 0; SqlHandlePtr _dbcHandle; std::chrono::steady_clock::time_point _lastUsed; diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 80a64226f..403673d15 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -51,10 +51,15 @@ 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), + _pool_id(s_next_pool_id.fetch_add(1)) {} std::shared_ptr ConnectionPool::acquire(const std::u16string& connStr, const py::dict& attrs_before, @@ -173,21 +178,22 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt 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 (_generation == candidate_generation && - _current_size > 0) { - --_current_size; + 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); + if (_current_size > 0) { + --_current_size; + } + return true; } - return true; - } - return false; - }), - _pool.end()); + return false; + }), + _pool.end()); + } } } } @@ -204,8 +210,9 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt std::lock_guard lock(_mutex); if (_generation == candidate_generation) { candidate->updateLastUsed(); - candidate->setPoolOrigin(this, _generation); + candidate->setPoolOrigin(_pool_id, _generation); valid_conn = candidate; + ++_checked_out; gen_valid = true; } } @@ -293,8 +300,9 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt std::lock_guard lock(_mutex); if (_generation == reservation_generation) { new_conn->updateLastUsed(); - new_conn->setPoolOrigin(this, _generation); + new_conn->setPoolOrigin(_pool_id, _generation); valid_conn = new_conn; + ++_checked_out; gen_valid = true; } } @@ -307,7 +315,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } catch (...) { // Construct/connect failed — release the reserved slot only if the pool // has not been reset in the meantime (#746). If close() ran while we were - // connecting outside the lock, close() already set _current_size = 0 and + // connecting outside the lock, close() already set _current_size = _checked_out and // bumped _generation; decrementing here would cancel another thread's // newer reservation instead of our own. { @@ -338,14 +346,27 @@ void ConnectionPool::release(std::shared_ptr conn) { return; } bool should_disconnect = false; - bool gen_matches = false; - uint64_t conn_gen = conn->originGeneration(); { std::lock_guard lock(_mutex); - gen_matches = conn->matchesPoolOrigin(this, _generation); - if (gen_matches && _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; + } + if (_current_size > 0) { + --_current_size; + } + conn->setPoolOrigin(0, 0); + } } else { should_disconnect = true; } @@ -358,12 +379,6 @@ void ConnectionPool::release(std::shared_ptr conn) { } catch (const std::exception& ex) { LOG("ConnectionPool::release: disconnect failed: %s", ex.what()); } - if (gen_matches) { - std::lock_guard lock(_mutex); - if (_generation == conn_gen && _current_size > 0) { - --_current_size; - } - } } } @@ -372,8 +387,9 @@ 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) { return false; } // Nothing checked out and the pool is empty: safe to drop immediately. @@ -407,7 +423,9 @@ void ConnectionPool::close() { to_close.push_back(_pool.front()); _pool.pop_front(); } - _current_size = 0; + // Retain reserved capacity for checked-out connections so a new acquire + // cannot exceed _max_size while old connections are still live (#746). + _current_size = _checked_out; ++_generation; } for (auto& conn : to_close) { diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index 6e9912f1b..5bd15b59d 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -42,15 +42,22 @@ class ConnectionPool { // dropped by the manager to reclaim memory (lazy eviction). bool canEvict(); - // Test accessors for pool generation and current size + // 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; + } 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) { @@ -71,7 +78,9 @@ class ConnectionPool { 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) 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::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 0700d3cf1..c500c017a 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6138,7 +6138,9 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("close", &ConnectionPool::close) .def("set_mock_mode", &ConnectionPool::set_mock_mode, py::arg("enable") = true) .def_property_readonly("current_size", &ConnectionPool::current_size) + .def_property_readonly("checked_out", &ConnectionPool::checked_out) .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) { @@ -6150,7 +6152,7 @@ PYBIND11_MODULE(ddbc_bindings, m) { conn->setTokenExpiry(expiry); } conn->updateLastUsed(); - conn->setPoolOrigin(&pool, pool.generation()); + conn->setPoolOrigin(pool.pool_id(), pool.generation()); pool.inject_candidate(conn); }, py::arg("conn_str"), py::arg("expiry") = 0, py::arg("mock") = false); diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 01a05bc9c..6172cfde5 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1414,46 +1414,54 @@ def test_pool_release_from_stale_generation_does_not_pollute_pool(conn_str): 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 + # 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 == 0 - assert pool.generation == 1 - - # 3. Acquire conn_2 under generation 1 (consumes the 1 slot of max_size=1) - 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 - # 4. Release conn_1 (from generation 0). - # The pool origin check ensures conn_1 is NOT added to the idle pool - # and current_size of generation 1 is NOT decremented. - pool.release(conn_1) - assert pool.current_size == 1, ( - f"Expected pool.current_size to stay 1, but got {pool.current_size}" - ) - - # 5. Since conn_2 is still checked out and max_size=1, acquire must be rejected + # 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 when max_size=1 capacity is occupied" + 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 - # 6. Release conn_2 (matches generation 1). It returns to the pool. + # 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 """, @@ -1461,6 +1469,62 @@ def test_pool_release_from_stale_generation_does_not_pollute_pool(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, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From b1fe514a3efd72185a0f9d5048b0202db5e65f55 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 19:13:56 -0400 Subject: [PATCH 08/12] FIX: Account for in-flight opens across close to prevent capacity overflow (#746) --- .../pybind/connection/connection_pool.cpp | 84 +++++-- .../pybind/connection/connection_pool.h | 5 + mssql_python/pybind/ddbc_bindings.cpp | 1 + tests/test_009_pooling.py | 238 +++++++++++------- 4 files changed, 224 insertions(+), 104 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 403673d15..ee63d53c3 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -59,6 +59,7 @@ ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs) _idle_timeout_secs(idle_timeout_secs), _current_size(0), _checked_out(0), + _in_flight(0), _pool_id(s_next_pool_id.fetch_add(1)) {} std::shared_ptr ConnectionPool::acquire(const std::u16string& connStr, @@ -124,6 +125,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // 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; @@ -137,6 +139,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } candidate = _pool.front(); _pool.pop_front(); + ++_in_flight; candidate_generation = _generation; } @@ -213,6 +216,9 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt candidate->setPoolOrigin(_pool_id, _generation); valid_conn = candidate; ++_checked_out; + if (_in_flight > 0) { + --_in_flight; + } gen_valid = true; } } @@ -220,17 +226,37 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt break; } // Pool was closed while validating candidate (#746); discard stale - // candidate and retry acquire. - to_disconnect.push_back(candidate); + // 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; } - // Candidate is dead, reset failed, or its token rotated — mark for - // disconnect and decrement the pool size if the pool generation still matches (#746). - to_disconnect.push_back(candidate); + // 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()); + } { std::lock_guard lock(_mutex); - if (_generation == candidate_generation && _current_size > 0) { + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { --_current_size; } } @@ -246,6 +272,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // 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; @@ -303,24 +330,42 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt 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; queue the stale connection - // for disconnect and retry acquire under the new generation. - to_disconnect.push_back(new_conn); + // 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). + 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 only if the pool - // has not been reset in the meantime (#746). If close() ran while we were - // connecting outside the lock, close() already set _current_size = _checked_out and - // bumped _generation; decrementing here would cancel another thread's - // newer reservation instead of our own. + // Construct/connect failed — release the reserved slot and in-flight count. { std::lock_guard lock(_mutex); - if (_generation == reservation_generation && _current_size > 0) { + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { --_current_size; } } @@ -389,7 +434,7 @@ bool ConnectionPool::canEvict() { // caller still holds one, so the pool must stay. 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) { + 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. @@ -423,9 +468,10 @@ void ConnectionPool::close() { to_close.push_back(_pool.front()); _pool.pop_front(); } - // Retain reserved capacity for checked-out connections so a new acquire - // cannot exceed _max_size while old connections are still live (#746). - _current_size = _checked_out; + // Retain reserved capacity for checked-out connections and in-flight opens + // so a new acquire cannot exceed _max_size while old connections or opens + // are still live (#746). + _current_size = _checked_out + _in_flight; ++_generation; } for (auto& conn : to_close) { diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index 5bd15b59d..abbc0c097 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -51,6 +51,10 @@ class ConnectionPool { 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; @@ -79,6 +83,7 @@ class ConnectionPool { 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}; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index c500c017a..cb0b765ed 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6139,6 +6139,7 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("set_mock_mode", &ConnectionPool::set_mock_mode, py::arg("enable") = true) .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( diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 6172cfde5..c2f7e5e9e 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1003,7 +1003,7 @@ def test_pool_size_accounting_race_on_close_interleave(conn_str): import threading from mssql_python import ddbc_bindings - pool = ddbc_bindings._TestConnectionPool(1, 600) + pool = ddbc_bindings._TestConnectionPool(2, 600) assert pool.current_size == 0 assert pool.generation == 0 @@ -1027,14 +1027,16 @@ def 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 the slot. Now pool.close() resets the pool - # and increments the pool generation counter. + # 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 == 0 + assert pool.current_size == 1 + assert pool.in_flight == 1 assert pool.generation == 1 - # Thread B initiates acquire and reserves the freed slot under the new generation. + # Thread B initiates acquire and reserves the 2nd slot under the new generation. in_factory_b = threading.Event() release_factory_b = threading.Event() @@ -1054,11 +1056,11 @@ def run_b(): 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 == 1 + assert pool.current_size == 2 + assert pool.in_flight == 2 - # Thread A now raises its error. With the generation counter fix, its cleanup - # detects that the pool generation changed (0 != 1) and does NOT decrement current_size. - # On buggy code without the fix, Thread A's cleanup would decrement current_size back to 0. + # 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]) @@ -1067,24 +1069,14 @@ def run_b(): assert pool.current_size == 1, ( f"Expected pool.current_size to be 1, but got {pool.current_size} (drift occurred!)" ) - - # Thread C now attempts to acquire on the same pool. - # Since max_size=1 and Thread B is still reserving the slot, Thread C must fail - # with 'pool size limit reached'. - thread_c_rejected = False - try: - pool.acquire("SERVER=dummy_test_746;", lambda: {}) - except RuntimeError as exc: - if "pool size limit reached" in str(exc): - thread_c_rejected = True - - assert thread_c_rejected, "Thread C should have been rejected due to pool capacity limit" + 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, @@ -1103,7 +1095,7 @@ def test_pool_size_accounting_race_on_candidate_validation_close_interleave(conn import threading from mssql_python import ddbc_bindings - pool = ddbc_bindings._TestConnectionPool(1, 600) + 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 @@ -1128,14 +1120,16 @@ def run_a(): 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. - # Now pool.close() clears idle connections, resets current_size=0, and bumps generation=1. + # Pool close retains capacity for in-flight validation: current_size stays 1. pool.close() - assert pool.current_size == 0 + assert pool.current_size == 1 + assert pool.in_flight == 1 assert pool.generation == 1 - # Thread B reserves the freed slot under generation 1. + # Thread B reserves the 2nd slot under generation 1. in_factory_b = threading.Event() release_factory_b = threading.Event() @@ -1155,12 +1149,12 @@ def run_b(): 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 == 1 + assert pool.current_size == 2 + assert pool.in_flight == 2 - # Thread A finishes factory_a (which raises an error / fails validation). - # Thread A's cleanup detects that the candidate was from generation 0 != 1, - # so it does NOT decrement current_size! - # On unpatched code, Thread A would decrement current_size to 0. + # 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) @@ -1168,20 +1162,12 @@ def run_b(): assert pool.current_size == 1, ( f"Expected pool.current_size to be 1, but got {pool.current_size} (drift occurred!)" ) - - # Thread C must be rejected because max_size=1 and Thread B holds the slot - thread_c_rejected = False - try: - pool.acquire("SERVER=dummy_test_746;", lambda: {}) - except RuntimeError as exc: - if "pool size limit reached" in str(exc): - thread_c_rejected = True - - assert thread_c_rejected, "Thread C should have been rejected due to pool capacity limit" + 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, @@ -1203,10 +1189,10 @@ def test_pool_size_accounting_race_on_successful_candidate_reuse_close_interleav import threading from mssql_python import ddbc_bindings - pool = ddbc_bindings._TestConnectionPool(1, 600) + pool = ddbc_bindings._TestConnectionPool(2, 600) pool.set_mock_mode(True) - # Inject candidate with near-expiry token to trigger token-factory validation - pool.inject_candidate("SERVER=dummy_test_746;", 1) + # 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 @@ -1231,13 +1217,15 @@ def run_a(): 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). Now pool.close() wipes the pool. + # Thread A popped candidate (generation 0). Pool close retains in-flight validation. pool.close() - assert pool.current_size == 0 + assert pool.current_size == 1 + assert pool.in_flight == 1 assert pool.generation == 1 - # Thread B reserves the freed slot under generation 1 + # Thread B reserves slot 2 under generation 1 in_factory_b = threading.Event() release_factory_b = threading.Event() @@ -1259,31 +1247,33 @@ def run_b(): 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 == 1 + assert pool.current_size == 2 + assert pool.in_flight == 2 - # Thread A's validation succeeds. Under the generation fix, Thread A detects - # generation mismatch (0 != 1), discards the stale candidate, and retries acquire. - # Since max_size=1 and Thread B currently holds the generation 1 reservation, - # Thread A's retry is rejected with 'pool size limit reached' instead of handing out - # an uncounted live connection. - release_factory_a.set() - t_a.join(timeout=5.0) - - assert len(t_a_error) == 1 and "pool size limit reached" in str(t_a_error[0]), ( - f"Expected Thread A to be rejected due to pool limit, got {t_a_error}" - ) - assert len(t_a_conn) == 0 - - # Thread B finishes connecting and successfully checks out its connection + # 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.current_size == 1 + assert pool.checked_out == 1 + assert pool.in_flight == 1 + assert pool.current_size == 2 - # Return Thread B's connection and clean up + # 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 + 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.current_size == 1 + assert pool.checked_out == 0 pool.close() assert pool.current_size == 0 """, @@ -1307,7 +1297,7 @@ def test_pool_size_accounting_race_on_successful_open_close_interleave(conn_str) import threading from mssql_python import ddbc_bindings - pool = ddbc_bindings._TestConnectionPool(1, 600) + pool = ddbc_bindings._TestConnectionPool(2, 600) pool.set_mock_mode(True) assert pool.current_size == 0 assert pool.generation == 0 @@ -1334,13 +1324,16 @@ def 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. Now pool.close() wipes the pool. + # 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 == 0 + assert pool.current_size == 1 + assert pool.in_flight == 1 assert pool.generation == 1 - # Thread B reserves the freed slot under 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() @@ -1362,31 +1355,106 @@ def run_b(): 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 == 1 - - # Thread A's open succeeds. Under the generation fix, Thread A detects - # reservation generation mismatch (0 != 1), discards the opened connection, - # and retries acquire under generation 1. - # Since max_size=1 and Thread B holds the generation 1 slot, Thread A's retry - # fails with 'pool size limit reached' rather than creating a 2nd concurrent connection. - release_factory_a.set() - t_a.join(timeout=5.0) + assert pool.current_size == 2 + assert pool.in_flight == 2 - assert len(t_a_error) == 1 and "pool size limit reached" in str(t_a_error[0]), ( - f"Expected Thread A to be rejected due to pool limit, got {t_a_error}" - ) - assert len(t_a_conn) == 0 - - # Thread B finishes connecting and successfully checks out its connection + # 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 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 + 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 - # Return Thread B's connection and clean up - pool.release(t_b_conn[0]) + # 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 """, From 3f16de23e271d4cf6c3d55edce68a9ad2aa6158c Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 19:59:02 -0400 Subject: [PATCH 09/12] FIX: Defer replacement pool creation while connections checked out (#746) --- mssql_python/pybind/connection/connection.cpp | 3 + mssql_python/pybind/connection/connection.h | 3 + .../pybind/connection/connection_pool.cpp | 48 ++++++++--- .../pybind/connection/connection_pool.h | 16 ++++ mssql_python/pybind/ddbc_bindings.cpp | 11 ++- tests/test_009_pooling.py | 84 ++++++++++++++++++- 6 files changed, 149 insertions(+), 16 deletions(-) diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index ee3f4e242..dcea07d75 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -659,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 d0bee3832..2273a31c4 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -164,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 ee63d53c3..9e88451c1 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -548,9 +548,20 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( } } } + // Defer replacement-pool creation if the existing pool still has live + // work. If the existing pool has finished all live work (canEvict() == true), + // evict it now and create a fresh replacement pool (#746). + auto it = _pools.find(key); + if (it != _pools.end() && it->second && it->second->canEvict()) { + evicted.push_back(it->second); + _pools.erase(it); + } 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; @@ -620,11 +631,10 @@ 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(). + // Under _manager_mutex, snapshot all pools to close their idle connections. + // We do not clear _pools immediately: close() disconnects ODBC handles + // (releasing the GIL), which must run outside _manager_mutex to avoid + // deadlock. std::vector> to_close; { std::lock_guard lock(_manager_mutex); @@ -634,12 +644,9 @@ void ConnectionPoolManager::closePools() { 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(); @@ -647,6 +654,27 @@ void ConnectionPoolManager::closePools() { LOG("ConnectionPoolManager::closePools: closing pool failed: %s", ex.what()); } } + { + std::lock_guard lock(_manager_mutex); + // Only evict pools that have no live work left (canEvict() == true). + // 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->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) { diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index abbc0c097..bcd2c35c9 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -127,6 +127,21 @@ 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; + } + private: ConnectionPoolManager() = default; ~ConnectionPoolManager() = default; @@ -146,6 +161,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 cb0b765ed..d95611923 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6109,9 +6109,14 @@ 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("disable_pooling", []() { // Disarm new-pool creation *before* closing so a connect racing this // disable cannot resurrect a pool after the map is cleared: any @@ -6124,7 +6129,9 @@ PYBIND11_MODULE(ddbc_bindings, m) { }, "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"); + py::class_>(m, "_TestPooledConnection") + .def_property_readonly("origin_generation", &Connection::originGeneration) + .def_property_readonly("origin_pool_id", &Connection::originPoolId); py::class_>(m, "_TestConnectionPool") .def(py::init(), py::arg("max_size") = 1, py::arg("idle_timeout_secs") = 600) .def( diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index c2f7e5e9e..6ffd6214d 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1198,10 +1198,13 @@ def test_pool_size_accounting_race_on_successful_candidate_reuse_close_interleav in_factory_a = threading.Event() release_factory_a = threading.Event() + factory_a_calls = [0] def factory_a(): - in_factory_a.set() - assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release 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 = [] @@ -1266,6 +1269,9 @@ def run_b(): 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 @@ -1304,10 +1310,13 @@ def test_pool_size_accounting_race_on_successful_open_close_interleave(conn_str) in_factory_a = threading.Event() release_factory_a = threading.Event() + factory_a_calls = [0] def factory_a(): - in_factory_a.set() - assert release_factory_a.wait(timeout=5.0), "Timed out waiting to release 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 = [] @@ -1364,6 +1373,7 @@ def run_b(): 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 @@ -1378,6 +1388,9 @@ def run_b(): 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 @@ -1593,6 +1606,69 @@ def test_pool_release_after_pool_recreation(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, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From f0c8804cf14712f791c5570b425ce8779540179e Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 21:46:34 -0400 Subject: [PATCH 10/12] FIX: Retain in-flight capacity during teardown and enforce use_count on eviction (#746) --- .../pybind/connection/connection_pool.cpp | 79 ++++++-- .../pybind/connection/connection_pool.h | 6 + mssql_python/pybind/ddbc_bindings.cpp | 17 +- tests/test_009_pooling.py | 188 +++++++++++++++++- 4 files changed, 262 insertions(+), 28 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 9e88451c1..f4dc71c7c 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -343,6 +343,14 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // 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). + std::function hook; + { + std::lock_guard lock(_mutex); + hook = _on_disconnect_hook; + } + if (hook) { + hook(); + } try { new_conn->disconnect(); } catch (const std::exception& ex) { @@ -391,6 +399,7 @@ void ConnectionPool::release(std::shared_ptr conn) { return; } bool should_disconnect = false; + bool decrement_in_flight = false; { std::lock_guard lock(_mutex); if (conn->originPoolId() == _pool_id) { @@ -406,9 +415,11 @@ void ConnectionPool::release(std::shared_ptr conn) { should_disconnect = true; if (_checked_out > 0) { --_checked_out; - } - if (_current_size > 0) { - --_current_size; + // 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); } @@ -419,11 +430,28 @@ 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) { + std::function hook; + { + std::lock_guard lock(_mutex); + hook = _on_disconnect_hook; + } + if (hook) { + hook(); + } try { conn->disconnect(); } catch (const std::exception& ex) { LOG("ConnectionPool::release: disconnect failed: %s", ex.what()); } + if (decrement_in_flight) { + std::lock_guard lock(_mutex); + if (_in_flight > 0) { + --_in_flight; + } + if (_current_size > 0) { + --_current_size; + } + } } } @@ -441,12 +469,7 @@ bool ConnectionPool::canEvict() { 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) { @@ -468,18 +491,36 @@ void ConnectionPool::close() { to_close.push_back(_pool.front()); _pool.pop_front(); } - // Retain reserved capacity for checked-out connections and in-flight opens - // so a new acquire cannot exceed _max_size while old connections or opens - // are still live (#746). + // 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) { + std::function hook; + { + std::lock_guard lock(_mutex); + hook = _on_disconnect_hook; + } + if (hook) { + hook(); + } 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; + } + } } } @@ -549,10 +590,12 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( } } // Defer replacement-pool creation if the existing pool still has live - // work. If the existing pool has finished all live work (canEvict() == true), - // evict it now and create a fresh replacement pool (#746). + // work. If the existing pool has finished all live work (canEvict() == true) + // and is not held by concurrent acquirers (use_count() == 1), evict it now + // and create a fresh replacement pool (#746). auto it = _pools.find(key); - if (it != _pools.end() && it->second && it->second->canEvict()) { + if (it != _pools.end() && it->second && it->second.use_count() == 1 && + it->second->canEvict()) { evicted.push_back(it->second); _pools.erase(it); } @@ -654,9 +697,11 @@ 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). + // 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 @@ -665,7 +710,7 @@ void ConnectionPoolManager::closePools() { // 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->canEvict()) { + if (!it->second || (it->second.use_count() == 1 && it->second->canEvict())) { it = _pools.erase(it); } else { ++it; diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index bcd2c35c9..fc1abeae4 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -77,6 +78,10 @@ class ConnectionPool { bool mock_mode() const { return _mock_mode; } + void set_on_disconnect_hook(std::function hook) { + std::lock_guard lock(_mutex); + _on_disconnect_hook = hook; + } private: size_t _max_size; // Maximum number of connections allowed @@ -87,6 +92,7 @@ class ConnectionPool { 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::function _on_disconnect_hook; std::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access }; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index d95611923..a36c81600 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6141,9 +6141,22 @@ PYBIND11_MODULE(ddbc_bindings, m) { return pool.acquire(connStr, py::dict(), token_factory); }, py::arg("conn_str"), py::arg("token_factory") = py::none()) - .def("release", &ConnectionPool::release, py::arg("conn")) - .def("close", &ConnectionPool::close) + .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 { + pool.set_on_disconnect_hook([hook]() { + py::gil_scoped_acquire gil; + hook(); + }); + } + }, + 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) diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 6ffd6214d..9366fc03b 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -987,13 +987,11 @@ def test_pooling_state_consistency(conn_str): 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(), the failed thread's - decrement must not cancel a newer generation's reservation on that pool. - With max_size=1, if Thread A's open failure wrongly decrements the counter - after Thread B reserved the slot under the new generation, Thread C would - be allowed to connect, exceeding max_size. The generation counter guarantees - that Thread A's cleanup only decrements if the pool generation still matches - its reservation. + 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. @@ -1087,8 +1085,10 @@ def test_pool_size_accounting_race_on_candidate_validation_close_interleave(conn """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 popped candidate's - cleanup must not decrement a reservation created under the new pool generation. + 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( """ @@ -1669,6 +1669,176 @@ def test_pool_manager_defers_replacement_while_connection_checked_out(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, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From fa8ae17c71c78004c7e295f52e12a5ef1efec531 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 21:55:00 -0400 Subject: [PATCH 11/12] FIX: Retain in-flight capacity for pruned and sibling connections, serialize replacement pool (#746) --- .../pybind/connection/connection_pool.cpp | 115 ++++++++++++++---- .../pybind/connection/connection_pool.h | 3 + mssql_python/pybind/ddbc_bindings.cpp | 37 +++++- tests/test_009_pooling.py | 85 +++++++++++++ 4 files changed, 210 insertions(+), 30 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index f4dc71c7c..1e0a4c668 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -62,11 +62,49 @@ ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs) _in_flight(0), _pool_id(s_next_pool_id.fetch_add(1)) {} +void ConnectionPool::drainDisconnectList(std::vector>& list) { + for (auto& conn : list) { + if (!conn) { + continue; + } + std::function hook; + { + std::lock_guard lock(_mutex); + hook = _on_disconnect_hook; + } + if (hook) { + hook(); + } + 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; while (valid_conn == nullptr) { @@ -97,13 +135,16 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt _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; + // 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 @@ -188,9 +229,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt [&](const std::shared_ptr& sibling) { if (sibling->currentAccessToken() == stale_token) { to_disconnect.push_back(sibling); - if (_current_size > 0) { - --_current_size; - } + ++_in_flight; return true; } return false; @@ -265,6 +304,8 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // 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 @@ -382,14 +423,8 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } } - // 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()); - } - } + // Phase 4: Disconnect expired/bad connections outside lock and decrement in-flight capacity. + drainDisconnectList(to_disconnect); return valid_conn; } @@ -539,6 +574,7 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( // 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; { @@ -591,23 +627,50 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( } // 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 now - // and create a fresh replacement pool (#746). + // 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()) { - evicted.push_back(it->second); + old_pool_to_close = it->second; _pools.erase(it); } - 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); + 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; + } + } + 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) { + return nullptr; + } + 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; } - created = true; + pool = pool_ref; } - 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 diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index fc1abeae4..627261d48 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -39,6 +39,9 @@ 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(); diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index a36c81600..ddc5a55c1 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6132,6 +6132,38 @@ PYBIND11_MODULE(ddbc_bindings, m) { py::class_>(m, "_TestPooledConnection") .def_property_readonly("origin_generation", &Connection::originGeneration) .def_property_readonly("origin_pool_id", &Connection::originPoolId); + struct GilSafeCallback { + py::object obj; + explicit GilSafeCallback(py::object o) : obj(std::move(o)) {} + GilSafeCallback(const GilSafeCallback& other) { + py::gil_scoped_acquire gil; + obj = other.obj; + } + GilSafeCallback(GilSafeCallback&& other) noexcept { + py::gil_scoped_acquire gil; + obj = std::move(other.obj); + } + GilSafeCallback& operator=(const GilSafeCallback& other) { + py::gil_scoped_acquire gil; + obj = other.obj; + return *this; + } + GilSafeCallback& operator=(GilSafeCallback&& other) noexcept { + py::gil_scoped_acquire gil; + obj = std::move(other.obj); + return *this; + } + ~GilSafeCallback() { + py::gil_scoped_acquire gil; + obj = py::object(); + } + void operator()() const { + py::gil_scoped_acquire gil; + if (obj && !obj.is_none()) { + obj(); + } + } + }; py::class_>(m, "_TestConnectionPool") .def(py::init(), py::arg("max_size") = 1, py::arg("idle_timeout_secs") = 600) .def( @@ -6150,10 +6182,7 @@ PYBIND11_MODULE(ddbc_bindings, m) { if (hook.is_none()) { pool.set_on_disconnect_hook(nullptr); } else { - pool.set_on_disconnect_hook([hook]() { - py::gil_scoped_acquire gil; - hook(); - }); + pool.set_on_disconnect_hook(GilSafeCallback(hook)); } }, py::arg("hook")) diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 9366fc03b..4818a6413 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1839,6 +1839,91 @@ def run_close(): ) +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, + ) + + # ============================================================================= # Native token-factory (lazy token acquisition) integration tests # ============================================================================= From 9eac00b571b7d162fac175ffc1e9495164f6bc62 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 17 Sep 2026 22:08:05 -0400 Subject: [PATCH 12/12] FIX: Coordinate same-key pool replacement with _closing_keys and use shared_ptr for GIL-safe hook (#746) --- .../pybind/connection/connection_pool.cpp | 290 ++++++++++-------- .../pybind/connection/connection_pool.h | 28 +- mssql_python/pybind/ddbc_bindings.cpp | 62 ++-- tests/test_009_pooling.py | 110 +++++++ 4 files changed, 329 insertions(+), 161 deletions(-) diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 1e0a4c668..f8ecad5f3 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -62,19 +62,27 @@ ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs) _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; } - std::function hook; - { - std::lock_guard lock(_mutex); - hook = _on_disconnect_hook; - } - if (hook) { - hook(); - } + invokeDisconnectHook(); try { conn->disconnect(); } catch (const std::exception& ex) { @@ -384,14 +392,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // 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). - std::function hook; - { - std::lock_guard lock(_mutex); - hook = _on_disconnect_hook; - } - if (hook) { - hook(); - } + invokeDisconnectHook(); try { new_conn->disconnect(); } catch (const std::exception& ex) { @@ -465,14 +466,7 @@ 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) { - std::function hook; - { - std::lock_guard lock(_mutex); - hook = _on_disconnect_hook; - } - if (hook) { - hook(); - } + invokeDisconnectHook(); try { conn->disconnect(); } catch (const std::exception& ex) { @@ -534,14 +528,7 @@ void ConnectionPool::close() { ++_generation; } for (auto& conn : to_close) { - std::function hook; - { - std::lock_guard lock(_mutex); - hook = _on_disconnect_hook; - } - if (hook) { - hook(); - } + invokeDisconnectHook(); try { conn->disconnect(); } catch (const std::exception& ex) { @@ -576,100 +563,154 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( std::shared_ptr pool; std::shared_ptr old_pool_to_close; bool created = false; - std::vector> evicted; - { - 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; - } - // 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; + 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(); } } - // 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); + 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(); } - 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); + } closing_guard{*this}; + + { + 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; + } } - created = true; } - pool = pool_ref; - } - } - 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()); + // 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; + } } - old_pool_to_close.reset(); - { - std::lock_guard lock(_manager_mutex); + 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); - if (_mock_mode) { - pool_ref->set_mock_mode(true); - } - 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 @@ -678,16 +719,6 @@ std::shared_ptr ConnectionPoolManager::acquireConnection( 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 @@ -737,13 +768,13 @@ void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) { } void ConnectionPoolManager::closePools() { + py::gil_scoped_release release_gil; // Under _manager_mutex, snapshot all pools to close their idle connections. - // We do not clear _pools immediately: close() disconnects ODBC handles - // (releasing the GIL), which must run outside _manager_mutex to avoid - // deadlock. + // 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) { @@ -788,4 +819,5 @@ void ConnectionPoolManager::closePools() { 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 627261d48..5bbaebb22 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -8,12 +8,14 @@ #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 @@ -81,12 +83,18 @@ class ConnectionPool { bool mock_mode() const { return _mock_mode; } - void set_on_disconnect_hook(std::function hook) { - std::lock_guard lock(_mutex); - _on_disconnect_hook = hook; + 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; @@ -95,7 +103,7 @@ class ConnectionPool { 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::function _on_disconnect_hook; + std::shared_ptr> _on_disconnect_hook; std::deque> _pool; // Available connections std::mutex _mutex; // Mutex for thread-safe access }; @@ -151,6 +159,13 @@ class ConnectionPoolManager { 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; @@ -158,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; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index ddc5a55c1..0e92f8076 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6117,6 +6117,9 @@ PYBIND11_MODULE(ddbc_bindings, m) { 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 @@ -6132,36 +6135,31 @@ PYBIND11_MODULE(ddbc_bindings, m) { py::class_>(m, "_TestPooledConnection") .def_property_readonly("origin_generation", &Connection::originGeneration) .def_property_readonly("origin_pool_id", &Connection::originPoolId); - struct GilSafeCallback { - py::object obj; - explicit GilSafeCallback(py::object o) : obj(std::move(o)) {} - GilSafeCallback(const GilSafeCallback& other) { - py::gil_scoped_acquire gil; - obj = other.obj; - } - GilSafeCallback(GilSafeCallback&& other) noexcept { - py::gil_scoped_acquire gil; - obj = std::move(other.obj); - } - GilSafeCallback& operator=(const GilSafeCallback& other) { - py::gil_scoped_acquire gil; - obj = other.obj; - return *this; - } - GilSafeCallback& operator=(GilSafeCallback&& other) noexcept { - py::gil_scoped_acquire gil; - obj = std::move(other.obj); - return *this; + 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; + } } - ~GilSafeCallback() { - py::gil_scoped_acquire gil; - obj = py::object(); + PyObjectHolder(const PyObjectHolder&) = delete; + PyObjectHolder& operator=(const PyObjectHolder&) = delete; + PyObjectHolder(PyObjectHolder&& other) noexcept : ptr(other.ptr) { + other.ptr = nullptr; } - void operator()() const { - py::gil_scoped_acquire gil; - if (obj && !obj.is_none()) { - obj(); + 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") @@ -6182,7 +6180,15 @@ PYBIND11_MODULE(ddbc_bindings, m) { if (hook.is_none()) { pool.set_on_disconnect_hook(nullptr); } else { - pool.set_on_disconnect_hook(GilSafeCallback(hook)); + 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")) diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 4818a6413..0225411eb 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1924,6 +1924,116 @@ def run_acquire(): ) +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 # =============================================================================