feat(storage): extend read hedging to mid-stream chunk reads

google/cloud/storage/internal/connection_impl.cc

@@ -420,16 +420,46 @@ StatusOr<std::unique_ptr<ObjectReadSource>> StorageConnectionImpl::ReadObject(
         *current, request, where);
   };
 
-  auto retry_source_factory =
-      [factory, current,
-       request]() -> StatusOr<std::unique_ptr<ObjectReadSource>> {
-    auto retry_policy = current->get<RetryPolicyOption>()->clone();
-    auto backoff_policy = current->get<BackoffPolicyOption>()->clone();
-    auto child = factory(request, *retry_policy, *backoff_policy);
+  HedgedObjectReadSource::Position position;
+  position.direction =
+      request.HasOption<ReadLast>() ? kFromEnd : kFromBeginning;
+  position.offset = position.direction == kFromEnd
+                        ? request.GetOption<ReadLast>().value()
+                        : request.StartingByte();
+  if (request.HasOption<ReadRange>()) {
+    position.end_offset = request.GetOption<ReadRange>().value().end;
+  }
+  if (request.HasOption<Generation>()) {
+    position.generation = request.GetOption<Generation>().value();
+  }
+
+  // Creates a `RetryObjectReadSource` positioned at `current_offset`, this is
+  // the same request rewrite `RetryObjectReadSource` applies when it resumes
+  // after a failure.
+  auto child_factory = [factory, current, request](
+                           std::int64_t current_offset,
+                           std::optional<std::int64_t> generation)
+      -> StatusOr<std::unique_ptr<ObjectReadSource>> {
+    ReadObjectRangeRequest req = request;
+    if (req.HasOption<ReadLast>()) {
+      req.set_option(ReadLast(current_offset));
+    } else if (current_offset != 0 || req.HasOption<ReadRange>() ||
+               req.HasOption<ReadFromOffset>()) {
+      req.set_option(ReadFromOffset(current_offset));
+    }
+    if (generation) {
+      req.set_option(Generation(*generation));
+    }
+    std::unique_ptr<RetryPolicy> retry_policy =
+        current->get<RetryPolicyOption>()->clone();
+    std::unique_ptr<BackoffPolicy> backoff_policy =
+        current->get<BackoffPolicyOption>()->clone();
+    StatusOr<std::unique_ptr<ObjectReadSource>> child =
+        factory(req, *retry_policy, *backoff_policy);
     if (!child) return child;
     return std::unique_ptr<ObjectReadSource>(
         std::make_unique<RetryObjectReadSource>(
-            factory, current, request, *std::move(child),
+            factory, current, std::move(req), *std::move(child),
             std::move(retry_policy), std::move(backoff_policy)));
   };
 
@@ -442,15 +472,15 @@ StatusOr<std::unique_ptr<ObjectReadSource>> StorageConnectionImpl::ReadObject(
       current->get<storage_experimental::MaximumHedgeBufferOption>();
 
   if (!enable_hedging || max_hedges <= 0 || !hedge_pool_ || !read_pool_) {
-    return retry_source_factory();
+    return child_factory(position.offset, position.generation);
   }
 
   // `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<ObjectReadSource>(
-      std::make_unique<HedgedObjectReadSource>(read_pool_, hedge_pool_,
-                                               std::move(retry_source_factory),
-                                               delay, max_hedges, max_buffer));
+      std::make_unique<HedgedObjectReadSource>(
+          read_pool_, hedge_pool_, std::move(child_factory), delay, max_hedges,
+          max_buffer, position));
 }
 
 StatusOr<ListObjectsResponse> StorageConnectionImpl::ListObjects(

google/cloud/storage/internal/connection_impl_test.cc

@@ -16,11 +16,13 @@
 #include "google/cloud/storage/internal/tracing_connection.h"
 #include "google/cloud/storage/options.h"
 #include "google/cloud/storage/testing/canonical_errors.h"
+#include "google/cloud/storage/testing/mock_client.h"
 #include "google/cloud/storage/testing/mock_generic_stub.h"
 #include "google/cloud/testing_util/chrono_literals.h"
 #include "google/cloud/testing_util/opentelemetry_matchers.h"
 #include "google/cloud/testing_util/status_matchers.h"
 #include <gmock/gmock.h>
+#include <functional>
 #include <memory>
 #include <string>
 #include <utility>
@@ -695,6 +697,95 @@ TEST(RetryClientTest, BackoffSpansUploadChunk) {
                   SpanNamed("storage::Client::WriteObject/UploadChunk")));
 }
 
+// `ReadObject()` positions the child stream it creates, using the same request
+// rewrite a resumed or hedged read applies. Verify the request that reaches
+// the stub for each way a caller can express a starting position.
+StatusOr<std::unique_ptr<ObjectReadSource>> ReadObjectWithRequest(
+    ReadObjectRangeRequest const& request,
+    std::function<void(ReadObjectRangeRequest const&)> check) {
+  auto mock = std::make_unique<MockGenericStub>();
+  EXPECT_CALL(*mock, options).Times(AtLeast(0));
+  EXPECT_CALL(*mock, ReadObject)
+      .WillOnce([check = std::move(check)](auto&, auto const&,
+                                           ReadObjectRangeRequest const& req) {
+        check(req);
+        return std::unique_ptr<ObjectReadSource>(
+            std::make_unique<testing::MockObjectReadSource>());
+      });
+  auto client = StorageConnectionImpl::Create(std::move(mock));
+  google::cloud::internal::OptionsSpan const span(BasicTestPolicies());
+  return client->ReadObject(request);
+}
+
+TEST(RetryClientTest, ReadObjectPositionsPlainRequest) {
+  ReadObjectRangeRequest request("test-bucket", "test-object");
+  EXPECT_THAT(
+      ReadObjectWithRequest(request,
+                            [](ReadObjectRangeRequest const& req) {
+                              // Offset 0 with no range: the rewrite must not
+                              // introduce a `ReadFromOffset(0)`, which would
+                              // change the request.
+                              EXPECT_FALSE(req.HasOption<ReadFromOffset>());
+                              EXPECT_FALSE(req.HasOption<ReadLast>());
+                              EXPECT_FALSE(req.HasOption<Generation>());
+                            }),
+      IsOk());
+}
+
+TEST(RetryClientTest, ReadObjectPositionsReadFromOffset) {
+  ReadObjectRangeRequest request("test-bucket", "test-object");
+  request.set_option(ReadFromOffset(1024));
+  EXPECT_THAT(ReadObjectWithRequest(
+                  request,
+                  [](ReadObjectRangeRequest const& req) {
+                    EXPECT_EQ(1024,
+                              req.GetOption<ReadFromOffset>().value_or(0));
+                  }),
+              IsOk());
+}
+
+TEST(RetryClientTest, ReadObjectPositionsReadRange) {
+  ReadObjectRangeRequest request("test-bucket", "test-object");
+  request.set_option(ReadRange(100, 200));
+  EXPECT_THAT(ReadObjectWithRequest(
+                  request,
+                  [](ReadObjectRangeRequest const& req) {
+                    // The range is preserved, and the offset is pinned to its
+                    // start so a resumed read picks up where this one left off.
+                    EXPECT_EQ(100, req.GetOption<ReadFromOffset>().value_or(0));
+                    ASSERT_TRUE(req.HasOption<ReadRange>());
+                    EXPECT_EQ(200, req.GetOption<ReadRange>().value().end);
+                  }),
+              IsOk());
+}
+
+TEST(RetryClientTest, ReadObjectPositionsReadLast) {
+  ReadObjectRangeRequest request("test-bucket", "test-object");
+  request.set_option(ReadLast(512));
+  EXPECT_THAT(ReadObjectWithRequest(
+                  request,
+                  [](ReadObjectRangeRequest const& req) {
+                    // `ReadLast` counts from the end, so it is rewritten in
+                    // place and must not become a `ReadFromOffset`.
+                    EXPECT_EQ(512, req.GetOption<ReadLast>().value_or(0));
+                    EXPECT_FALSE(req.HasOption<ReadFromOffset>());
+                  }),
+              IsOk());
+}
+
+TEST(RetryClientTest, ReadObjectPinsGeneration) {
+  ReadObjectRangeRequest request("test-bucket", "test-object");
+  request.set_option(Generation(12345));
+  request.set_option(ReadFromOffset(64));
+  EXPECT_THAT(ReadObjectWithRequest(
+                  request,
+                  [](ReadObjectRangeRequest const& req) {
+                    EXPECT_EQ(12345, req.GetOption<Generation>().value_or(0));
+                    EXPECT_EQ(64, req.GetOption<ReadFromOffset>().value_or(0));
+                  }),
+              IsOk());
+}
+
 }  // namespace
 }  // namespace internal
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END