fix(storage): address review comments on read hedging and thread pool

Author: ajayky-os

google/cloud/storage/internal/hedged_object_read_source.cc

@@ -45,6 +45,9 @@ void RunAttempt(std::shared_ptr<RaceState> const& state,
                 HedgedObjectReadSource::ChildFactory const& factory,
                 std::size_t n, bool resolve_on_open_error,
                 std::shared_ptr<HedgingThreadPool> release_slot) {
+  // Releases the acquired hedge concurrency slot upon function exit across
+  // all code paths (early return on open/allocation error, race winner, or
+  // race loser). For primary attempts, release_slot is nullptr.
   struct SlotGuard {
     std::shared_ptr<HedgingThreadPool> pool;
     ~SlotGuard() {
@@ -113,7 +116,9 @@ StatusOr<HttpResponse> HedgedObjectReadSource::Close() {
 
 StatusOr<ReadSourceResult> HedgedObjectReadSource::Read(char* buf,
                                                         std::size_t n) {
-  if (is_closed_) return ReadSourceResult{};
+  if (is_closed_) {
+    return ReadSourceResult{0, HttpResponse{HttpStatusCode::kOk, {}, {}}};
+  }
 
   // Only the stream open is hedged. Once a child has won the race all
   // subsequent reads continue on it, at its current offset, without any
@@ -142,9 +147,20 @@ StatusOr<ReadSourceResult> HedgedObjectReadSource::Read(char* buf,
   // complete either way.
   if (!read_pool_->Enqueue(primary)) primary();
 
-  for (int hedges_dispatched = 0; hedges_dispatched != max_hedges_;) {
+  for (int hedges_dispatched = 0; hedges_dispatched < max_hedges_;) {
     if (future.wait_for(delay_) != std::future_status::timeout) break;
-    if (!hedge_pool_->TryAcquireHedgeToken()) continue;
+    if (!hedge_pool_->TryAcquireHedgeToken()) {
+      // When delay_ is 0ms (or token acquisition fails), back off briefly on
+      // the future instead of busy-spinning if tokens or concurrency slots are
+      // temporarily exhausted.
+      if (delay_ == std::chrono::milliseconds::zero()) {
+        if (future.wait_for(std::chrono::milliseconds(10)) !=
+            std::future_status::timeout) {
+          break;
+        }
+      }
+      continue;
+    }
     auto hedge = [state, factory = child_factory_, n, pool = hedge_pool_] {
       RunAttempt(state, factory, n, /*resolve_on_open_error=*/false, pool);
     };

google/cloud/storage/internal/hedged_object_read_source_test.cc

@@ -271,6 +271,134 @@ TEST(HedgedObjectReadSourceTest,
   primary_closed->get_future().get();
 }
 
+TEST(HedgedObjectReadSourceTest, HedgeOpenFailureReleasesSlot) {
+  // Verify that if a hedge attempt fails during stream opening (factory()
+  // error), the hedge concurrency slot is released via RAII (SlotGuard) and is
+  // not leaked.
+  auto unblock_primary = std::make_shared<std::promise<void>>();
+  auto calls = std::make_shared<std::atomic<int>>(0);
+  auto factory = [unblock_primary,
+                  calls]() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
+    int call_count = ++*calls;
+    if (call_count == 1) {
+      // Primary attempt: stalls until unblocked.
+      auto mock = std::make_unique<MockObjectReadSource>();
+      EXPECT_CALL(*mock, Read)
+          .WillOnce([unblock_primary](char* buf, std::size_t) {
+            unblock_primary->get_future().get();
+            std::string const payload = "primary";
+            std::copy(payload.begin(), payload.end(), buf);
+            return MakeReadResult(payload);
+          });
+      return std::unique_ptr<ObjectReadSource>(std::move(mock));
+    }
+    // Hedge attempt: fails to open.
+    return Status(StatusCode::kUnavailable, "open failed");
+  };
+
+  auto hedge_pool = std::make_shared<HedgingThreadPool>(
+      /*max_threads=*/2, /*rate_limit=*/0.0, /*capacity=*/0.0,
+      /*max_concurrent=*/1);
+
+  HedgedObjectReadSource source(MakeUnlimitedReadPool(), hedge_pool, factory,
+                                std::chrono::milliseconds(1),
+                                /*max_hedges=*/1, kUnlimitedBuffer);
+
+  std::vector<char> buffer(100);
+  std::thread unblocker([unblock_primary] {
+    std::this_thread::sleep_for(std::chrono::milliseconds(30));
+    unblock_primary->set_value();
+  });
+
+  auto result = source.Read(buffer.data(), buffer.size());
+  unblocker.join();
+
+  ASSERT_THAT(result, IsOk());
+  EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+              Eq("primary"));
+  EXPECT_THAT(calls->load(), Eq(2));
+
+  // If the slot leaked on open failure, TryAcquireHedgeToken would fail because
+  // max_concurrent is 1.
+  EXPECT_TRUE(hedge_pool->TryAcquireHedgeToken());
+  hedge_pool->ReleaseHedgeSlot();
+}
+
+TEST(HedgedObjectReadSourceTest, ZeroDelayBacksOffOnHedgeTokenExhaustion) {
+  // Verify that when delay_ == 0ms and TryAcquireHedgeToken() returns false,
+  // the hedging loop backs off instead of busy-spinning, allowing the primary
+  // read to complete normally.
+  auto unblock_primary = std::make_shared<std::promise<void>>();
+  auto calls = std::make_shared<std::atomic<int>>(0);
+  auto factory = [unblock_primary,
+                  calls]() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
+    ++*calls;
+    auto mock = std::make_unique<MockObjectReadSource>();
+    EXPECT_CALL(*mock, Read)
+        .WillOnce([unblock_primary](char* buf, std::size_t) {
+          unblock_primary->get_future().get();
+          std::string const payload = "primary_data";
+          std::copy(payload.begin(), payload.end(), buf);
+          return MakeReadResult(payload);
+        });
+    return std::unique_ptr<ObjectReadSource>(std::move(mock));
+  };
+
+  auto hedge_pool = std::make_shared<HedgingThreadPool>(
+      /*max_threads=*/1, /*rate_limit=*/0.0, /*capacity=*/0.0,
+      /*max_concurrent=*/1);
+  // Exhaust all hedge slots so TryAcquireHedgeToken fails.
+  ASSERT_TRUE(hedge_pool->TryAcquireHedgeToken());
+
+  HedgedObjectReadSource source(MakeUnlimitedReadPool(), hedge_pool, factory,
+                                std::chrono::milliseconds(0),
+                                /*max_hedges=*/2, kUnlimitedBuffer);
+
+  std::vector<char> buffer(100);
+  std::thread unblocker([unblock_primary] {
+    std::this_thread::sleep_for(std::chrono::milliseconds(30));
+    unblock_primary->set_value();
+  });
+
+  auto result = source.Read(buffer.data(), buffer.size());
+  unblocker.join();
+
+  ASSERT_THAT(result, IsOk());
+  EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+              Eq("primary_data"));
+  EXPECT_THAT(calls->load(), Eq(1));
+
+  hedge_pool->ReleaseHedgeSlot();
+}
+
+TEST(HedgedObjectReadSourceTest, NonPositiveMaxHedgesDoesNotHedge) {
+  // Verify that negative or zero max_hedges values defensively result in 0
+  // hedge attempts, running only the primary attempt.
+  auto calls = std::make_shared<std::atomic<int>>(0);
+  auto factory = [calls]() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
+    ++*calls;
+    auto mock = std::make_unique<MockObjectReadSource>();
+    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<ObjectReadSource>(std::move(mock));
+  };
+
+  HedgedObjectReadSource source(MakeUnlimitedReadPool(),
+                                MakeUnlimitedHedgePool(), factory,
+                                std::chrono::milliseconds(0),
+                                /*max_hedges=*/-1, kUnlimitedBuffer);
+
+  std::vector<char> buffer(100);
+  auto result = source.Read(buffer.data(), buffer.size());
+  ASSERT_THAT(result, IsOk());
+  EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+              Eq("primary_only"));
+  EXPECT_THAT(calls->load(), Eq(1));
+}
+
 TEST(HedgedObjectReadSourceTest, PrimaryOpenErrorPropagates) {
   auto factory = []() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
     return Status(StatusCode::kPermissionDenied, "uh-oh");

google/cloud/storage/internal/hedging_thread_pool.h

@@ -135,15 +135,28 @@ class ThreadPool {
 };
 
 /**
- * A dedicated thread pool with integrated hedge throttling.
+ * Coordinates and bounds speculative hedged requests across a storage client.
  *
  * 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`.
+ * limits: a maximum number of concurrently active hedges (when
+ * `max_concurrent > 0`), and a maximum rate of new hedges per second via a
+ * token bucket (when `rate_limit > 0.0`). Setting `rate_limit <= 0.0` disables
+ * rate limiting (unlimited hedges per second). Task execution is dispatched
+ * onto a dedicated internal `ThreadPool`.
  */
 class HedgingThreadPool {
  public:
+  /**
+   * Constructs a `HedgingThreadPool`.
+   *
+   * @param max_threads the worker pool thread limit. Clamped to at least 1.
+   * @param rate_limit the token bucket refill rate in tokens/sec. When <= 0.0,
+   *     rate limiting is disabled (unlimited hedges per second).
+   * @param capacity the maximum burst capacity in tokens. Clamped to at
+   *     least 1.0.
+   * @param max_concurrent the ceiling on concurrently active hedges. When <= 0,
+   *     concurrency limiting is disabled.
+   */
   HedgingThreadPool(std::size_t max_threads, double rate_limit, double capacity,
                     std::int64_t max_concurrent)
       : rate_limit_(rate_limit),
@@ -175,7 +188,8 @@ class HedgingThreadPool {
    * On success the caller *must* eventually call `ReleaseHedgeSlot()`.
    */
   bool TryAcquireHedgeToken() {
-    // Gate 1: the ceiling on concurrently active hedges.
+    // Gate 1: the ceiling on concurrently active hedges. When
+    // max_concurrent_hedges_ <= 0, concurrency limiting is disabled.
     if (max_concurrent_hedges_ > 0) {
       std::int64_t current =
           active_concurrent_hedges_.load(std::memory_order_relaxed);
@@ -185,7 +199,8 @@ class HedgingThreadPool {
           current, current + 1, std::memory_order_relaxed));
     }
 
-    // Gate 2: the rate limit on new hedges (token bucket).
+    // Gate 2: the rate limit on new hedges (token bucket). When
+    // rate_limit_ <= 0.0, rate limiting is disabled.
     if (rate_limit_ > 0.0) {
       std::lock_guard<std::mutex> lock(limiter_mutex_);
       Refill();
@@ -219,14 +234,14 @@ class HedgingThreadPool {
     tokens_ = (std::min)(tokens_capacity_, tokens_ + elapsed * rate_limit_);
   }
 
-  // Token bucket rate limiter.
+  // Token bucket rate limiter. A rate_limit_ <= 0.0 disables rate limiting.
   double rate_limit_;
   double tokens_capacity_;
   double tokens_;
   std::chrono::steady_clock::time_point last_refill_;
   std::mutex limiter_mutex_;
 
-  // Concurrency limiter.
+  // Concurrency limiter. A max_concurrent_hedges_ <= 0 disables limit.
   std::int64_t const max_concurrent_hedges_;
   std::atomic<std::int64_t> active_concurrent_hedges_{0};
 

google/cloud/storage/internal/hedging_thread_pool_test.cc

@@ -179,6 +179,16 @@ TEST(HedgingThreadPoolTest, FractionalRateLimiter) {
   EXPECT_FALSE(pool.TryAcquireHedgeToken());
 }
 
+TEST(HedgingThreadPoolTest, ZeroRateLimitDisablesRateLimiting) {
+  // A rate limit of 0.0 disables rate limiting, allowing unlimited
+  // acquisitions.
+  HedgingThreadPool pool(5, 0.0, 0.0, 0);
+
+  for (int i = 0; i < 100; ++i) {
+    EXPECT_TRUE(pool.TryAcquireHedgeToken());
+  }
+}
+
 TEST(HedgingThreadPoolTest, SafeDestructionOnWorkerThread) {
   TestSafeDestructionOnWorkerThread(
       std::make_shared<HedgingThreadPool>(1, 0.0, 0.0, 0));

google/cloud/storage/options.h

@@ -115,7 +115,8 @@ struct MaxReadHedgesOption {
  * 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.
+ * Set to 0 for automatic sizing (defaults to at least 64 threads or 4x hardware
+ * concurrency).
  *
  * @ingroup storage-options
  */
@@ -127,8 +128,8 @@ struct ReadThreadPoolSizeOption {
  * 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.
+ * Set to 0 for automatic sizing (defaults to `MaxConcurrentHedgesOption` if
+ * set, or at least 16 threads or 2x hardware concurrency).
  *
  * @ingroup storage-options
  */