feat(bigtable): add support for PartialResultSetResume (#15637)

* feat(bigtable): add support for PartialResultSetResume

Author: mpeddada1

google/cloud/bigtable/CMakeLists.txt

@@ -206,6 +206,8 @@ add_library(
     internal/operation_context_factory.cc
     internal/operation_context_factory.h
     internal/partial_result_set_reader.h
+    internal/partial_result_set_resume.cc
+    internal/partial_result_set_resume.h
     internal/partial_result_set_source.cc
     internal/partial_result_set_source.h
     internal/prefix_range_end.cc
@@ -476,6 +478,7 @@ if (BUILD_TESTING)
         internal/mutate_rows_limiter_test.cc
         internal/operation_context_factory_test.cc
         internal/operation_context_test.cc
+        internal/partial_result_set_resume_test.cc
         internal/partial_result_set_source_test.cc
         internal/prefix_range_end_test.cc
         internal/rate_limiter_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

@@ -64,6 +64,7 @@ bigtable_client_unit_tests = [
     "internal/mutate_rows_limiter_test.cc",
     "internal/operation_context_factory_test.cc",
     "internal/operation_context_test.cc",
+    "internal/partial_result_set_resume_test.cc",
     "internal/partial_result_set_source_test.cc",
     "internal/prefix_range_end_test.cc",
     "internal/rate_limiter_test.cc",

google/cloud/bigtable/google_cloud_cpp_bigtable.bzl

@@ -102,6 +102,7 @@ google_cloud_cpp_bigtable_hdrs = [
     "internal/operation_context.h",
     "internal/operation_context_factory.h",
     "internal/partial_result_set_reader.h",
+    "internal/partial_result_set_resume.h",
     "internal/partial_result_set_source.h",
     "internal/prefix_range_end.h",
     "internal/rate_limiter.h",
@@ -214,6 +215,7 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/mutate_rows_limiter.cc",
     "internal/operation_context.cc",
     "internal/operation_context_factory.cc",
+    "internal/partial_result_set_resume.cc",
     "internal/partial_result_set_source.cc",
     "internal/prefix_range_end.cc",
     "internal/rate_limiter.cc",

google/cloud/bigtable/internal/partial_result_set_resume.cc

@@ -0,0 +1,69 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/bigtable/internal/partial_result_set_resume.h"
+#include <thread>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+void PartialResultSetResume::TryCancel() { reader_->TryCancel(); }
+
+bool PartialResultSetResume::Read(
+    absl::optional<std::string> const& resume_token,
+    UnownedPartialResultSet& result) {
+  bool resumption = false;
+  do {
+    if (reader_->Read(resume_token, result)) {
+      // Let the caller know if we recreated the PartialResultSetReader using
+      // the resume_token so that they might discard any previous results that
+      // will be contained in the new stream.
+      if (resumption) result.resumption = true;
+      return true;
+    }
+    auto status = Finish();
+    if (status.ok()) return false;
+    if (!resume_token) {
+      // Our caller has requested that we not try to resume the stream,
+      // probably because they have already delivered previous results that
+      // would otherwise be replayed.
+      return false;
+    }
+    if (idempotency_ == google::cloud::Idempotency::kNonIdempotent ||
+        !retry_policy_prototype_->OnFailure(status)) {
+      return false;
+    }
+    std::this_thread::sleep_for(backoff_policy_prototype_->OnCompletion());
+    resumption = true;
+    last_status_.reset();
+    reader_ = factory_(*resume_token);
+  } while (!retry_policy_prototype_->IsExhausted());
+  return false;
+}
+
+Status PartialResultSetResume::Finish() {
+  // Finish() can be called only once, so cache the last result.
+  if (last_status_.has_value()) {
+    return *last_status_;
+  }
+  last_status_ = reader_->Finish();
+  return *last_status_;
+}
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google

google/cloud/bigtable/internal/partial_result_set_resume.h

@@ -0,0 +1,74 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_PARTIAL_RESULT_SET_RESUME_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_PARTIAL_RESULT_SET_RESUME_H
+
+#include "google/cloud/bigtable/internal/partial_result_set_reader.h"
+#include "google/cloud/bigtable/rpc_backoff_policy.h"
+#include "google/cloud/bigtable/rpc_retry_policy.h"
+#include "google/cloud/bigtable/version.h"
+#include "google/cloud/backoff_policy.h"
+#include "absl/types/optional.h"
+#include <functional>
+#include <memory>
+#include <string>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+/// Create a new PartialResultSetReader given a resume token value.
+using PartialResultSetReaderFactory =
+    std::function<std::unique_ptr<PartialResultSetReader>(std::string)>;
+
+/**
+ * A PartialResultSetReader that resumes the streaming RPC on retryable errors.
+ */
+class PartialResultSetResume : public PartialResultSetReader {
+ public:
+  PartialResultSetResume(
+      PartialResultSetReaderFactory factory,
+      google::cloud::Idempotency idempotency,
+      std::unique_ptr<bigtable::RPCRetryPolicy> retry_policy,
+      std::unique_ptr<bigtable::RPCBackoffPolicy> backoff_policy)
+      : factory_(std::move(factory)),
+        idempotency_(idempotency),
+        retry_policy_prototype_(std::move(retry_policy)),
+        backoff_policy_prototype_(std::move(backoff_policy)),
+        reader_(factory_(std::string{})) {}
+
+  ~PartialResultSetResume() override = default;
+
+  void TryCancel() override;
+  bool Read(absl::optional<std::string> const& resume_token,
+            UnownedPartialResultSet& result) override;
+  Status Finish() override;
+
+ private:
+  PartialResultSetReaderFactory factory_;
+  google::cloud::Idempotency idempotency_;
+  std::unique_ptr<bigtable::RPCRetryPolicy> retry_policy_prototype_;
+  std::unique_ptr<bigtable::RPCBackoffPolicy> backoff_policy_prototype_;
+  std::unique_ptr<PartialResultSetReader> reader_;
+  absl::optional<Status> last_status_;
+};
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_PARTIAL_RESULT_SET_RESUME_H