4 changed files with 584 additions and 999 deletions
@ -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 |
||||||
Loading…
Reference in new issue