feat(storage): support pre-warming read ranges in AsyncConnection

google/cloud/storage/internal/async/connection_impl.cc

@@ -26,6 +26,7 @@
 #include "google/cloud/storage/internal/async/object_descriptor_impl.h"
 #include "google/cloud/storage/internal/async/open_object.h"
 #include "google/cloud/storage/internal/async/open_stream.h"
+#include "google/cloud/storage/internal/async/options.h"
 #include "google/cloud/storage/internal/async/read_payload_impl.h"
 #include "google/cloud/storage/internal/async/reader_connection_impl.h"
 #include "google/cloud/storage/internal/async/reader_connection_resume.h"
@@ -58,6 +59,7 @@
 #include "google/cloud/internal/make_status.h"
 #include <grpcpp/grpcpp.h>
 #include <memory>
+#include <set>
 #include <utility>
 
 namespace google {
@@ -220,6 +222,24 @@ AsyncConnectionImpl::Open(OpenParams p) {
   auto initial_request = google::storage::v2::BidiReadObjectRequest{};
   *initial_request.mutable_read_object_spec() = p.read_spec;
   auto current = internal::MakeImmutableOptions(p.options);
+  // If pre-warmed ranges are configured, populate the initial request
+  // with these ranges to start downloading them as soon as the stream opens.
+  if (current->has<ReadRangesOption>()) {
+    auto const& ranges = current->get<ReadRangesOption>();
+    std::int64_t id = 0;
+    std::set<std::pair<std::int64_t, std::int64_t>> seen_ranges;
+    for (auto const& r : ranges) {
+      auto range_key = std::make_pair(r.offset, r.length);
+      if (!seen_ranges.insert(range_key).second) {
+        continue;  // Skip duplicate range.
+      }
+      auto* proto_range = initial_request.add_read_ranges();
+      proto_range->set_read_offset(r.offset);
+      proto_range->set_read_length(r.length);
+      // Generate sequential IDs starting at 1. The receiver must match these.
+      proto_range->set_read_id(++id);
+    }
+  }
   // Get the policy factory and immediately create a policy.
   auto resume_policy = current->get<storage::ResumePolicyOption>()();
 

google/cloud/storage/internal/async/connection_impl_open_test.cc

@@ -16,6 +16,7 @@
 #include "google/cloud/storage/async/retry_policy.h"
 #include "google/cloud/storage/internal/async/connection_impl.h"
 #include "google/cloud/storage/internal/async/default_options.h"
+#include "google/cloud/storage/internal/async/options.h"
 #include "google/cloud/storage/internal/grpc/ctype_cord_workaround.h"
 #include "google/cloud/storage/testing/canonical_errors.h"
 #include "google/cloud/storage/testing/mock_resume_policy.h"
@@ -405,6 +406,150 @@ TEST(AsyncConnectionImplTest, TooManyTransienErrors) {
   ASSERT_THAT(pending.get(), StatusIs(TransientError().code()));
 }
 
+TEST(AsyncConnectionImplTest, OpenWithReadRanges) {
+  auto constexpr kExpectedRequestSpec = R"pb(
+    bucket: "test-only-invalid"
+    object: "test-object"
+    generation: 42
+    if_metageneration_match: 7
+  )pb";
+  auto constexpr kExpectedRequest = R"pb(
+    read_object_spec {
+      bucket: "test-only-invalid"
+      object: "test-object"
+      generation: 42
+      if_metageneration_match: 7
+    }
+    read_ranges { read_offset: 0 read_length: 1024 read_id: 1 }
+    read_ranges { read_offset: 2048 read_length: 4096 read_id: 2 }
+  )pb";
+  auto constexpr kMetadataText = R"pb(
+    bucket: "projects/_/buckets/test-bucket"
+    name: "test-object"
+    generation: 42
+  )pb";
+
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_shared<storage::testing::MockStorageStub>();
+  EXPECT_CALL(*mock, AsyncBidiReadObject)
+      .WillOnce([&](CompletionQueue const&,
+                    std::shared_ptr<grpc::ClientContext> const&,
+                    google::cloud::internal::ImmutableOptions const& options) {
+        EXPECT_EQ(options->get<AuthorityOption>(), kAuthority);
+
+        auto stream = std::make_unique<MockStream>();
+        EXPECT_CALL(*stream, Start).WillOnce([&sequencer]() {
+          return sequencer.PushBack("Start").then(
+              [](auto f) { return f.get(); });
+        });
+        EXPECT_CALL(*stream, Write)
+            .WillOnce(
+                [=, &sequencer](
+                    google::storage::v2::BidiReadObjectRequest const& request,
+                    grpc::WriteOptions) {
+                  auto expected = google::storage::v2::BidiReadObjectRequest{};
+                  EXPECT_TRUE(
+                      TextFormat::ParseFromString(kExpectedRequest, &expected));
+                  EXPECT_THAT(request, IsProtoEqual(expected));
+                  return sequencer.PushBack("Write").then(
+                      [](auto f) { return f.get(); });
+                });
+        EXPECT_CALL(*stream, Read)
+            .WillOnce([&]() {
+              return sequencer.PushBack("Read").then(
+                  [=](auto f) -> absl::optional<
+                                  google::storage::v2::BidiReadObjectResponse> {
+                    if (!f.get()) return absl::nullopt;
+                    auto constexpr kHandleText = R"pb(
+                      handle: "handle-12345"
+                    )pb";
+                    auto response =
+                        google::storage::v2::BidiReadObjectResponse{};
+                    EXPECT_TRUE(TextFormat::ParseFromString(
+                        kMetadataText, response.mutable_metadata()));
+                    EXPECT_TRUE(TextFormat::ParseFromString(
+                        kHandleText, response.mutable_read_handle()));
+                    return response;
+                  });
+            })
+            .WillOnce([&sequencer]() {
+              return sequencer.PushBack("Read[N]").then(
+                  [](auto f) -> absl::optional<
+                                 google::storage::v2::BidiReadObjectResponse> {
+                    if (!f.get()) return absl::nullopt;
+                    return google::storage::v2::BidiReadObjectResponse{};
+                  });
+            });
+        EXPECT_CALL(*stream, Cancel).WillOnce([&sequencer]() {
+          (void)sequencer.PushBack("Cancel");
+        });
+        EXPECT_CALL(*stream, Finish).WillOnce([&sequencer]() {
+          return sequencer.PushBack("Finish").then(
+              [](auto) { return Status{}; });
+        });
+
+        return std::unique_ptr<BidiReadStream>(std::move(stream));
+      })
+      .WillRepeatedly([](CompletionQueue const&,
+                         std::shared_ptr<grpc::ClientContext> const&,
+                         google::cloud::internal::ImmutableOptions const&) {
+        auto stream = std::make_unique<NiceMock<MockStream>>();
+        ON_CALL(*stream, Start).WillByDefault(InvokeWithoutArgs([] {
+          return make_ready_future(false);
+        }));
+        ON_CALL(*stream, Finish).WillByDefault(InvokeWithoutArgs([] {
+          return make_ready_future(Status{});
+        }));
+        ON_CALL(*stream, Cancel).WillByDefault([] {});
+        return std::unique_ptr<BidiReadStream>(std::move(stream));
+      });
+
+  auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  EXPECT_CALL(*mock_cq, MakeRelativeTimer)
+      .WillRepeatedly([](std::chrono::nanoseconds) {
+        return make_ready_future(
+            StatusOr<std::chrono::system_clock::time_point>(
+                std::chrono::system_clock::now()));
+      });
+  auto connection = std::make_shared<AsyncConnectionImpl>(
+      CompletionQueue(mock_cq), std::shared_ptr<GrpcChannelRefresh>(), mock,
+      TestOptions());
+
+  auto request = google::storage::v2::BidiReadObjectSpec{};
+  ASSERT_TRUE(TextFormat::ParseFromString(kExpectedRequestSpec, &request));
+  auto options =
+      connection->options().set<ReadRangesOption>({{0, 1024}, {2048, 4096}});
+  auto pending = connection->Open({std::move(request), std::move(options)});
+
+  auto next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Start");
+  next.first.set_value(true);
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Write");
+  next.first.set_value(true);
+
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Read");
+  next.first.set_value(true);
+
+  auto p = pending.get();
+  ASSERT_THAT(p, IsOkAndHolds(NotNull()));
+  auto descriptor = *std::move(p);
+
+  descriptor.reset();
+
+  auto last_read = sequencer.PopFrontWithName();
+  EXPECT_EQ(last_read.second, "Read[N]");
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Cancel");
+  next.first.set_value(true);
+  last_read.first.set_value(false);
+
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Finish");
+  next.first.set_value(true);
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal

google/cloud/storage/internal/async/object_descriptor_impl.cc

@@ -18,6 +18,7 @@
 #include "google/cloud/storage/internal/async/handle_redirect_error.h"
 #include "google/cloud/storage/internal/async/multi_stream_manager.h"
 #include "google/cloud/storage/internal/async/object_descriptor_reader_tracing.h"
+#include "google/cloud/storage/internal/async/options.h"
 #include "google/cloud/storage/internal/grpc/object_metadata_parser.h"
 #include "google/cloud/storage/internal/hash_function.h"
 #include "google/cloud/storage/internal/hash_function_impl.h"
@@ -28,6 +29,7 @@
 #include "google/cloud/grpc_error_delegate.h"
 #include "google/cloud/internal/opentelemetry.h"
 #include "google/rpc/status.pb.h"
+#include <algorithm>
 #include <limits>
 #include <memory>
 #include <utility>
@@ -52,6 +54,34 @@ ObjectDescriptorImpl::ObjectDescriptorImpl(
       []() -> std::shared_ptr<ReadStream> { return nullptr; },  // NOLINT
       std::make_shared<ReadStream>(std::move(stream),
                                    resume_policy_prototype_->clone()));
+  // If pre-warmed ranges are specified, initialize their `ReadRange` objects,
+  // register them as active on the initial stream, and cache them.
+  if (options_.has<ReadRangesOption>()) {
+    auto const& ranges = options_.get<ReadRangesOption>();
+    auto it = stream_manager_->GetFirstStream();
+    if (it != stream_manager_->End()) {
+      std::int64_t id = 0;
+      std::set<std::pair<std::int64_t, std::int64_t>> seen_ranges;
+      for (auto const& r : ranges) {
+        auto range_key = std::make_pair(r.offset, r.length);
+        if (!seen_ranges.insert(range_key).second) {
+          continue;  // Skip duplicate range.
+        }
+        ++id;
+        auto range = std::make_shared<ReadRange>(r.offset, r.length,
+                                                 read_object_spec_.bucket(),
+                                                 read_object_spec_.object());
+        // Registering on the stream allows `OnRead` to route incoming data to
+        // these ranges.
+        it->active_ranges.emplace(id, range);
+        // Cache them so subsequent `Read()` calls can claim them.
+        prewarmed_ranges_.emplace(range_key, PrewarmedRange{range, id});
+      }
+      // Ensure new dynamically requested ranges use IDs that don't conflict
+      // with pre-warmed ones.
+      read_id_generator_ = id;
+    }
+  }
 }
 
 ObjectDescriptorImpl::~ObjectDescriptorImpl() { Cancel(); }
@@ -193,6 +223,21 @@ std::unique_ptr<storage::AsyncReaderConnection> ObjectDescriptorImpl::Read(
       read_object_spec_.bucket(), read_object_spec_.object());
 
   std::unique_lock<std::mutex> lk(mu_);
+  // Check if this range matches a pre-warmed range.
+  auto cache_key = std::make_pair(p.start, p.length);
+  auto cache_it = prewarmed_ranges_.find(cache_key);
+  if (cache_it != prewarmed_ranges_.end()) {
+    // Cache hit. Claim the pre-warmed range and return it to the user.
+    auto prewarmed = std::move(cache_it->second);
+    prewarmed_ranges_.erase(cache_it);
+    lk.unlock();
+    if (!internal::TracingEnabled(options_)) {
+      return std::unique_ptr<storage::AsyncReaderConnection>(
+          std::make_unique<ObjectDescriptorReader>(std::move(prewarmed.range)));
+    }
+    return MakeTracingObjectDescriptorReader(std::move(prewarmed.range));
+  }
+
   if (stream_manager_->Empty()) {
     lk.unlock();
     range->OnFinish(Status(StatusCode::kFailedPrecondition,
@@ -366,9 +411,17 @@ void ObjectDescriptorImpl::OnRead(
     auto id = range_data.read_range().read_id();
     auto const l = copy.find(id);
     if (l == copy.end()) continue;
-    // TODO(#15104) - Consider returning if the range is done, and then
-    // skipping CleanupDoneRanges().
-    l->second->OnRead(std::move(range_data), is_transcoded, object_size);
+
+    auto range = l->second;
+    bool active = false;
+    lk.lock();
+    active = it->active_ranges.count(id) != 0;
+    lk.unlock();
+    if (active) {
+      // TODO(#15104) - Consider returning if the range is done, and then
+      // skipping CleanupDoneRanges().
+      range->OnRead(std::move(range_data), is_transcoded, object_size);
+    }
   }
   lk.lock();
   stream_manager_->CleanupDoneRanges(it);

google/cloud/storage/internal/async/object_descriptor_impl.h

@@ -27,9 +27,11 @@
 #include "google/storage/v2/storage.pb.h"
 #include <cstdint>
 #include <functional>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <optional>
+#include <set>
 #include <unordered_map>
 
 namespace google {
@@ -124,6 +126,15 @@ class ObjectDescriptorImpl
   std::optional<google::storage::v2::Object> metadata_;
   std::int64_t read_id_generator_ = 0;
 
+  // Information about a pre-warmed range that has not been claimed yet.
+  struct PrewarmedRange {
+    std::shared_ptr<ReadRange> range;
+    std::int64_t read_id;
+  };
+  // Cache of pre-warmed ranges, keyed by (offset, length).
+  std::map<std::pair<std::int64_t, std::int64_t>, PrewarmedRange>
+      prewarmed_ranges_;
+
   Options options_;
   std::unique_ptr<StreamManager> stream_manager_;
   // The future for the proactive background stream.