adding c++ support for pooling

gargsaumya · gargsaumya · commit d373e9b37061 · 2025-05-30T02:03:01.000+05:30
diff --git a/mssql_python/pybind/CMakeLists.txt b/mssql_python/pybind/CMakeLists.txt
@@ -90,7 +90,7 @@ execute_process(
 )
 
 # Add module library
-add_library(ddbc_bindings MODULE ddbc_bindings.cpp connection/connection.cpp)
+add_library(ddbc_bindings MODULE ddbc_bindings.cpp connection/connection.cpp connection/connection_pool.cpp)
 
 # Add include directories for your project
 target_include_directories(ddbc_bindings PRIVATE
diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp
@@ -5,6 +5,7 @@
 //             taken up in future
 
 #include "connection.h"
+#include "connection_pool.h"
 #include <vector>
 #include <pybind11/pybind11.h>
 
@@ -16,8 +17,8 @@ SqlHandlePtr Connection::_envHandle = nullptr;
 // This class wraps low-level ODBC operations like connect/disconnect,
 // transaction control, and autocommit configuration.
 //-------------------------------------------------------------------------------------------------
-Connection::Connection(const std::wstring& conn_str, bool autocommit)
-    : _connStr(conn_str) , _autocommit(autocommit) {
+Connection::Connection(const std::wstring& conn_str, bool use_pool)
+    : _connStr(conn_str), _autocommit(false), _fromPool(use_pool) {
     if (!_envHandle) {
         LOG("Allocating environment handle");
         SQLHANDLE env = nullptr;
@@ -64,6 +65,7 @@ void Connection::connect(const py::dict& attrs_before) {
         (SQLWCHAR*)_connStr.c_str(), SQL_NTS,
         nullptr, 0, nullptr, SQL_DRIVER_NOPROMPT);
     checkError(ret);
+    updateLastUsed();
 }
 
 void Connection::disconnect() {
@@ -91,6 +93,7 @@ void Connection::commit() {
     if (!_dbcHandle) {
         ThrowStdException("Connection handle not allocated");
     }
+    updateLastUsed();
     LOG("Committing transaction");
     SQLRETURN ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbcHandle->get(), SQL_COMMIT);
     checkError(ret);
@@ -100,6 +103,7 @@ void Connection::rollback() {
     if (!_dbcHandle) {
         ThrowStdException("Connection handle not allocated");
     }
+    updateLastUsed();
     LOG("Rolling back transaction");
     SQLRETURN ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbcHandle->get(), SQL_ROLLBACK);
     checkError(ret);
@@ -132,6 +136,7 @@ SqlHandlePtr Connection::allocStatementHandle() {
     if (!_dbcHandle) {
         ThrowStdException("Connection handle not allocated");
     }
+    updateLastUsed();
     LOG("Allocating statement handle");
     SQLHANDLE stmt = nullptr;
     SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_STMT, _dbcHandle->get(), &stmt);
@@ -185,4 +190,99 @@ void Connection::applyAttrsBefore(const py::dict& attrs) {
             }
         }
     }
-}
+}
+
+bool Connection::isAlive() const {
+    if (!_dbcHandle) {
+        ThrowStdException("Connection handle not allocated");
+    }
+    SQLUINTEGER status;
+    SQLRETURN ret = SQLGetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_CONNECTION_DEAD,
+        &status, 0, nullptr);
+    return SQL_SUCCEEDED(ret) && status == SQL_CD_FALSE;
+}
+
+bool Connection::reset() {
+    if (!_dbcHandle) {
+        ThrowStdException("Connection handle not allocated");
+    }
+    LOG("Resetting connection via SQL_ATTR_RESET_CONNECTION");
+    SQLULEN reset = SQL_TRUE;
+    SQLRETURN ret = SQLSetConnectAttr_ptr(
+        _dbcHandle->get(),
+        SQL_ATTR_RESET_CONNECTION,
+        (SQLPOINTER)SQL_RESET_CONNECTION_YES,
+        SQL_IS_INTEGER);
+    if (!SQL_SUCCEEDED(ret)) {
+        LOG("Failed to reset connection. Marking as dead.");
+        disconnect();
+        return false;
+    }
+    updateLastUsed();
+    return true;
+}
+
+void Connection::updateLastUsed() {
+    _lastUsed = std::chrono::steady_clock::now();
+}
+
+std::chrono::steady_clock::time_point Connection::lastUsed() const {
+    return _lastUsed;
+}
+
+ConnectionHandle::ConnectionHandle(const std::wstring& connStr, bool usePool, const py::dict& attrsBefore)
+    : _connStr(connStr), _usePool(usePool) {
+    if (_usePool) {
+        _conn = ConnectionPoolManager::getInstance().acquireConnection(connStr, attrsBefore);
+    } else {
+        _conn = std::make_shared<Connection>(connStr, false);
+        _conn->connect(attrsBefore);
+    }
+}
+
+void ConnectionHandle::close() {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    if (_usePool) {
+        ConnectionPoolManager::getInstance().returnConnection(_connStr, _conn);
+    } else {
+        _conn->disconnect();
+    }
+    _conn = nullptr;
+}
+
+void ConnectionHandle::commit() {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    _conn->commit();
+}
+
+void ConnectionHandle::rollback() {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    _conn->rollback();
+}
+
+void ConnectionHandle::setAutocommit(bool enabled) {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    _conn->setAutocommit(enabled);
+}
+
+bool ConnectionHandle::getAutocommit() const {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    return _conn->getAutocommit();
+}
+
+SqlHandlePtr ConnectionHandle::allocStatementHandle() {
+    if (!_conn) {
+        ThrowStdException("Connection object is not initialized");
+    }
+    return _conn->allocStatementHandle();
+}
diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h
@@ -13,7 +13,7 @@
 
 class Connection {
 public:
-    Connection(const std::wstring& conn_str, bool autocommit = false);
+    Connection(const std::wstring& connStr, bool fromPool);
     ~Connection();
 
     // Establish the connection using the stored connection string.
@@ -34,8 +34,16 @@ class Connection {
     //  Check whether autocommit is enabled.
     bool getAutocommit() const;
 
+    bool isAlive() const;
+
+    bool reset();
+
+    void updateLastUsed();
+
+    std::chrono::steady_clock::time_point lastUsed() const;
+
     // Allocate a new statement handle on this connection.
-    SqlHandlePtr allocStatementHandle();
+    SqlHandlePtr allocStatementHandle();    
 
 private:
     void allocateDbcHandle();
@@ -44,9 +52,26 @@ class Connection {
     void applyAttrsBefore(const py::dict& attrs_before);
 
     std::wstring _connStr;
-    bool _usePool = false;
+    bool _fromPool = false;
     bool _autocommit = true;
     SqlHandlePtr _dbcHandle;
-
     static SqlHandlePtr _envHandle;
+    std::chrono::steady_clock::time_point _lastUsed;
 };
+
+class ConnectionHandle {
+public:
+    ConnectionHandle(const std::wstring& connStr, bool usePool, const py::dict& attrsBefore = py::dict());
+
+    void close();
+    void commit();
+    void rollback();
+    void setAutocommit(bool enabled);
+    bool getAutocommit() const;
+    SqlHandlePtr allocStatementHandle();
+
+private:
+    std::shared_ptr<Connection> _conn;
+    bool _usePool;
+    std::wstring _connStr;
+};
diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp
@@ -0,0 +1,85 @@
+#include "connection_pool.h"
+#include <iostream>
+#include <exception>
+
+ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs)
+    : _max_size(max_size),  _idle_timeout_secs(idle_timeout_secs), _current_size(0) {}
+
+std::shared_ptr<Connection> ConnectionPool::acquire(const std::wstring& connStr, const py::dict& attrs_before) {
+    std::lock_guard<std::mutex> lock(_mutex);
+    auto now = std::chrono::steady_clock::now();
+    size_t before = _pool.size();
+    _pool.erase(std::remove_if(_pool.begin(), _pool.end(), [&](const std::shared_ptr<Connection>& conn) {
+        auto idle_time = std::chrono::duration_cast<std::chrono::seconds>(now - conn->lastUsed()).count();
+        if (idle_time > _idle_timeout_secs) {
+            conn->disconnect();
+            return true;
+        }
+        return false;
+    }), _pool.end());
+    size_t pruned = before - _pool.size();
+    _current_size = (_current_size >= pruned) ? (_current_size - pruned) : 0;
+
+    while (!_pool.empty()) {
+        auto conn = _pool.front();
+        _pool.pop_front();
+        if (conn->isAlive()) {
+            if (!conn->reset()) {
+                continue;
+            }
+            return conn;
+        } else {
+            conn->disconnect();
+            --_current_size;
+        }
+    }
+    if (_current_size < _max_size) {
+        auto conn = std::make_shared<Connection>(connStr, true);
+        conn->connect(attrs_before);
+        return conn;
+    } else {
+        LOG("Cannot acquire connection: pool size limit reached");
+        return nullptr;
+    }
+}
+
+void ConnectionPool::release(std::shared_ptr<Connection> conn) {
+    std::lock_guard<std::mutex> lock(_mutex);
+    if (_pool.size() < _max_size) {
+        conn->updateLastUsed();
+        _pool.push_back(conn);
+    }
+    else {
+        conn->disconnect();
+        if (_current_size > 0) --_current_size;
+    }
+}
+
+ConnectionPoolManager& ConnectionPoolManager::getInstance() {
+    static ConnectionPoolManager manager;
+    return manager;
+}
+
+std::shared_ptr<Connection> ConnectionPoolManager::acquireConnection(const std::wstring& connStr, const py::dict& attrs_before) {
+    std::lock_guard<std::mutex> lock(_manager_mutex);
+
+    auto& pool = _pools[connStr];
+    if (!pool) {
+        LOG("Creating new connection pool");
+        pool = std::make_shared<ConnectionPool>(_default_max_size, _default_idle_secs);
+    }
+    return pool->acquire(connStr, attrs_before);
+}
+
+void ConnectionPoolManager::returnConnection(const std::wstring& conn_str, const std::shared_ptr<Connection> conn) {
+    std::lock_guard<std::mutex> lock(_manager_mutex);
+    if (_pools.find(conn_str) != _pools.end()) {
+        _pools[conn_str]->release((conn));
+    }
+}
+
+void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) {
+    std::lock_guard<std::mutex> lock(_manager_mutex);
+    _default_max_size = max_size;
+    _default_idle_secs = idle_timeout_secs;
+}
diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h
@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+// INFO|TODO - Note that is file is Windows specific right now. Making it arch agnostic will be
+//             taken up in future.
+
+#pragma once
+#include <deque>
+#include <unordered_map>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <chrono>
+#include "connection.h"
+
+// Manages a fixed-size pool of reusable database connections for a single connection string
+class ConnectionPool {
+public:
+    ConnectionPool(size_t max_size, int idle_timeout_secs);
+
+    // Acquires a connection from the pool or creates a new one if under limit
+    std::shared_ptr<Connection> acquire(const std::wstring& connStr, const py::dict& attrs_before = py::dict());
+
+    // Returns a connection to the pool for reuse
+    void release(std::shared_ptr<Connection> conn);
+
+private:
+    size_t _max_size;       // Maximum number of connections allowed
+    int _idle_timeout_secs; // Idle time before connections are considered stale
+    size_t _current_size = 0;
+    std::deque<std::shared_ptr<Connection>> _pool;  // Available connections
+    std::mutex _mutex;      // Mutex for thread-safe access
+};
+
+// Singleton manager that handles multiple pools keyed by connection string
+class ConnectionPoolManager {
+public:
+    // Returns the singleton instance of the manager
+    static ConnectionPoolManager& getInstance();
+
+    void configure(int max_size, int idle_timeout);
+
+    // Gets a connection from the appropriate pool (creates one if none exists)
+    std::shared_ptr<Connection> acquireConnection(const std::wstring& conn_str, const py::dict& attrs_before = py::dict());
+
+    // Returns a connection to its original pool
+    void returnConnection(const std::wstring& conn_str, std::shared_ptr<Connection> conn);
+
+private:
+    ConnectionPoolManager() = default;    
+    ~ConnectionPoolManager() = default;
+
+    // Map from connection string to connection pool
+    std::unordered_map<std::wstring, std::shared_ptr<ConnectionPool>> _pools;
+
+    // Protects access to the _pools map
+    std::mutex _manager_mutex;
+    size_t _default_max_size = 10;
+    int _default_idle_secs = 300;
+    
+    // Prevent copying
+    ConnectionPoolManager(const ConnectionPoolManager&) = delete;
+    ConnectionPoolManager& operator=(const ConnectionPoolManager&) = delete;
+};
diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp

Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ execute_process(`
`90`	`90`	`)`
`91`	`91`
`92`	`92`	`# Add module library`
`93`		`-add_library(ddbc_bindings MODULE ddbc_bindings.cpp connection/connection.cpp)`
	`93`	`+add_library(ddbc_bindings MODULE ddbc_bindings.cpp connection/connection.cpp connection/connection_pool.cpp)`
`94`	`94`
`95`	`95`	`# Add include directories for your project`
`96`	`96`	`target_include_directories(ddbc_bindings PRIVATE`