fix(storage): hedged read attempts must not own the connection

google/cloud/storage/internal/connection_impl_test.cc

@@ -25,6 +25,8 @@
 #include <gmock/gmock.h>
 #include <opentelemetry/metrics/provider.h>
 #include <algorithm>
+#include <atomic>
+#include <chrono>
 #include <cstdint>
 #include <functional>
 #include <future>
@@ -824,7 +826,6 @@ ReadSourceResult MakeReadResult(std::string const& payload, char* buf) {
 }
 
 TEST(RetryClientTest, HedgedReadRecordsMetricsOnGlobalMeterProvider) {
-  GTEST_SKIP() << "Skipping test that is leaking mock objects #16522";
   auto dispatched = std::make_unique<MockCounter<std::uint64_t>>();
   auto won = std::make_unique<MockCounter<std::uint64_t>>();
   EXPECT_CALL(*dispatched, Add(std::uint64_t{1})).Times(1);
@@ -852,37 +853,44 @@ TEST(RetryClientTest, HedgedReadRecordsMetricsOnGlobalMeterProvider) {
   // A single read thread, so the first read tells us which thread opens every
   // primary attempt. Any other thread is running a hedge.
   auto primary_thread = std::make_shared<std::thread::id>();
+  auto primary_opens = std::make_shared<std::atomic<int>>(0);
   auto unblock_primary = std::make_shared<std::promise<void>>();
-  auto primary_closed = std::make_shared<std::promise<void>>();
+  std::shared_future<void> const primary_unblocked =
+      unblock_primary->get_future().share();
   auto mock = std::make_unique<MockGenericStub>();
   EXPECT_CALL(*mock, options).Times(AtLeast(0));
   EXPECT_CALL(*mock, ReadObject)
-      .WillOnce([primary_thread](auto&, auto const&,
-                                 ReadObjectRangeRequest const&) {
-        *primary_thread = std::this_thread::get_id();
-        auto source = std::make_unique<testing::MockObjectReadSource>();
-        EXPECT_CALL(*source, Read).WillOnce([](char* buf, std::size_t) {
-          return MakeReadResult("warm-up", buf);
-        });
-        return StatusOr<std::unique_ptr<ObjectReadSource>>(std::move(source));
-      })
-      .WillRepeatedly([primary_thread, unblock_primary, primary_closed](
+      .WillRepeatedly([primary_thread, primary_opens, primary_unblocked](
                           auto&, auto const&, ReadObjectRangeRequest const&) {
         auto source = std::make_unique<testing::MockObjectReadSource>();
-        if (std::this_thread::get_id() == *primary_thread) {
-          EXPECT_CALL(*source, Read)
-              .WillOnce([unblock_primary](char* buf, std::size_t) {
-                unblock_primary->get_future().wait();
-                return MakeReadResult("slow", buf);
-              });
-          EXPECT_CALL(*source, Close).WillOnce([primary_closed] {
-            primary_closed->set_value();
-            return make_status_or(HttpResponse{HttpStatusCode::kOk, {}, {}});
-          });
-        } else {
+        if (*primary_opens == 0) *primary_thread = std::this_thread::get_id();
+        if (std::this_thread::get_id() != *primary_thread) {
           EXPECT_CALL(*source, Read).WillOnce([](char* buf, std::size_t) {
             return MakeReadResult("hedge", buf);
           });
+          return StatusOr<std::unique_ptr<ObjectReadSource>>(std::move(source));
+        }
+        switch ((*primary_opens)++) {
+          case 0:
+            EXPECT_CALL(*source, Read).WillOnce([](char* buf, std::size_t) {
+              return MakeReadResult("warm-up", buf);
+            });
+            break;
+          case 1:
+            EXPECT_CALL(*source, Read)
+                .WillOnce([primary_unblocked](char* buf, std::size_t) {
+                  primary_unblocked.wait();
+                  return MakeReadResult("slow", buf);
+                });
+            EXPECT_CALL(*source, Close).WillOnce([] {
+              return make_status_or(HttpResponse{HttpStatusCode::kOk, {}, {}});
+            });
+            break;
+          default:
+            EXPECT_CALL(*source, Read).WillOnce([](char* buf, std::size_t) {
+              return MakeReadResult("drain", buf);
+            });
+            break;
         }
         return StatusOr<std::unique_ptr<ObjectReadSource>>(std::move(source));
       });
@@ -908,18 +916,40 @@ TEST(RetryClientTest, HedgedReadRecordsMetricsOnGlobalMeterProvider) {
   }
 
   // Primary stalls, the hedge wins.
-  google::cloud::internal::OptionsSpan const span(
-      client->options().set<storage_experimental::ReadHedgeDelayOption>(
-          std::chrono::milliseconds(1)));
-  StatusOr<std::unique_ptr<ObjectReadSource>> source =
-      client->ReadObject(ReadObjectRangeRequest("test-bucket", "test-object"));
-  ASSERT_THAT(source, IsOk());
-  StatusOr<ReadSourceResult> result =
-      (*source)->Read(buffer.data(), buffer.size());
+  StatusOr<std::unique_ptr<ObjectReadSource>> source;
+  {
+    google::cloud::internal::OptionsSpan const span(
+        client->options().set<storage_experimental::ReadHedgeDelayOption>(
+            std::chrono::milliseconds(1)));
+    source = client->ReadObject(
+        ReadObjectRangeRequest("test-bucket", "test-object"));
+    ASSERT_THAT(source, IsOk());
+    StatusOr<ReadSourceResult> const result =
+        (*source)->Read(buffer.data(), buffer.size());
+    ASSERT_THAT(result, IsOk());
+    EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+                Eq("hedge"));
+  }
+
+  // The losing primary is still closing its `RetryObjectReadSource` (which
+  // holds a reference to `client`) on the single read thread. A synchronous
+  // read on that same thread can only run after the loser's task has finished,
+  // so `client` is destroyed on this thread rather than on a detached worker.
   unblock_primary->set_value();
-  primary_closed->get_future().wait();
-  ASSERT_THAT(result, IsOk());
-  EXPECT_THAT(std::string(buffer.data(), result->bytes_received), Eq("hedge"));
+  {
+    google::cloud::internal::OptionsSpan const span(
+        client->options().set<storage_experimental::ReadHedgeDelayOption>(
+            std::chrono::seconds(30)));
+    StatusOr<std::unique_ptr<ObjectReadSource>> drain =
+        client->ReadObject(ReadObjectRangeRequest("test-bucket", "drain"));
+    ASSERT_THAT(drain, IsOk());
+    StatusOr<ReadSourceResult> const result =
+        (*drain)->Read(buffer.data(), buffer.size());
+    ASSERT_THAT(result, IsOk());
+    EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+                Eq("drain"));
+  }
+  EXPECT_EQ(3, primary_opens->load());
 }
 
 }  // namespace

google/cloud/storage/internal/hedged_object_read_source.cc

@@ -115,7 +115,8 @@ struct RaceState {
 // loser closes its own child. A failed attempt only resolves the race if it
 // is the last one standing, or if it is the primary failing permanently.
 void RunAttempt(std::shared_ptr<RaceState> const& state,
-                HedgedObjectReadSource::ChildFactory const& factory,
+                std::weak_ptr<HedgedObjectReadSource::ChildFactory const> const&
+                    weak_factory,
                 std::unique_ptr<ObjectReadSource> child,
                 std::unique_ptr<char[]> buffer, std::size_t buffer_capacity,
                 std::int64_t offset, std::optional<std::int64_t> generation,
@@ -141,8 +142,20 @@ void RunAttempt(std::shared_ptr<RaceState> const& state,
   } guard{std::move(release_slot)};
 
   if (!child) {
-    StatusOr<std::unique_ptr<ObjectReadSource>> source =
-        factory(offset, generation);
+    // Lock `weak_factory` only for the open call: the factory captures the
+    // connection, which owns the thread pools this task runs on, so the task
+    // must not keep it alive past `promise.set_value()`.
+    StatusOr<std::unique_ptr<ObjectReadSource>> source = [&] {
+      std::shared_ptr<HedgedObjectReadSource::ChildFactory const> const
+          factory = weak_factory.lock();
+      if (!factory) {
+        return StatusOr<std::unique_ptr<ObjectReadSource>>(
+            google::cloud::internal::CancelledError(
+                "stream destroyed before the read attempt started",
+                GCP_ERROR_INFO()));
+      }
+      return (*factory)(offset, generation);
+    }();
     if (!source) return state->Fail(std::move(source).status(), is_primary);
     child = *std::move(source);
   }
@@ -310,9 +323,10 @@ StatusOr<ReadSourceResult> HedgedObjectReadSource::ReadRaced(char* buf,
   staging_buffer_.reset();
   staging_buffer_capacity_ = 0;
 
-  auto primary = [state, factory = child_factory_, offset = current_offset_,
-                  gen = generation_, n] {
-    RunAttempt(state, *factory, std::move(state->primary_child),
+  auto primary = [state,
+                  factory = std::weak_ptr<ChildFactory const>(child_factory_),
+                  offset = current_offset_, gen = generation_, n] {
+    RunAttempt(state, factory, std::move(state->primary_child),
                std::move(state->primary_buffer), state->primary_buffer_capacity,
                offset, gen, n,
                /*is_primary=*/true, std::weak_ptr<HedgingThreadPool>{});
@@ -339,10 +353,11 @@ StatusOr<ReadSourceResult> HedgedObjectReadSource::ReadRaced(char* buf,
       continue;
     }
     state->active_attempts.fetch_add(1);
-    auto hedge = [state, factory = child_factory_, offset = current_offset_,
-                  gen = generation_, n,
+    auto hedge = [state,
+                  factory = std::weak_ptr<ChildFactory const>(child_factory_),
+                  offset = current_offset_, gen = generation_, n,
                   pool = std::weak_ptr<HedgingThreadPool>(hedge_pool_)] {
-      RunAttempt(state, *factory, /*child=*/nullptr, /*buffer=*/nullptr,
+      RunAttempt(state, factory, /*child=*/nullptr, /*buffer=*/nullptr,
                  /*buffer_capacity=*/0, offset, gen, n, /*is_primary=*/false,
                  pool);
     };

google/cloud/storage/internal/hedged_object_read_source.h

@@ -121,7 +121,7 @@ class HedgedObjectReadSource : public ObjectReadSource {
 
   std::shared_ptr<ThreadPool> read_pool_;
   std::shared_ptr<HedgingThreadPool> hedge_pool_;
-  // Shared with the racing attempts, which may outlive this object.
+  // Held by the racing attempts through a `std::weak_ptr`, see `RunAttempt()`.
   std::shared_ptr<ChildFactory const> child_factory_;
   std::chrono::milliseconds delay_;
   int max_hedges_;

google/cloud/storage/internal/hedged_object_read_source_test.cc

@@ -23,6 +23,7 @@
 #include <chrono>
 #include <future>
 #include <memory>
+#include <optional>
 #include <string>
 #include <thread>
 #include <vector>
@@ -1527,6 +1528,53 @@ TEST(HedgedObjectReadSourceTest,
   // the losing hedge's mock before returning.
 }
 
+TEST(HedgedObjectReadSourceTest, QueuedAttemptAfterDestructionDoesNotOpen) {
+  // An attempt that starts running only after the `HedgedObjectReadSource` has
+  // been destroyed must retire without calling the factory.
+  std::optional<PrimaryReadPool> read_pool;
+  read_pool.emplace();
+  std::thread::id const primary_thread = read_pool->worker_id();
+
+  // Hold the single read worker so the primary attempt stays queued until
+  // `source` has been destroyed.
+  auto release_worker = std::make_shared<std::promise<void>>();
+  ASSERT_TRUE(read_pool->pool()->Enqueue(
+      [release_worker] { release_worker->get_future().wait(); }));
+
+  auto primary_opens = std::make_shared<std::atomic<int>>(0);
+  auto factory =
+      [primary_opens,
+       primary_thread]() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
+    auto mock = std::make_unique<MockObjectReadSource>();
+    if (OnPrimaryThread(primary_thread)) {
+      ++*primary_opens;
+      EXPECT_CALL(*mock, Read).WillOnce(ImmediateRead("stale"));
+      EXPECT_CALL(*mock, Close)
+          .Times(AtMost(1))
+          .WillRepeatedly(Return(
+              make_status_or(HttpResponse{HttpStatusCode::kOk, {}, {}})));
+      return std::unique_ptr<ObjectReadSource>(std::move(mock));
+    }
+    EXPECT_CALL(*mock, Read).WillOnce(ImmediateRead("hedge"));
+    return std::unique_ptr<ObjectReadSource>(std::move(mock));
+  };
+
+  {
+    HedgedObjectReadSource source(read_pool->pool(), MakeUnlimitedHedgePool(),
+                                  Adapt(factory), std::chrono::milliseconds(1),
+                                  /*max_hedges=*/1, kUnlimitedBuffer);
+    std::vector<char> buffer(100);
+    StatusOr<ReadSourceResult> result =
+        source.Read(buffer.data(), buffer.size());
+    ASSERT_THAT(result, IsOk());
+    EXPECT_THAT(std::string(buffer.data(), result->bytes_received),
+                Eq("hedge"));
+  }
+  release_worker->set_value();
+  read_pool.reset();
+  EXPECT_THAT(primary_opens->load(), Eq(0));
+}
+
 // The hedging counters backed by mocks. The counters are owned by `metrics`,
 // the raw pointers stay valid for as long as it lives.
 struct MockHedgingCounters {