transiton to libpqxx & connection pool for central controllers

controller/ConnectionPool.hpp

@@ -0,0 +1,161 @@
+/*
+ * Copyright (c)2021 ZeroTier, Inc.
+ *
+ * Use of this software is governed by the Business Source License included
+ * in the LICENSE.TXT file in the project's root directory.
+ *
+ * Change Date: 2025-01-01
+ *
+ * On the date above, in accordance with the Business Source License, use
+ * of this software will be governed by version 2.0 of the Apache License.
+ */
+/****/
+
+#ifndef ZT_CONNECTION_POOL_H_
+#define ZT_CONNECTION_POOL_H_
+
+
+#ifndef _DEBUG
+	#define _DEBUG(x)
+#endif
+
+#include <deque>
+#include <set>
+#include <memory>
+#include <mutex>
+#include <exception>
+#include <string>
+
+namespace ZeroTier {
+
+struct ConnectionUnavailable : std::exception { 
+    char const* what() const throw() {
+        return "Unable to allocate connection";
+    }; 
+};
+
+
+class Connection {
+public:
+    virtual ~Connection() {};
+};
+
+class ConnectionFactory {
+public:
+    virtual ~ConnectionFactory() {};
+    virtual std::shared_ptr<Connection> create()=0;
+};
+
+struct ConnectionPoolStats {
+    size_t pool_size;
+    size_t borrowed_size;
+};
+
+template<class T>
+class ConnectionPool {
+public:
+    ConnectionPool(size_t max_pool_size, size_t min_pool_size, std::shared_ptr<ConnectionFactory> factory)
+        : m_maxPoolSize(max_pool_size)
+        , m_minPoolSize(min_pool_size)
+        , m_factory(factory)
+    {
+        while(m_pool.size() < m_minPoolSize){
+            m_pool.push_back(m_factory->create());
+        }
+    };
+
+    ConnectionPoolStats get_stats() {
+        std::unique_lock<std::mutex> lock(m_poolMutex);
+
+        ConnectionPoolStats stats;
+        stats.pool_size = m_pool.size();
+        stats.borrowed_size = m_borrowed.size();			
+
+        return stats;
+    };
+
+    ~ConnectionPool() {
+    };
+
+    /**
+     * Borrow
+     *
+     * Borrow a connection for temporary use
+     *
+     * When done, either (a) call unborrow() to return it, or (b) (if it's bad) just let it go out of scope.  This will cause it to automatically be replaced.
+     * @retval a shared_ptr to the connection object
+     */
+    std::shared_ptr<T> borrow() {
+        std::unique_lock<std::mutex> l(m_poolMutex);
+        
+        while((m_pool.size() + m_borrowed.size()) < m_minPoolSize) {
+            std::shared_ptr<Connection> conn = m_factory->create();
+            m_pool.push_back(conn);
+        }
+
+        if(m_pool.size()==0){
+            
+            if ((m_pool.size() + m_borrowed.size()) <= m_maxPoolSize) {
+                try {
+                    std::shared_ptr<Connection> conn = m_factory->create();
+                    m_borrowed.insert(conn);
+                    return std::static_pointer_cast<T>(conn);
+                } catch (std::exception &e) {
+                    throw ConnectionUnavailable();
+                }
+            } else {
+                for(auto it = m_borrowed.begin(); it != m_borrowed.end(); ++it){
+                    if((*it).unique()) {
+                        // This connection has been abandoned! Destroy it and create a new connection
+                        try {
+                            // If we are able to create a new connection, return it
+                            _DEBUG("Creating new connection to replace discarded connection");
+                            std::shared_ptr<Connection> conn = m_factory->create();
+                            m_borrowed.erase(it);
+                            m_borrowed.insert(conn);
+                            return std::static_pointer_cast<T>(conn);
+                        } catch(std::exception& e) {
+                            // Error creating a replacement connection
+                            throw ConnectionUnavailable();
+                        }
+                    }
+                }
+                // Nothing available
+                throw ConnectionUnavailable();
+            }
+        }
+
+        // Take one off the front
+        std::shared_ptr<Connection> conn = m_pool.front();
+        m_pool.pop_front();
+        // Add it to the borrowed list
+        m_borrowed.insert(conn);
+        return std::static_pointer_cast<T>(conn);
+    };
+
+    /**
+     * Unborrow a connection
+     *
+     * Only call this if you are returning a working connection.  If the connection was bad, just let it go out of scope (so the connection manager can replace it).
+     * @param the connection
+     */
+    void unborrow(std::shared_ptr<T> conn) {
+        // Lock
+        std::unique_lock<std::mutex> lock(m_poolMutex);
+        m_borrowed.erase(conn);
+        if ((m_pool.size() + m_borrowed.size()) < m_maxPoolSize) {
+            m_pool.push_back(conn);
+        }
+    };
+protected:
+    size_t m_maxPoolSize;
+    size_t m_minPoolSize;
+    std::shared_ptr<ConnectionFactory> m_factory;
+    std::deque<std::shared_ptr<Connection> > m_pool;
+    std::set<std::shared_ptr<Connection> > m_borrowed;
+    std::mutex m_poolMutex;
+};
+
+}
+
+#endif