Skip to content

Commit fb0e7c1

Browse files
authored
FIX: Ensure graceful cleanup of connection pooling to avoid application hang (#120)
### ADO Work Item Reference <!-- Insert your ADO Work Item ID below (e.g. AB#37452) --> > [AB#37851](https://sqlclientdrivers.visualstudio.com/c6d89619-62de-46a0-8b46-70b92a84d85e/_workitems/edit/37851) ------------------------------------------------------------------- ### Summary <!-- Insert your Copilot Generated Summary below --> This pull request introduces enhancements to the connection pooling system in the `mssql_python` module, focusing on resource cleanup during application shutdown. The key changes include adding functionality to close all connection pools gracefully and ensuring pooling is disabled automatically when the program exits. ### Enhancements to connection pooling: * **Automatic resource cleanup on shutdown**: - Added an `atexit` hook in `mssql_python/pooling.py` to call `ddbc_bindings.close_pooling()` when the program exits, ensuring connection pools are closed properly. * **New `close` method for `ConnectionPool`**: - Implemented a method in `connection_pool.cpp` to close all connections in a pool, releasing resources safely. Connections are disconnected in a loop, and exceptions during disconnection are logged. - Declared the new `close` method in `connection_pool.h`. * **New `closePools` method for `ConnectionPoolManager`**: - Added functionality in `connection_pool.cpp` to iterate through all pools managed by `ConnectionPoolManager` and close them. The `_pools` map is cleared afterward. - Declared the `closePools` method in `connection_pool.h`. ### Integration with Python bindings: * **Python binding for `close_pooling`**: - Added a new Python-accessible function `close_pooling` in `ddbc_bindings.cpp`, which invokes `ConnectionPoolManager::closePools()` to close all connection pools from the Python layer. These changes improve the robustness of the connection pooling system by ensuring resources are properly released during shutdown, preventing potential memory leaks or dangling connections. <!-- ### PR Title Guide > For feature requests FEAT: (short-description) > For non-feature requests like test case updates, config updates , dependency updates etc CHORE: (short-description) > For Fix requests FIX: (short-description) > For doc update requests DOC: (short-description) > For Formatting, indentation, or styling update STYLE: (short-description) > For Refactor, without any feature changes REFACTOR: (short-description) > For release related changes, without any feature changes RELEASE: #<RELEASE_VERSION> (short-description) -->
1 parent 94d5387 commit fb0e7c1

File tree

4 files changed

+43
-1
lines changed

4 files changed

+43
-1
lines changed

mssql_python/pooling.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# mssql_python/pooling.py
2+
import atexit
23
from mssql_python import ddbc_bindings
34
import threading
45

@@ -27,3 +28,8 @@ def enable(cls, max_size=100, idle_timeout=600):
2728
@classmethod
2829
def is_enabled(cls):
2930
return cls._enabled
31+
32+
@atexit.register
33+
def shutdown_pooling():
34+
if PoolingManager.is_enabled():
35+
ddbc_bindings.close_pooling()

mssql_python/pybind/connection/connection_pool.cpp

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,25 @@ void ConnectionPool::release(std::shared_ptr<Connection> conn) {
8383
}
8484
}
8585

86+
void ConnectionPool::close() {
87+
std::vector<std::shared_ptr<Connection>> to_close;
88+
{
89+
std::lock_guard<std::mutex> lock(_mutex);
90+
while (!_pool.empty()) {
91+
to_close.push_back(_pool.front());
92+
_pool.pop_front();
93+
}
94+
_current_size = 0;
95+
}
96+
for (auto& conn : to_close) {
97+
try {
98+
conn->disconnect();
99+
} catch (const std::exception& ex) {
100+
LOG("ConnectionPool::close: disconnect failed: {}", ex.what());
101+
}
102+
}
103+
}
104+
86105
ConnectionPoolManager& ConnectionPoolManager::getInstance() {
87106
static ConnectionPoolManager manager;
88107
return manager;
@@ -110,4 +129,14 @@ void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) {
110129
std::lock_guard<std::mutex> lock(_manager_mutex);
111130
_default_max_size = max_size;
112131
_default_idle_secs = idle_timeout_secs;
113-
}
132+
}
133+
134+
void ConnectionPoolManager::closePools() {
135+
std::lock_guard<std::mutex> lock(_manager_mutex);
136+
for (auto& [conn_str, pool] : _pools) {
137+
if (pool) {
138+
pool->close();
139+
}
140+
}
141+
_pools.clear();
142+
}

mssql_python/pybind/connection/connection_pool.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ class ConnectionPool {
2424
// Returns a connection to the pool for reuse
2525
void release(std::shared_ptr<Connection> conn);
2626

27+
// Closes all connections in the pool, releasing resources
28+
void close();
29+
2730
private:
2831
size_t _max_size; // Maximum number of connections allowed
2932
int _idle_timeout_secs; // Idle time before connections are considered stale
@@ -46,6 +49,9 @@ class ConnectionPoolManager {
4649
// Returns a connection to its original pool
4750
void returnConnection(const std::wstring& conn_str, std::shared_ptr<Connection> conn);
4851

52+
// Closes all pools and their connections
53+
void closePools();
54+
4955
private:
5056
ConnectionPoolManager() = default;
5157
~ConnectionPoolManager() = default;

mssql_python/pybind/ddbc_bindings.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2109,6 +2109,7 @@ PYBIND11_MODULE(ddbc_bindings, m) {
21092109
.def("get_autocommit", &ConnectionHandle::getAutocommit)
21102110
.def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle);
21112111
m.def("enable_pooling", &enable_pooling, "Enable global connection pooling");
2112+
m.def("close_pooling", []() {ConnectionPoolManager::getInstance().closePools();});
21122113
m.def("DDBCSQLExecDirect", &SQLExecDirect_wrap, "Execute a SQL query directly");
21132114
m.def("DDBCSQLExecute", &SQLExecute_wrap, "Prepare and execute T-SQL statements");
21142115
m.def("DDBCSQLRowCount", &SQLRowCount_wrap,

0 commit comments

Comments
 (0)