feat(storage): expose initial read ranges in AsyncClient (#16341)

Author: kalragauri

google/cloud/storage/async/client.cc

@@ -16,6 +16,7 @@
 #include "google/cloud/storage/internal/async/connection_impl.h"
 #include "google/cloud/storage/internal/async/connection_tracing.h"
 #include "google/cloud/storage/internal/async/default_options.h"
+#include "google/cloud/storage/internal/async/options.h"
 #include "google/cloud/storage/internal/grpc/stub.h"
 #include "google/cloud/grpc_options.h"
 #include <memory>
@@ -84,6 +85,27 @@ future<StatusOr<ObjectDescriptor>> AsyncClient::Open(
       });
 }
 
+future<StatusOr<ObjectDescriptor>> AsyncClient::Open(
+    BucketName const& bucket_name, std::string object_name,
+    InitialReadRanges const& config, Options opts) {
+  auto spec = google::storage::v2::BidiReadObjectSpec{};
+  spec.set_bucket(bucket_name.FullName());
+  spec.set_object(std::move(object_name));
+
+  // Convert the user-facing `InitialReadRanges` to the internal
+  // `ReadRangesOption` so it can be propagated down to the connection
+  // implementation.
+  if (!config.initial_ranges.empty()) {
+    std::vector<storage_internal::ReadRangeConfig> internal_ranges;
+    internal_ranges.reserve(config.initial_ranges.size());
+    for (auto const& r : config.initial_ranges) {
+      internal_ranges.push_back({r.offset, r.length});
+    }
+    opts.set<storage_internal::ReadRangesOption>(std::move(internal_ranges));
+  }
+  return Open(std::move(spec), std::move(opts));
+}
+
 future<StatusOr<std::pair<AsyncReader, AsyncToken>>> AsyncClient::ReadObject(
     BucketName const& bucket_name, std::string object_name, Options opts) {
   auto request = google::storage::v2::ReadObjectRequest{};

google/cloud/storage/async/client.h

@@ -83,6 +83,27 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  */
 class AsyncClient {
  public:
+  /**
+   * Specifies a byte range for a read request.
+   */
+  struct ByteRange {
+    std::int64_t offset = 0;
+    std::int64_t length = 0;
+  };
+
+  /**
+   * Specifies initial byte ranges to request concurrently when opening an
+   * object.
+   *
+   * Passing initial read ranges allows the client to begin fetching expected
+   * byte ranges during connection setup, which may improve first-byte retrieval
+   * times for known access patterns.
+   */
+  struct InitialReadRanges {
+    // The initial read ranges.
+    std::vector<ByteRange> initial_ranges;
+  };
+
   /// Create a new client configured with @p options.
   explicit AsyncClient(Options options = {});
   /// Create a new client using @p connection. This is often used for mocking.
@@ -273,6 +294,29 @@ class AsyncClient {
                                           std::string object_name,
                                           Options opts = {});
 
+  /**
+   * Open an object descriptor, requesting specified initial read ranges
+   * concurrently.
+   *
+   * @par Example
+   * @snippet storage_async_samples.cc open-object-initial-read-ranges
+   *
+   * @par Idempotency
+   * This is a read-only operation and is always idempotent. The operation will
+   * retry until the descriptor is successfully created. The descriptor itself
+   * will resume any incomplete ranged reads if the connection(s) are
+   * interrupted. Use `ResumePolicyOption` and `ResumePolicy` to control this.
+   *
+   * @param bucket_name the name of the bucket that contains the object.
+   * @param object_name the name of the object to be read.
+   * @param config initial byte ranges to request during connection setup.
+   * @param opts options controlling the behavior of this RPC.
+   */
+  future<StatusOr<ObjectDescriptor>> Open(BucketName const& bucket_name,
+                                          std::string object_name,
+                                          InitialReadRanges const& config,
+                                          Options opts = {});
+
   /**
    * Open an object descriptor to perform one or more ranged reads.
    *

google/cloud/storage/async/client_test.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/storage/async/client.h"
+#include "google/cloud/storage/internal/async/options.h"
 #include "google/cloud/storage/mocks/mock_async_connection.h"
 #include "google/cloud/storage/mocks/mock_async_object_descriptor_connection.h"
 #include "google/cloud/storage/mocks/mock_async_reader_connection.h"
@@ -290,6 +291,43 @@ TEST(AsyncClient, Open) {
              "empty response", [](auto const& p) { return p.size(); }, 0)));
 }
 
+TEST(AsyncClient, OpenWithInitialReadRanges) {
+  auto constexpr kExpectedRequest = R"pb(
+    bucket: "projects/_/buckets/test-bucket"
+    object: "test-object"
+  )pb";
+  auto mock = std::make_shared<MockAsyncConnection>();
+  EXPECT_CALL(*mock, options).WillRepeatedly(Return(Options{}));
+
+  EXPECT_CALL(*mock, Open).WillOnce([&](AsyncConnection::OpenParams const& p) {
+    EXPECT_TRUE(p.options.has<storage_internal::ReadRangesOption>());
+    auto const& ranges = p.options.get<storage_internal::ReadRangesOption>();
+    EXPECT_EQ(ranges.size(), 2);
+    if (ranges.size() >= 2) {
+      EXPECT_EQ(ranges[0].offset, 0);
+      EXPECT_EQ(ranges[0].length, 100);
+      EXPECT_EQ(ranges[1].offset, 1000);
+      EXPECT_EQ(ranges[1].length, 200);
+    }
+
+    auto expected = google::storage::v2::BidiReadObjectSpec{};
+    EXPECT_TRUE(TextFormat::ParseFromString(kExpectedRequest, &expected));
+    EXPECT_THAT(p.read_spec, IsProtoEqual(expected));
+
+    auto descriptor = std::make_shared<MockAsyncObjectDescriptorConnection>();
+    return make_ready_future(make_status_or(
+        std::shared_ptr<ObjectDescriptorConnection>(std::move(descriptor))));
+  });
+
+  auto client = AsyncClient(mock);
+  AsyncClient::InitialReadRanges config;
+  config.initial_ranges = {{0, 100}, {1000, 200}};
+  auto descriptor =
+      client.Open(BucketName("test-bucket"), "test-object", std::move(config))
+          .get();
+  ASSERT_STATUS_OK(descriptor);
+}
+
 TEST(AsyncClient, OpenWithInvalidBucket) {
   auto constexpr kExpectedRequest = R"pb(
     bucket: "test-only-invalid"

google/cloud/storage/examples/storage_async_samples.cc

@@ -246,6 +246,55 @@ void OpenObjectMultipleRangedRead(google::cloud::storage::AsyncClient& client,
   std::cout << "The ranges contain " << count << " newlines\n";
 }
 
+void OpenObjectWithInitialReadRanges(
+    google::cloud::storage::AsyncClient& client,
+    std::vector<std::string> const& argv) {
+  //! [open-object-initial-read-ranges]
+  // [START storage_open_object_initial_read_ranges]
+  namespace gcs = google::cloud::storage;
+
+  // Helper coroutine to count newlines returned by an AsyncReader.
+  auto count_newlines =
+      [](gcs::AsyncReader reader,
+         gcs::AsyncToken token) -> google::cloud::future<std::uint64_t> {
+    std::uint64_t count = 0;
+    while (token.valid()) {
+      auto [payload, t] = (co_await reader.Read(std::move(token))).value();
+      token = std::move(t);
+      for (auto const& buffer : payload.contents()) {
+        count += std::count(buffer.begin(), buffer.end(), '\n');
+      }
+    }
+    co_return count;
+  };
+
+  auto coro =
+      [&count_newlines](
+          gcs::AsyncClient& client, std::string bucket_name,
+          std::string object_name) -> google::cloud::future<std::uint64_t> {
+    gcs::AsyncClient::InitialReadRanges config;
+    config.initial_ranges = {{0, 1024}, {1024, 1024}};
+
+    auto descriptor =
+        (co_await client.Open(gcs::BucketName(std::move(bucket_name)),
+                              std::move(object_name), std::move(config)))
+            .value();
+
+    auto [r1, t1] = descriptor.Read(0, 1024);
+    auto [r2, t2] = descriptor.Read(1024, 1024);
+
+    auto c1 = count_newlines(std::move(r1), std::move(t1));
+    auto c2 = count_newlines(std::move(r2), std::move(t2));
+    co_return (co_await std::move(c1)) + (co_await std::move(c2));
+  };
+  // [END storage_open_object_initial_read_ranges]
+  //! [open-object-initial-read-ranges]
+  // The example is easier to test and run if we call the coroutine and block
+  // until it completes.
+  auto const count = coro(client, argv.at(0), argv.at(1)).get();
+  std::cout << "The pre-warmed ranges contain " << count << " newlines\n";
+}
+
 void OpenObjectReadFullObject(google::cloud::storage::AsyncClient& client,
                               std::vector<std::string> const& argv) {
   //! [open-object-read-full-object]
@@ -997,6 +1046,11 @@ void OpenObjectMultipleRangedRead(google::cloud::storage::AsyncClient&,
   std::cerr << "AsyncClient::Open() example requires coroutines\n";
 }
 
+void OpenObjectWithInitialReadRanges(google::cloud::storage::AsyncClient&,
+                                     std::vector<std::string> const&) {
+  std::cerr << "AsyncClient::Open() example requires coroutines\n";
+}
+
 void OpenMultipleObjectsRangedRead(google::cloud::storage::AsyncClient&,
                                    std::vector<std::string> const&) {
   std::cerr << "AsyncClient::Open() example requires coroutines\n";
@@ -1306,6 +1360,10 @@ void AutoRun(std::vector<std::string> const& argv) {
             << std::endl;
   OpenObjectMultipleRangedRead(client, {bucket_name, composed_name});
 
+  std::cout << "Running the OpenObjectWithInitialReadRanges() example"
+            << std::endl;
+  OpenObjectWithInitialReadRanges(client, {bucket_name, composed_name});
+
   std::cout << "Running the OpenMultipleObjectsRangedRead() example"
             << std::endl;
   auto const multi_read_o1 =
@@ -1563,6 +1621,8 @@ int main(int argc, char* argv[]) try {
                  OpenObjectSingleRangedRead),
       make_entry("open-object-multiple-ranged-read", {},
                  OpenObjectMultipleRangedRead),
+      make_entry("open-object-initial-read-ranges", {},
+                 OpenObjectWithInitialReadRanges),
       make_entry("open-object-read-full-object", {}, OpenObjectReadFullObject),
       make_entry("open-multiple-objects-ranged-read",
                  {"<object-name-1>", "<object-name-2>", "<object-name-3>"},