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