From edddb774b3e8c7ca9876b04897dd10c390771135 Mon Sep 17 00:00:00 2001 From: Ajay Y Date: Thu, 27 Aug 2026 10:21:32 +0000 Subject: [PATCH] feat(storage): separate read and hedging thread pools - Extract lazy, dynamically scaling ThreadPool primitive from HedgingThreadPool. - Separate StorageConnectionImpl thread pool into a dedicated ReadThreadPool (for primary stream opens) and a HedgingThreadPool (for speculative secondary hedges). - Add ReadThreadPoolSizeOption and HedgingThreadPoolSizeOption with auto-scaling defaults to prevent read bottlenecking under high concurrency. - Extract DefaultReadThreadPoolSize() and DefaultHedgingThreadPoolSize() helpers to share sizing logic between DefaultOptions() and connection initialization. - Enqueue primary read attempt to ReadThreadPool and speculative hedge attempts to HedgingThreadPool, ensuring complete fault and stall isolation. - Clamp ThreadPool capacity to at least 1 to prevent deadlock on zero sizing. - Add unit tests verifying thread pool execution, default sizes, zero-size handling, lazy spawning, and pool isolation under saturation. --- .../cloud/storage/internal/connection_impl.cc | 36 ++-- .../cloud/storage/internal/connection_impl.h | 1 + .../internal/hedged_object_read_source.cc | 13 +- .../internal/hedged_object_read_source.h | 4 +- .../hedged_object_read_source_test.cc | 187 +++++++++++++---- .../storage/internal/hedging_thread_pool.h | 189 +++++++++++++----- .../internal/hedging_thread_pool_test.cc | 132 ++++++++---- google/cloud/storage/options.h | 28 +++ 8 files changed, 438 insertions(+), 152 deletions(-) diff --git a/google/cloud/storage/internal/connection_impl.cc b/google/cloud/storage/internal/connection_impl.cc index 41c9e1297c639..bac46cfd22f53 100644 --- a/google/cloud/storage/internal/connection_impl.cc +++ b/google/cloud/storage/internal/connection_impl.cc @@ -159,23 +159,29 @@ StorageConnectionImpl::StorageConnectionImpl( : stub_(std::move(stub)), options_(MergeOptions(std::move(options), stub_->options())) { if (options_.get()) { - // The pool only runs stream-open attempts: one primary and (at most) a few - // hedges per stream being opened. Size it to the number of connections the - // REST layer can use, falling back to the hardware concurrency when the - // connection pool is unbounded (`ConnectionPoolSizeOption == 0`). - auto pool_size = options_.get(); - if (pool_size == 0) { - pool_size = - (std::max)(4, std::thread::hardware_concurrency()); + // `DefaultOptions()` normally resolves these, but a connection can be + // built without it, in which case the option is left at 0 ("automatic"). + // A pool sized 0 would accept reads it never runs, hanging the caller. + std::size_t read_threads = + options_.get(); + if (read_threads == 0) read_threads = DefaultReadThreadPoolSize(); + // The read pool only ever has one thread per in-flight application read, + // and each one blocks inside a synchronous read. Once it saturates, new + // primaries queue behind blocked ones and reads degrade to hedge-only. + read_pool_ = std::make_shared(read_threads); + + std::int64_t const max_concurrent = + options_.get(); + std::size_t hedge_threads = + options_.get(); + if (hedge_threads == 0) { + hedge_threads = DefaultHedgingThreadPoolSize(max_concurrent); } - auto const max_threads = 2 * pool_size; - auto const rate_limit = + double const rate_limit = options_.get(); - auto const max_concurrent = - options_.get(); // Allow bursts of up to one second worth of hedges. hedge_pool_ = std::make_shared( - max_threads, rate_limit, rate_limit, max_concurrent); + hedge_threads, rate_limit, rate_limit, max_concurrent); } } @@ -435,14 +441,14 @@ StatusOr> StorageConnectionImpl::ReadObject( auto const max_buffer = current->get(); - if (!enable_hedging || max_hedges <= 0 || !hedge_pool_) { + if (!enable_hedging || max_hedges <= 0 || !hedge_pool_ || !read_pool_) { return retry_source_factory(); } // `max_buffer` bounds the size of an individual read, which is only known // when the application calls `Read()`; the source applies it there. return std::unique_ptr( - std::make_unique(hedge_pool_, + std::make_unique(read_pool_, hedge_pool_, std::move(retry_source_factory), delay, max_hedges, max_buffer)); } diff --git a/google/cloud/storage/internal/connection_impl.h b/google/cloud/storage/internal/connection_impl.h index b1e2b36ae7cc7..895e943a93227 100644 --- a/google/cloud/storage/internal/connection_impl.h +++ b/google/cloud/storage/internal/connection_impl.h @@ -188,6 +188,7 @@ class StorageConnectionImpl std::unique_ptr stub_; Options options_; + std::shared_ptr read_pool_; std::shared_ptr hedge_pool_; google::cloud::internal::InvocationIdGenerator invocation_id_generator_; }; diff --git a/google/cloud/storage/internal/hedged_object_read_source.cc b/google/cloud/storage/internal/hedged_object_read_source.cc index ebd0a8441c2a1..7d1d87e18eac0 100644 --- a/google/cloud/storage/internal/hedged_object_read_source.cc +++ b/google/cloud/storage/internal/hedged_object_read_source.cc @@ -55,7 +55,7 @@ void RunAttempt(std::shared_ptr const& state, auto source = factory(); if (!source) { if (!resolve_on_open_error) return; - auto expected = false; + bool expected = false; if (state->resolved.compare_exchange_strong(expected, true)) { state->promise.set_value( RaceResult{std::move(source).status(), nullptr, {}}); @@ -65,7 +65,7 @@ void RunAttempt(std::shared_ptr const& state, std::unique_ptr buffer(new (std::nothrow) char[n]); if (!buffer) { if (!resolve_on_open_error) return; - auto expected = false; + bool expected = false; if (state->resolved.compare_exchange_strong(expected, true)) { state->promise.set_value(RaceResult{ google::cloud::internal::ResourceExhaustedError( @@ -76,7 +76,7 @@ void RunAttempt(std::shared_ptr const& state, return; } auto result = (*source)->Read(buffer.get(), n); - auto expected = false; + bool expected = false; if (state->resolved.compare_exchange_strong(expected, true)) { state->promise.set_value( RaceResult{std::move(result), *std::move(source), std::move(buffer)}); @@ -88,9 +88,11 @@ void RunAttempt(std::shared_ptr const& state, } // namespace HedgedObjectReadSource::HedgedObjectReadSource( + std::shared_ptr read_pool, std::shared_ptr hedge_pool, ChildFactory child_factory, std::chrono::milliseconds delay, int max_hedges, std::size_t max_buffer) - : hedge_pool_(std::move(hedge_pool)), + : read_pool_(std::move(read_pool)), + hedge_pool_(std::move(hedge_pool)), child_factory_(std::move(child_factory)), delay_(delay), max_hedges_(max_hedges), @@ -135,9 +137,10 @@ StatusOr HedgedObjectReadSource::Read(char* buf, auto primary = [state, factory = child_factory_, n] { RunAttempt(state, factory, n, /*resolve_on_open_error=*/true, nullptr); }; + // The primary attempt is scheduled on the dedicated read pool. // If the pool is shutting down run the attempt inline, the read must // complete either way. - if (!hedge_pool_->Enqueue(primary)) primary(); + if (!read_pool_->Enqueue(primary)) primary(); for (int i = 0; i != max_hedges_; ++i) { if (future.wait_for(delay_) != std::future_status::timeout) break; diff --git a/google/cloud/storage/internal/hedged_object_read_source.h b/google/cloud/storage/internal/hedged_object_read_source.h index b7c8930fcfc7b..c224b3ecedafb 100644 --- a/google/cloud/storage/internal/hedged_object_read_source.h +++ b/google/cloud/storage/internal/hedged_object_read_source.h @@ -54,7 +54,8 @@ class HedgedObjectReadSource : public ObjectReadSource { using ChildFactory = std::function>()>; - HedgedObjectReadSource(std::shared_ptr hedge_pool, + HedgedObjectReadSource(std::shared_ptr read_pool, + std::shared_ptr hedge_pool, ChildFactory child_factory, std::chrono::milliseconds delay, int max_hedges, std::size_t max_buffer); @@ -66,6 +67,7 @@ class HedgedObjectReadSource : public ObjectReadSource { StatusOr Read(char* buf, std::size_t n) override; private: + std::shared_ptr read_pool_; std::shared_ptr hedge_pool_; ChildFactory child_factory_; std::chrono::milliseconds delay_; diff --git a/google/cloud/storage/internal/hedged_object_read_source_test.cc b/google/cloud/storage/internal/hedged_object_read_source_test.cc index ab28f2e1a2dae..aee4e67d187c9 100644 --- a/google/cloud/storage/internal/hedged_object_read_source_test.cc +++ b/google/cloud/storage/internal/hedged_object_read_source_test.cc @@ -36,23 +36,53 @@ using ::google::cloud::storage::testing::MockObjectReadSource; using ::google::cloud::testing_util::IsOk; using ::google::cloud::testing_util::StatusIs; using ::testing::Eq; -using ::testing::Return; // Large enough that no test read is treated as oversized. -auto constexpr kUnlimitedBuffer = std::size_t{1} << 30; +std::size_t constexpr kUnlimitedBuffer = std::size_t{1} << 30; -std::shared_ptr MakeUnlimitedPool() { +std::shared_ptr MakeUnlimitedReadPool() { + return std::make_shared(/*max_threads=*/4); +} + +std::shared_ptr MakeUnlimitedHedgePool() { return std::make_shared( /*max_threads=*/4, /*rate_limit=*/0.0, /*capacity=*/0.0, /*max_concurrent=*/0); } ReadSourceResult MakeReadResult(std::string const& payload) { - auto result = - ReadSourceResult{payload.size(), HttpResponse{HttpStatusCode::kOk, - /*payload=*/{}, - /*headers=*/{}}}; - return result; + return ReadSourceResult{payload.size(), + HttpResponse{HttpStatusCode::kOk, {}, {}}}; +} + +auto MakeStallingPrimaryFactory( + std::shared_ptr> const& unblock_primary, + std::shared_ptr> const& primary_closed, + std::shared_ptr> const& calls) { + return [unblock_primary, primary_closed, + calls]() -> StatusOr> { + auto mock = std::make_unique(); + if (++*calls == 1) { + EXPECT_CALL(*mock, Read) + .WillOnce([unblock_primary](char* buf, std::size_t) { + unblock_primary->get_future().get(); + std::string const payload = "slow"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }); + EXPECT_CALL(*mock, Close).WillOnce([primary_closed]() { + primary_closed->set_value(); + return make_status_or(HttpResponse{HttpStatusCode::kOk, {}, {}}); + }); + } else { + EXPECT_CALL(*mock, Read).WillOnce([](char* buf, std::size_t) { + std::string const payload = "hedge"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }); + } + return std::unique_ptr(std::move(mock)); + }; } TEST(HedgedObjectReadSourceTest, PrimaryWins) { @@ -66,7 +96,8 @@ TEST(HedgedObjectReadSourceTest, PrimaryWins) { return std::unique_ptr(std::move(mock)); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(500), /*max_hedges=*/2, kUnlimitedBuffer); @@ -88,12 +119,21 @@ TEST(HedgedObjectReadSourceTest, SubsequentReadsContinueOnWinner) { ++*factory_calls; auto mock = std::make_unique(); EXPECT_CALL(*mock, Read) - .WillOnce(Return(MakeReadResult("chunk-1"))) - .WillOnce(Return(MakeReadResult("chunk-2"))); + .WillOnce([](char* buf, std::size_t) { + std::string const payload = "chunk-1"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }) + .WillOnce([](char* buf, std::size_t) { + std::string const payload = "chunk-2"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }); return std::unique_ptr(std::move(mock)); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(500), /*max_hedges=*/2, kUnlimitedBuffer); @@ -110,43 +150,92 @@ TEST(HedgedObjectReadSourceTest, HedgeWinsWhenPrimaryStalls) { auto unblock_primary = std::make_shared>(); auto primary_closed = std::make_shared>(); auto calls = std::make_shared>(0); - auto factory = [unblock_primary, primary_closed, - calls]() -> StatusOr> { - auto mock = std::make_unique(); - if (++*calls == 1) { - EXPECT_CALL(*mock, Read).WillOnce([unblock_primary](char*, std::size_t) { - unblock_primary->get_future().get(); - return MakeReadResult("slow"); - }); - EXPECT_CALL(*mock, Close).WillOnce([primary_closed]() { - primary_closed->set_value(); - return make_status_or(HttpResponse{HttpStatusCode::kOk, {}, {}}); - }); - } else { - EXPECT_CALL(*mock, Read).WillOnce(Return(MakeReadResult("hedge"))); - } - return std::unique_ptr(std::move(mock)); - }; + auto factory = + MakeStallingPrimaryFactory(unblock_primary, primary_closed, calls); auto source = std::make_unique( - MakeUnlimitedPool(), factory, std::chrono::milliseconds(1), + MakeUnlimitedReadPool(), MakeUnlimitedHedgePool(), factory, + std::chrono::milliseconds(1), /*max_hedges=*/2, kUnlimitedBuffer); std::vector buffer(100); auto result = source->Read(buffer.data(), buffer.size()); ASSERT_THAT(result, IsOk()); EXPECT_THAT(result->bytes_received, Eq(5)); + EXPECT_THAT(std::string(buffer.data(), result->bytes_received), Eq("hedge")); + + unblock_primary->set_value(); + primary_closed->get_future().get(); +} + +TEST(HedgedObjectReadSourceTest, ReadPoolSaturationDoesNotBlockHedges) { + // Verify thread pool isolation: If the read pool is busy with slow reads, + // speculative hedge attempts on hedge_pool_ can still execute immediately. + auto unblock_primary = std::make_shared>(); + auto primary_closed = std::make_shared>(); + auto calls = std::make_shared>(0); + auto factory = + MakeStallingPrimaryFactory(unblock_primary, primary_closed, calls); + + HedgedObjectReadSource source(std::make_shared(/*max_threads=*/1), + MakeUnlimitedHedgePool(), factory, + std::chrono::milliseconds(1), + /*max_hedges=*/2, kUnlimitedBuffer); + + std::vector buffer(100); + auto result = source.Read(buffer.data(), buffer.size()); + ASSERT_THAT(result, IsOk()); + EXPECT_THAT(result->bytes_received, Eq(5)); + EXPECT_THAT(std::string(buffer.data(), result->bytes_received), Eq("hedge")); unblock_primary->set_value(); primary_closed->get_future().get(); } +TEST(HedgedObjectReadSourceTest, HedgePoolExhaustionDoesNotBlockPrimary) { + // Verify that if the hedge pool is fully exhausted / rate limited (0 tokens), + // the primary attempt on read_pool still completes successfully. + auto read_pool = MakeUnlimitedReadPool(); + auto hedge_pool = std::make_shared( + /*max_threads=*/1, /*rate_limit=*/0.0, /*capacity=*/0.0, + /*max_concurrent=*/1); + // Acquire the only slot so hedge pool has 0 available capacity. + ASSERT_TRUE(hedge_pool->TryAcquireHedgeToken()); + + auto calls = std::make_shared>(0); + auto factory = [calls]() -> StatusOr> { + ++*calls; + auto mock = std::make_unique(); + EXPECT_CALL(*mock, Read).WillOnce([](char* buf, std::size_t) { + std::string const payload = "primary_only"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }); + return std::unique_ptr(std::move(mock)); + }; + + HedgedObjectReadSource source(read_pool, hedge_pool, factory, + std::chrono::milliseconds(10), + /*max_hedges=*/2, kUnlimitedBuffer); + + std::vector buffer(100); + auto result = source.Read(buffer.data(), buffer.size()); + ASSERT_THAT(result, IsOk()); + EXPECT_THAT(result->bytes_received, Eq(12)); + EXPECT_THAT(std::string(buffer.data(), result->bytes_received), + Eq("primary_only")); + EXPECT_THAT(calls->load(), Eq(1)); + + hedge_pool->ReleaseHedgeSlot(); +} + TEST(HedgedObjectReadSourceTest, PrimaryOpenErrorPropagates) { auto factory = []() -> StatusOr> { return Status(StatusCode::kPermissionDenied, "uh-oh"); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(500), /*max_hedges=*/2, kUnlimitedBuffer); @@ -159,7 +248,8 @@ TEST(HedgedObjectReadSourceTest, CloseWithoutReadSucceeds) { auto factory = []() -> StatusOr> { return Status(StatusCode::kUnimplemented, "never called"); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(500), /*max_hedges=*/2, kUnlimitedBuffer); EXPECT_TRUE(source.IsOpen()); @@ -167,19 +257,21 @@ TEST(HedgedObjectReadSourceTest, CloseWithoutReadSucceeds) { } TEST(HedgedObjectReadSourceTest, CloseBeforeRead) { - auto pool = std::make_shared(1, 0.0, 0.0, 0); + auto read_pool = std::make_shared(1); + auto hedge_pool = std::make_shared(1, 0.0, 0.0, 0); auto factory = []() { return std::unique_ptr( std::make_unique()); }; - HedgedObjectReadSource source(pool, factory, std::chrono::milliseconds(10), 2, + HedgedObjectReadSource source(read_pool, hedge_pool, factory, + std::chrono::milliseconds(10), 2, kUnlimitedBuffer); EXPECT_TRUE(source.IsOpen()); - EXPECT_STATUS_OK(source.Close()); + EXPECT_THAT(source.Close(), IsOk()); EXPECT_FALSE(source.IsOpen()); auto const res = source.Read(nullptr, 1024); - EXPECT_TRUE(res.ok()); - EXPECT_EQ(res->bytes_received, 0); + EXPECT_THAT(res, IsOk()); + EXPECT_THAT(res->bytes_received, Eq(0)); } TEST(HedgedObjectReadSourceTest, OversizedReadIsNotHedged) { @@ -199,7 +291,8 @@ TEST(HedgedObjectReadSourceTest, OversizedReadIsNotHedged) { // A zero delay would let a hedge start immediately if the limit were not // honored, so any race would be observable as extra factory calls. - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(0), /*max_hedges=*/2, /*max_buffer=*/8); @@ -216,7 +309,8 @@ TEST(HedgedObjectReadSourceTest, OversizedReadPropagatesOpenError) { return Status(StatusCode::kPermissionDenied, "uh-oh"); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(0), /*max_hedges=*/2, /*max_buffer=*/8); @@ -233,12 +327,21 @@ TEST(HedgedObjectReadSourceTest, SubsequentReadsIgnoreBufferLimit) { ++*calls; auto mock = std::make_unique(); EXPECT_CALL(*mock, Read) - .WillOnce(Return(MakeReadResult("small"))) - .WillOnce(Return(MakeReadResult("large"))); + .WillOnce([](char* buf, std::size_t) { + std::string const payload = "small"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }) + .WillOnce([](char* buf, std::size_t) { + std::string const payload = "large"; + std::copy(payload.begin(), payload.end(), buf); + return MakeReadResult(payload); + }); return std::unique_ptr(std::move(mock)); }; - HedgedObjectReadSource source(MakeUnlimitedPool(), factory, + HedgedObjectReadSource source(MakeUnlimitedReadPool(), + MakeUnlimitedHedgePool(), factory, std::chrono::milliseconds(500), /*max_hedges=*/2, /*max_buffer=*/64); diff --git a/google/cloud/storage/internal/hedging_thread_pool.h b/google/cloud/storage/internal/hedging_thread_pool.h index c6ef44dc32d58..1e5b1e8cadb3a 100644 --- a/google/cloud/storage/internal/hedging_thread_pool.h +++ b/google/cloud/storage/internal/hedging_thread_pool.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -35,14 +36,17 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN namespace internal { /** - * A lazy, dynamically-scaling thread pool with integrated hedge throttling. + * A lazy, dynamically-scaling thread pool for asynchronous background tasks. * - * The pool starts with no threads and spawns workers on demand, up to - * @p max_threads. Hedged requests are gated by `TryAcquireHedgeToken()`, - * which enforces two limits: a maximum number of concurrently active hedges, - * and a maximum rate of new hedges per second (a token bucket). + * The pool starts with no threads and spawns workers on demand up to + * @p max_threads. Idle workers wait on a condition variable until tasks arrive + * or the pool is shut down. + * + * @p max_threads is clamped to at least 1. A pool that can never spawn a + * worker would accept tasks it silently never runs, and callers that block on + * a task's side effects would wait forever. */ -class HedgingThreadPool { +class ThreadPool { private: struct State { std::size_t const max_threads; @@ -52,24 +56,15 @@ class HedgingThreadPool { std::condition_variable cv; bool stop = false; - // Concurrency limiter. - std::int64_t const max_concurrent_hedges; - std::atomic active_concurrent_hedges{0}; - - explicit State(std::size_t mt, std::int64_t mc) - : max_threads(mt), max_concurrent_hedges(mc) {} + explicit State(std::size_t mt) : max_threads(mt) {} }; public: - HedgingThreadPool(std::size_t max_threads, double rate_limit, double capacity, - std::int64_t max_concurrent) - : state_(std::make_shared(max_threads, max_concurrent)), - rate_limit_(rate_limit), - tokens_capacity_((std::max)(1.0, capacity)), - tokens_((std::max)(1.0, capacity)), - last_refill_(std::chrono::steady_clock::now()) {} + explicit ThreadPool(std::size_t max_threads) + : state_( + std::make_shared((std::max)(1, max_threads))) {} - ~HedgingThreadPool() { + ~ThreadPool() { { std::lock_guard lock(state_->queue_mutex); state_->stop = true; @@ -86,6 +81,11 @@ class HedgingThreadPool { } } + ThreadPool(ThreadPool const&) = delete; + ThreadPool& operator=(ThreadPool const&) = delete; + ThreadPool(ThreadPool&&) = delete; + ThreadPool& operator=(ThreadPool&&) = delete; + /** * Schedule @p task to run on a pool thread. * @@ -108,6 +108,67 @@ class HedgingThreadPool { return true; } + std::size_t max_threads() const { return state_->max_threads; } + + private: + void SpawnWorker() { + workers_.emplace_back([state = state_]() mutable { + while (true) { + std::function task; + { + std::unique_lock lock(state->queue_mutex); + ++state->idle_threads; + state->cv.wait( + lock, [&state] { return state->stop || !state->tasks.empty(); }); + --state->idle_threads; + if (state->stop && state->tasks.empty()) return; + task = std::move(state->tasks.front()); + state->tasks.pop(); + } + task(); + } + }); + } + + std::shared_ptr state_; + std::vector workers_; +}; + +/** + * A dedicated thread pool with integrated hedge throttling. + * + * Hedged requests are gated by `TryAcquireHedgeToken()`, which enforces two + * limits: a maximum number of concurrently active hedges, and a maximum rate of + * new hedges per second (a token bucket). Task execution is dispatched onto a + * dedicated internal `ThreadPool`. + */ +class HedgingThreadPool { + public: + HedgingThreadPool(std::size_t max_threads, double rate_limit, double capacity, + std::int64_t max_concurrent) + : rate_limit_(rate_limit), + tokens_capacity_((std::max)(1.0, capacity)), + tokens_((std::max)(1.0, capacity)), + last_refill_(std::chrono::steady_clock::now()), + max_concurrent_hedges_(max_concurrent), + pool_(max_threads) {} + + ~HedgingThreadPool() = default; + + HedgingThreadPool(HedgingThreadPool const&) = delete; + HedgingThreadPool& operator=(HedgingThreadPool const&) = delete; + HedgingThreadPool(HedgingThreadPool&&) = delete; + HedgingThreadPool& operator=(HedgingThreadPool&&) = delete; + + /** + * Schedule @p task to run on a pool thread. + * + * Returns false if the pool is shutting down. + */ + bool Enqueue(std::function task) { + return pool_.Enqueue(std::move(task)); + } + /** * Try to reserve capacity for one hedged request. * @@ -115,12 +176,12 @@ class HedgingThreadPool { */ bool TryAcquireHedgeToken() { // Gate 1: the ceiling on concurrently active hedges. - if (state_->max_concurrent_hedges > 0) { - auto current = - state_->active_concurrent_hedges.load(std::memory_order_relaxed); + if (max_concurrent_hedges_ > 0) { + std::int64_t current = + active_concurrent_hedges_.load(std::memory_order_relaxed); do { - if (current >= state_->max_concurrent_hedges) return false; - } while (!state_->active_concurrent_hedges.compare_exchange_weak( + if (current >= max_concurrent_hedges_) return false; + } while (!active_concurrent_hedges_.compare_exchange_weak( current, current + 1, std::memory_order_relaxed)); } @@ -129,10 +190,7 @@ class HedgingThreadPool { std::lock_guard lock(limiter_mutex_); Refill(); if (tokens_ < 1.0) { - if (state_->max_concurrent_hedges > 0) { - state_->active_concurrent_hedges.fetch_sub(1, - std::memory_order_relaxed); - } + ReleaseHedgeSlot(); return false; } tokens_ -= 1.0; @@ -142,34 +200,18 @@ class HedgingThreadPool { } void ReleaseHedgeSlot() { - if (state_->max_concurrent_hedges > 0) { - state_->active_concurrent_hedges.fetch_sub(1, std::memory_order_relaxed); + if (max_concurrent_hedges_ > 0) { + active_concurrent_hedges_.fetch_sub(1, std::memory_order_relaxed); } } - private: - void SpawnWorker() { - workers_.emplace_back([state = state_]() mutable { - while (true) { - std::function task; - { - std::unique_lock lock(state->queue_mutex); - ++state->idle_threads; - state->cv.wait( - lock, [&state] { return state->stop || !state->tasks.empty(); }); - --state->idle_threads; - if (state->stop && state->tasks.empty()) return; - task = std::move(state->tasks.front()); - state->tasks.pop(); - } - task(); - } - }); - } + std::size_t max_threads() const { return pool_.max_threads(); } + private: void Refill() { - auto now = std::chrono::steady_clock::now(); - auto const elapsed = + std::chrono::steady_clock::time_point const now = + std::chrono::steady_clock::now(); + double const elapsed = std::chrono::duration_cast>(now - last_refill_) .count(); @@ -177,17 +219,54 @@ class HedgingThreadPool { tokens_ = (std::min)(tokens_capacity_, tokens_ + elapsed * rate_limit_); } - std::shared_ptr state_; - std::vector workers_; - // Token bucket rate limiter. double rate_limit_; double tokens_capacity_; double tokens_; std::chrono::steady_clock::time_point last_refill_; std::mutex limiter_mutex_; + + // Concurrency limiter. + std::int64_t const max_concurrent_hedges_; + std::atomic active_concurrent_hedges_{0}; + + // Declared last so the pool (and its worker threads) is destroyed and joined + // first, before any other member variables are torn down. + ThreadPool pool_; }; +/** + * The automatic size for the pool running primary read attempts. + * + * These threads block inside a synchronous read, so the pool needs roughly one + * thread per concurrent application read. That has little to do with the core + * count; the floor is what serves small hosts. + * + * Used by `StorageConnectionImpl` when `ReadThreadPoolSizeOption` is unset (0). + */ +inline std::size_t DefaultReadThreadPoolSize() { + static std::size_t const kCores = std::thread::hardware_concurrency(); + return (std::max)(64, 4 * kCores); +} + +/** + * The automatic size for the pool running speculative hedge attempts. + * + * When @p max_concurrent_hedges is set it is already a ceiling on how many + * hedges can run at once, so a larger pool could never use the extra threads. + * + * Used by `StorageConnectionImpl` when `HedgingThreadPoolSizeOption` is + * unset (0). + */ +inline std::size_t DefaultHedgingThreadPoolSize( + std::int64_t max_concurrent_hedges) { + if (max_concurrent_hedges > 0) { + return static_cast(max_concurrent_hedges); + } + static std::size_t const kCores = std::thread::hardware_concurrency(); + return (std::max)(16, 2 * kCores); +} + } // namespace internal GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage diff --git a/google/cloud/storage/internal/hedging_thread_pool_test.cc b/google/cloud/storage/internal/hedging_thread_pool_test.cc index 043f40232ce17..117be2ecd0345 100644 --- a/google/cloud/storage/internal/hedging_thread_pool_test.cc +++ b/google/cloud/storage/internal/hedging_thread_pool_test.cc @@ -15,9 +15,13 @@ #include "google/cloud/storage/internal/hedging_thread_pool.h" #include #include +#include #include +#include #include +#include #include +#include namespace google { namespace cloud { @@ -26,6 +30,96 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN namespace internal { namespace { +using ::testing::Eq; +using ::testing::Ge; + +TEST(ThreadPoolTest, EnqueueAndExecute) { + std::promise p1; + std::promise p2; + std::future f1 = p1.get_future(); + std::future f2 = p2.get_future(); + + ThreadPool pool(2); + EXPECT_TRUE(pool.Enqueue([&p1] { p1.set_value(); })); + EXPECT_TRUE(pool.Enqueue([&p2] { p2.set_value(); })); + + f1.get(); + f2.get(); +} + +TEST(ThreadPoolTest, ConcurrentTasks) { + std::size_t const task_count = 50; + std::atomic completed{0}; + std::vector> promises(task_count); + std::vector> futures; + futures.reserve(task_count); + // Declare the pool after the promises so workers are joined before the + // promises and atomic counter they reference go out of scope. + ThreadPool pool(8); + for (std::size_t i = 0; i != task_count; ++i) { + futures.push_back(promises[i].get_future()); + EXPECT_TRUE(pool.Enqueue([&promises, &completed, i] { + ++completed; + promises[i].set_value(); + })); + } + + for (auto& f : futures) { + f.get(); + } + EXPECT_THAT(completed.load(), Eq(static_cast(task_count))); +} + +TEST(ThreadPoolTest, ZeroThreadsStillRunsTasks) { + // A pool that could never spawn a worker would queue the task, report + // success, and leave anyone waiting on the task's side effects blocked + // forever. The size is clamped to 1 instead. + std::promise p; + std::future f = p.get_future(); + ThreadPool pool(0); + EXPECT_THAT(pool.max_threads(), Eq(std::size_t{1})); + EXPECT_TRUE(pool.Enqueue([&p] { p.set_value(); })); + f.get(); +} + +TEST(ThreadPoolTest, DefaultSizes) { + // `StorageConnectionImpl` sizes its pools with these when the options are + // unset (0), so the floors are the contract. + EXPECT_THAT(DefaultReadThreadPoolSize(), Ge(std::size_t{64})); + EXPECT_THAT(DefaultHedgingThreadPoolSize(0), Ge(std::size_t{16})); + // An explicit hedge ceiling already bounds concurrency, so the pool matches + // it rather than the (larger) automatic size. + EXPECT_THAT(DefaultHedgingThreadPoolSize(3), Eq(std::size_t{3})); + EXPECT_THAT(DefaultHedgingThreadPoolSize(1), Eq(std::size_t{1})); +} + +template +void TestSafeDestructionOnWorkerThread(std::shared_ptr pool) { + auto started = std::make_shared>(); + auto destroyed = std::make_shared>(); + auto release = std::make_shared>(); + std::future started_future = started->get_future(); + std::future destroyed_future = destroyed->get_future(); + std::shared_future release_future = release->get_future().share(); + + ASSERT_TRUE(pool->Enqueue( + [pool_copy = pool, started, destroyed, release_future]() mutable { + started->set_value(); + release_future.wait(); + pool_copy.reset(); + destroyed->set_value(); + })); + + started_future.get(); + pool.reset(); + release->set_value(); + destroyed_future.get(); +} + +TEST(ThreadPoolTest, SafeDestructionOnWorkerThread) { + TestSafeDestructionOnWorkerThread(std::make_shared(1)); +} + TEST(HedgingThreadPoolTest, EnqueueAndExecute) { // Declare the promises *before* the pool. `get()` returns as soon as the // shared state is ready, which may be before the worker has returned from @@ -33,8 +127,8 @@ TEST(HedgingThreadPoolTest, EnqueueAndExecute) { // means the workers are done before the promises they reference go away. std::promise p1; std::promise p2; - auto f1 = p1.get_future(); - auto f2 = p2.get_future(); + std::future f1 = p1.get_future(); + std::future f2 = p2.get_future(); HedgingThreadPool pool(2, 0.0, 0.0, 0); EXPECT_TRUE(pool.Enqueue([&p1] { p1.set_value(); })); @@ -86,38 +180,8 @@ TEST(HedgingThreadPoolTest, FractionalRateLimiter) { } TEST(HedgingThreadPoolTest, SafeDestructionOnWorkerThread) { - // Dropping the last reference to the pool from inside a task runs the pool - // destructor on one of its own worker threads. It must detach that thread - // rather than join itself. - // - // Every synchronization object is shared and captured by value: this worker - // is detached, so it can still be running after this function returns and - // must not reference anything on the test's stack. - auto started = std::make_shared>(); - auto destroyed = std::make_shared>(); - auto release = std::make_shared>(); - auto started_future = started->get_future(); - auto destroyed_future = destroyed->get_future(); - auto release_future = release->get_future().share(); - - auto pool = std::make_shared(1, 0.0, 0.0, 0); - ASSERT_TRUE(pool->Enqueue( - [pool_copy = pool, started, destroyed, release_future]() mutable { - started->set_value(); - release_future.wait(); - // The test thread has dropped its reference by now, so this is the - // last one: the pool destructor runs on this worker thread. - pool_copy.reset(); - destroyed->set_value(); - })); - - started_future.get(); - pool.reset(); - release->set_value(); - // Wait for the destructor to finish on the worker thread. This replaces a - // timing-based sleep: the test cannot return while the pool is still being - // destroyed. - destroyed_future.get(); + TestSafeDestructionOnWorkerThread( + std::make_shared(1, 0.0, 0.0, 0)); } } // namespace diff --git a/google/cloud/storage/options.h b/google/cloud/storage/options.h index a7d192a2463b0..fd8e94a3691b6 100644 --- a/google/cloud/storage/options.h +++ b/google/cloud/storage/options.h @@ -24,6 +24,7 @@ #include "google/cloud/internal/rest_options.h" #include "google/cloud/options.h" #include +#include #include #include #include @@ -110,6 +111,31 @@ struct MaxReadHedgesOption { using Type = int; }; +/** + * The maximum number of threads in the thread pool used for primary reads + * when `EnableReadHedgingOption` is enabled. + * + * Sizing defaults to at least 64 threads or 4x hardware concurrency. + * + * @ingroup storage-options + */ +struct ReadThreadPoolSizeOption { + using Type = std::size_t; +}; + +/** + * The maximum number of threads in the thread pool used for speculative + * hedged requests when `EnableReadHedgingOption` is enabled. + * + * Sizing defaults to `MaxConcurrentHedgesOption` if set, or at least 16 + * threads or 2x hardware concurrency. + * + * @ingroup storage-options + */ +struct HedgingThreadPoolSizeOption { + using Type = std::size_t; +}; + /** * Set the HTTP version used by the client. * @@ -494,6 +520,8 @@ using ClientOptionList = ::google::cloud::OptionList< storage_experimental::MaximumHedgeBufferOption, storage_experimental::ReadHedgeDelayOption, storage_experimental::MaxReadHedgesOption, + storage_experimental::ReadThreadPoolSizeOption, + storage_experimental::HedgingThreadPoolSizeOption, storage_experimental::OTelSpanEnrichmentOption>; GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END