diff --git a/src/include/dbconnector/pool/connection_pool.hpp b/src/include/dbconnector/pool/connection_pool.hpp index bb9457c..1fabc17 100644 --- a/src/include/dbconnector/pool/connection_pool.hpp +++ b/src/include/dbconnector/pool/connection_pool.hpp @@ -18,45 +18,46 @@ namespace pool { template class ConnectionPool : public std::enable_shared_from_this> { public: - static constexpr size_t DEFAULT_POOL_SIZE = 4; - static constexpr size_t DEFAULT_POOL_TIMEOUT_MS = 30000; + static constexpr uint64_t DEFAULT_POOL_SIZE = 4; + static constexpr uint64_t DEFAULT_POOL_TIMEOUT_MS = 30000; enum class ThreadLocalCacheState { CACHE_ENABLED, CACHE_DISABLED, }; - ConnectionPool(size_t max_connections = DEFAULT_POOL_SIZE, size_t timeout_ms = DEFAULT_POOL_TIMEOUT_MS, + ConnectionPool(uint64_t max_connections = DEFAULT_POOL_SIZE, uint64_t wait_timeout_ms = DEFAULT_POOL_TIMEOUT_MS, ThreadLocalCacheState tl_cache_state = ThreadLocalCacheState::CACHE_DISABLED); virtual ~ConnectionPool(); - PooledConnection Acquire(); + PooledConnection WaitAcquire(); PooledConnection TryAcquire(); PooledConnection ForceAcquire(); - void Return(std::unique_ptr conn, std::chrono::steady_clock::time_point created_at); - void Discard(); + bool IsShutdown() const; void Shutdown(); - void SetMaxConnections(size_t new_max); - size_t GetMaxConnections() const; - size_t GetAvailableConnections() const; - size_t GetTotalConnections() const; - bool IsShutdown() const; + uint64_t GetMaxConnections() const; + void SetMaxConnections(uint64_t new_max); + uint64_t GetWaitTimeoutMs() const; + void SetWaitTimeoutMs(uint64_t timeout_ms); + + uint64_t GetAvailableConnections() const; + uint64_t GetTotalConnections() const; + + bool IsThreadLocalCacheEnabled() const; + void SetThreadLocalCacheEnabled(bool enabled); + uint64_t GetThreadLocalCacheHits() const; + uint64_t GetThreadLocalCacheMisses() const; - void SetMaxLifetimeSeconds(uint64_t new_max_lifetime_seconds); uint64_t GetMaxLifetimeSeconds() const; - void SetIdleTimeoutSeconds(uint64_t new_idle_timeout_seconds); + void SetMaxLifetimeSeconds(uint64_t new_max_lifetime_seconds); uint64_t GetIdleTimeoutSeconds() const; + void SetIdleTimeoutSeconds(uint64_t new_idle_timeout_seconds); bool EnsureReaperRunning(); void ShutdownReaper(); - size_t GetThreadLocalCacheHits() const; - size_t GetThreadLocalCacheMisses() const; - void SetThreadLocalCacheEnabled(bool enabled); - bool IsThreadLocalCacheEnabled() const; - protected: virtual std::unique_ptr CreateNewConnection() = 0; virtual bool CheckConnectionHealthy(ConnectionT &conn) = 0; @@ -71,6 +72,7 @@ class ConnectionPool : public std::enable_shared_from_this; friend struct ThreadLocalConnectionCache; static ThreadLocalConnectionCache &GetThreadLocalCache() { @@ -78,6 +80,9 @@ class ConnectionPool : public std::enable_shared_from_this conn, std::chrono::steady_clock::time_point created_at); + void Discard(); + bool TimeoutEnabled() const; std::chrono::steady_clock::time_point GetNowForTimeoutPurposes(); static bool TimePointExpired(std::chrono::steady_clock::time_point point, uint64_t timeout, @@ -96,15 +101,18 @@ class ConnectionPool : public std::enable_shared_from_this cached_conn); - size_t max_connections = 0; - size_t timeout_ms = 0; + void DecrementTotalConnections(); + + std::atomic max_connections {0}; + std::atomic wait_timeout_ms {0}; mutable std::mutex pool_lock; std::condition_variable pool_cv; std::deque> available; - size_t total_connections = 0; - bool shutdown_flag = false; + std::atomic available_connections {0}; + std::atomic total_connections {0}; + std::atomic shutdown_flag {false}; std::atomic max_lifetime_seconds {0}; std::atomic idle_timeout_seconds {0}; @@ -113,9 +121,9 @@ class ConnectionPool : public std::enable_shared_from_this reaper_shutdown_flag {false}; - std::atomic tl_cache_enabled; - std::atomic tl_cache_hits {0}; - std::atomic tl_cache_misses {0}; + std::atomic tl_cache_enabled {false}; + std::atomic tl_cache_hits {0}; + std::atomic tl_cache_misses {0}; }; } // namespace pool diff --git a/src/include/dbconnector/pool/connection_pool_impl.hpp b/src/include/dbconnector/pool/connection_pool_impl.hpp index d29a9f5..0ecd1a7 100644 --- a/src/include/dbconnector/pool/connection_pool_impl.hpp +++ b/src/include/dbconnector/pool/connection_pool_impl.hpp @@ -13,9 +13,9 @@ namespace dbconnector { namespace pool { template -ConnectionPool::ConnectionPool(size_t max_connections_p, size_t timeout_ms_p, +ConnectionPool::ConnectionPool(uint64_t max_connections_p, uint64_t timeout_ms_p, ThreadLocalCacheState tl_cache_state) - : max_connections(max_connections_p), timeout_ms(timeout_ms_p), total_connections(0), shutdown_flag(false), + : max_connections(max_connections_p), wait_timeout_ms(timeout_ms_p), total_connections(0), shutdown_flag(false), tl_cache_enabled(ThreadLocalCacheState::CACHE_ENABLED == tl_cache_state) { } @@ -29,11 +29,12 @@ void ConnectionPool::Shutdown() { { std::unique_lock lock(pool_lock); ShutdownReaperInternal(lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { return; } - shutdown_flag = true; + shutdown_flag.store(true, std::memory_order_relaxed); available.clear(); + available_connections.store(available.size(), std::memory_order_relaxed); } pool_cv.notify_all(); //! Other threads' TL caches self-cleanup via ThreadLocalConnectionCache destructors @@ -42,8 +43,7 @@ void ConnectionPool::Shutdown() { template bool ConnectionPool::IsShutdown() const { - std::lock_guard lock(pool_lock); - return shutdown_flag; + return shutdown_flag.load(std::memory_order_relaxed); } template @@ -104,10 +104,11 @@ bool ConnectionPool::TryReturnToThreadLocal(std::unique_ptr lock(pool_lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { return false; } - if (total_connections >= max_connections && available.empty()) { + if (total_connections.load(std::memory_order_relaxed) >= max_connections.load(std::memory_order_relaxed) && + available.empty()) { return false; } cache.cached_conn = CachedConnection(std::move(conn), created_at, returned_at); @@ -127,22 +128,21 @@ void ConnectionPool::ReturnFromThreadLocalCache(CachedConnection lock(pool_lock); - if (expired || shutdown_flag) { - if (total_connections > 0) { - total_connections--; - } + if (expired || shutdown_flag.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); return; } available.emplace_back(std::move(cached_conn)); + available_connections.store(available.size(), std::memory_order_relaxed); } pool_cv.notify_one(); } template -PooledConnection ConnectionPool::Acquire() { +PooledConnection ConnectionPool::WaitAcquire() { { std::lock_guard lock(pool_lock); - if (max_connections == 0) { + if (max_connections.load(std::memory_order_relaxed) == 0) { throw PoolException("Connection pool is disabled (pool_size=0). Use " "the force acquire mode to create connections without pooling."); } @@ -156,37 +156,35 @@ PooledConnection ConnectionPool::Acquire() { std::unique_lock lock(pool_lock); - auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(wait_timeout_ms.load(std::memory_order_relaxed)); while (true) { - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { throw PoolException("Connection pool has been shut down"); } while (!available.empty()) { auto cached_conn = std::move(available.front()); available.pop_front(); + available_connections.store(available.size(), std::memory_order_relaxed); now = GetNowForTimeoutPurposes(); bool healthy = CheckConnectionNotExpiredAndHealthy(lock, cached_conn, now); - if (shutdown_flag) { - if (total_connections > 0) { - total_connections--; - } + if (shutdown_flag.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); throw PoolException("Connection pool has been shut down"); } if (healthy) { return PooledConnection(this->shared_from_this(), std::move(cached_conn)); } - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); } - if (total_connections < max_connections) { - total_connections++; + if (total_connections.load(std::memory_order_relaxed) < max_connections.load(std::memory_order_relaxed)) { + total_connections.fetch_add(1, std::memory_order_relaxed); lock.unlock(); try { @@ -194,9 +192,7 @@ PooledConnection ConnectionPool::Acquire() { return PooledConnection(this->shared_from_this(), std::move(conn), now); } catch (...) { lock.lock(); - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); pool_cv.notify_one(); throw; } @@ -204,8 +200,10 @@ PooledConnection ConnectionPool::Acquire() { // Spurious wakeups re-evaluate all conditions above. The deadline is not reset. if (pool_cv.wait_until(lock, deadline) == std::cv_status::timeout) { - throw PoolException("Connection pool timeout: all " + std::to_string(max_connections) + - " connections in use, waited " + std::to_string(timeout_ms) + "ms"); + throw PoolException("Connection pool timeout: all " + + std::to_string(max_connections.load(std::memory_order_relaxed)) + + " connections in use, waited " + + std::to_string(wait_timeout_ms.load(std::memory_order_relaxed)) + "ms"); } } } @@ -214,7 +212,7 @@ template PooledConnection ConnectionPool::TryAcquire() { { std::lock_guard lock(pool_lock); - if (max_connections == 0) { + if (max_connections.load(std::memory_order_relaxed) == 0) { throw PoolException("Connection pool is disabled (pool_size=0). Use " "the force acquire mode to create connections without pooling."); } @@ -229,33 +227,30 @@ PooledConnection ConnectionPool::TryAcquire() { std::unique_lock lock(pool_lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { return PooledConnection(); } while (!available.empty()) { auto cached_conn = std::move(available.front()); available.pop_front(); + available_connections.store(available.size(), std::memory_order_relaxed); bool healthy = CheckConnectionNotExpiredAndHealthy(lock, cached_conn, now); - if (shutdown_flag) { - if (total_connections > 0) { - total_connections--; - } + if (shutdown_flag.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); return PooledConnection(); } if (healthy) { return PooledConnection(this->shared_from_this(), std::move(cached_conn)); } - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); } - if (total_connections < max_connections) { - total_connections++; + if (total_connections.load(std::memory_order_relaxed) < max_connections.load(std::memory_order_relaxed)) { + total_connections.fetch_add(1, std::memory_order_relaxed); lock.unlock(); try { @@ -263,9 +258,7 @@ PooledConnection ConnectionPool::TryAcquire() { return PooledConnection(this->shared_from_this(), std::move(conn), now); } catch (...) { lock.lock(); - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); pool_cv.notify_one(); throw; } @@ -288,10 +281,10 @@ PooledConnection ConnectionPool::ForceAcquire() { bool pooling_disabled = false; { std::lock_guard lock(pool_lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { throw PoolException("Connection pool has been shut down"); } - pooling_disabled = (max_connections == 0); + pooling_disabled = (max_connections.load(std::memory_order_relaxed) == 0); } if (pooling_disabled) { @@ -302,32 +295,29 @@ PooledConnection ConnectionPool::ForceAcquire() { { std::unique_lock lock(pool_lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { throw PoolException("Connection pool has been shut down"); } while (!available.empty()) { auto cached_conn = std::move(available.front()); available.pop_front(); + available_connections.store(available.size(), std::memory_order_relaxed); bool healthy = CheckConnectionNotExpiredAndHealthy(lock, cached_conn, now); - if (shutdown_flag) { - if (total_connections > 0) { - total_connections--; - } + if (shutdown_flag.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); throw PoolException("Connection pool has been shut down"); } if (healthy) { return PooledConnection(this->shared_from_this(), std::move(cached_conn)); } - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); } - total_connections++; + total_connections.fetch_add(1, std::memory_order_relaxed); } try { @@ -335,9 +325,7 @@ PooledConnection ConnectionPool::ForceAcquire() { return PooledConnection(this->shared_from_this(), std::move(conn), now); } catch (...) { std::lock_guard lock(pool_lock); - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); pool_cv.notify_one(); throw; } @@ -377,20 +365,17 @@ void ConnectionPool::Return(std::unique_ptr conn, { std::lock_guard lock(pool_lock); - if (shutdown_flag) { - if (total_connections > 0) { - total_connections--; - } + if (shutdown_flag.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); return; } - if (total_connections > max_connections) { - if (total_connections > 0) { - total_connections--; - } + if (total_connections.load(std::memory_order_relaxed) > max_connections.load(std::memory_order_relaxed)) { + DecrementTotalConnections(); return; } CachedConnection cached_conn(std::move(conn), created_at, now); available.emplace_back(std::move(cached_conn)); + available_connections.store(available.size(), std::memory_order_relaxed); } pool_cv.notify_one(); } @@ -399,44 +384,52 @@ template void ConnectionPool::Discard() { { std::lock_guard lock(pool_lock); - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); } pool_cv.notify_one(); } template -void ConnectionPool::SetMaxConnections(size_t new_max) { +uint64_t ConnectionPool::GetMaxConnections() const { + return max_connections.load(std::memory_order_relaxed); +} + +template +void ConnectionPool::SetMaxConnections(uint64_t new_max) { std::deque> to_evict; { std::lock_guard lock(pool_lock); - max_connections = new_max; - while (!available.empty() && total_connections > max_connections) { + this->max_connections.store(new_max, std::memory_order_relaxed); + while (!available.empty() && + total_connections.load(std::memory_order_relaxed) > max_connections.load(std::memory_order_relaxed)) { to_evict.push_back(std::move(available.back())); available.pop_back(); - total_connections--; + available_connections.store(available.size(), std::memory_order_relaxed); + DecrementTotalConnections(); } } pool_cv.notify_all(); } template -size_t ConnectionPool::GetMaxConnections() const { - std::lock_guard lock(pool_lock); - return max_connections; +uint64_t ConnectionPool::GetWaitTimeoutMs() const { + return wait_timeout_ms.load(std::memory_order_relaxed); } template -size_t ConnectionPool::GetAvailableConnections() const { +void ConnectionPool::SetWaitTimeoutMs(uint64_t timeout_ms) { std::lock_guard lock(pool_lock); - return available.size(); + wait_timeout_ms.store(timeout_ms, std::memory_order_relaxed); } template -size_t ConnectionPool::GetTotalConnections() const { - std::lock_guard lock(pool_lock); - return total_connections; +uint64_t ConnectionPool::GetAvailableConnections() const { + return available_connections.load(std::memory_order_relaxed); +} + +template +uint64_t ConnectionPool::GetTotalConnections() const { + return total_connections.load(std::memory_order_relaxed); } template @@ -576,13 +569,12 @@ void ConnectionPool::ReaperLoop() { TimePointExpired(cached_conn.GetReturnedAt(), idle_timeout_val, now)) { expired.emplace_back(std::move(cached_conn)); it = available.erase(it); // erase returns iterator to next element - if (total_connections > 0) { - total_connections--; - } + DecrementTotalConnections(); } else { ++it; } } + available_connections.store(available.size(), std::memory_order_relaxed); if (expired.size() > 0) { // release lock while destroying expired connections @@ -602,7 +594,7 @@ bool ConnectionPool::EnsureReaperRunning() { std::unique_lock lock(pool_lock); - if (shutdown_flag) { + if (shutdown_flag.load(std::memory_order_relaxed)) { return false; } @@ -651,12 +643,12 @@ void ConnectionPool::ShutdownReaper() { } template -size_t ConnectionPool::GetThreadLocalCacheHits() const { +uint64_t ConnectionPool::GetThreadLocalCacheHits() const { return tl_cache_hits.load(std::memory_order_relaxed); } template -size_t ConnectionPool::GetThreadLocalCacheMisses() const { +uint64_t ConnectionPool::GetThreadLocalCacheMisses() const { return tl_cache_misses.load(std::memory_order_relaxed); } @@ -679,5 +671,14 @@ void ConnectionPool::ForEachIdleConnection(Fn &&fn) { } } +template +void ConnectionPool::DecrementTotalConnections() { + // Note, this function is only called when holding the lock + // it is not supposed to be atomic. + if (total_connections.load(std::memory_order_relaxed) > 0) { + total_connections.fetch_sub(1, std::memory_order_relaxed); + } +} + } // namespace pool } // namespace dbconnector diff --git a/test/test_pool.cpp b/test/test_pool.cpp index d8b1668..cbab949 100644 --- a/test/test_pool.cpp +++ b/test/test_pool.cpp @@ -52,16 +52,29 @@ class TestConnectionPool : public dbconnector::pool::ConnectionPool(); - REQUIRE(pool->Acquire()); + REQUIRE(pool->WaitAcquire()); REQUIRE(pool->TryAcquire()); REQUIRE(pool->ForceAcquire()); + pool->SetMaxConnections(42); + REQUIRE(pool->GetMaxConnections() == 42); + pool->SetWaitTimeoutMs(43); + REQUIRE(pool->GetWaitTimeoutMs() == 43); + REQUIRE(pool->GetAvailableConnections() == 0); + REQUIRE(pool->GetTotalConnections() == 1); + REQUIRE(pool->IsThreadLocalCacheEnabled()); + pool->SetThreadLocalCacheEnabled(false); + REQUIRE(!pool->IsThreadLocalCacheEnabled()); + pool->SetMaxLifetimeSeconds(44); + REQUIRE(pool->GetMaxLifetimeSeconds() == 44); + pool->SetIdleTimeoutSeconds(45); + REQUIRE(pool->GetIdleTimeoutSeconds() == 45); } TEST_CASE("Test pool size no thread-local", group_name) { auto pool = std::make_shared(2, 500, false); { - auto conn_main = pool->Acquire(); + auto conn_main = pool->WaitAcquire(); REQUIRE(conn_main); REQUIRE(1 == pool->GetTotalConnections()); } @@ -83,7 +96,7 @@ TEST_CASE("Test pool size no thread-local", group_name) { bool timeout_thrown = false; try { - auto conn_main = pool->Acquire(); + auto conn_main = pool->WaitAcquire(); } catch (const std::exception &e) { REQUIRE(0 == std::string(e.what()).find("Connection pool timeout")); timeout_thrown = true; @@ -119,7 +132,7 @@ TEST_CASE("Test pool size with thread-local", group_name) { REQUIRE(3 == pool->GetTotalConnections()); { - auto conn_main = pool->Acquire(); + auto conn_main = pool->WaitAcquire(); REQUIRE(conn_main); REQUIRE(conn_main.GetConnection().GetId() == conn_main_id); REQUIRE(3 == pool->GetTotalConnections()); @@ -130,7 +143,7 @@ TEST_CASE("Test pool size with thread-local", group_name) { TEST_CASE("Test pool disabled", group_name) { auto pool = std::make_shared(0); - REQUIRE_THROWS(pool->Acquire()); + REQUIRE_THROWS(pool->WaitAcquire()); REQUIRE_THROWS(pool->TryAcquire()); uint64_t conn1_id = 0; @@ -172,7 +185,7 @@ TEST_CASE("Test pool disable running", group_name) { REQUIRE(pool->GetAvailableConnections() == 1); pool->SetMaxConnections(0); - REQUIRE_THROWS(pool->Acquire()); + REQUIRE_THROWS(pool->WaitAcquire()); REQUIRE_THROWS(pool->TryAcquire()); REQUIRE(pool->GetTotalConnections() == 2); @@ -219,7 +232,7 @@ TEST_CASE("Test pool with a reaper", group_name) { REQUIRE(pool->EnsureReaperRunning()); { - auto conn = pool->Acquire(); + auto conn = pool->WaitAcquire(); REQUIRE(conn); REQUIRE(1 == pool->GetTotalConnections()); } @@ -235,7 +248,7 @@ TEST_CASE("Test pool with a reaper", group_name) { pool->SetMaxLifetimeSeconds(0); { - auto conn = pool->Acquire(); + auto conn = pool->WaitAcquire(); REQUIRE(conn); REQUIRE(1 == pool->GetTotalConnections()); } @@ -256,7 +269,7 @@ TEST_CASE("Test pool with a reaper restart", group_name) { pool->ShutdownReaper(); { - auto conn = pool->Acquire(); + auto conn = pool->WaitAcquire(); REQUIRE(conn); REQUIRE(1 == pool->GetTotalConnections()); }