feat(bigtable): add AsyncWaitForConsistency() helper for Table Admin (#9310)

dbolduc · web-flow · commit a85d425f1a83 · 2022-06-20T16:09:15.000-04:00
diff --git a/google/cloud/bigtable/CMakeLists.txt b/google/cloud/bigtable/CMakeLists.txt
@@ -205,8 +205,6 @@ add_library(
     internal/rpc_policy_parameters.h
     internal/rpc_policy_parameters.inc
     internal/unary_client_utils.h
-    internal/wait_for_consistency.cc
-    internal/wait_for_consistency.h
     metadata_update_policy.cc
     metadata_update_policy.h
     mutation_batcher.cc
@@ -241,7 +239,9 @@ add_library(
     table_config.h
     version.cc
     version.h
-    version_info.h)
+    version_info.h
+    wait_for_consistency.cc
+    wait_for_consistency.h)
 target_link_libraries(
     google_cloud_cpp_bigtable
     PUBLIC absl::memory
@@ -356,7 +356,6 @@ if (BUILD_TESTING)
         internal/legacy_row_reader_test.cc
         internal/logging_data_client_test.cc
         internal/prefix_range_end_test.cc
-        internal/wait_for_consistency_test.cc
         legacy_table_test.cc
         metadata_update_policy_test.cc
         mocks/mock_row_reader_test.cc
@@ -381,7 +380,8 @@ if (BUILD_TESTING)
         table_sample_row_keys_test.cc
         table_test.cc
         testing/cleanup_stale_resources_test.cc
-        testing/random_names_test.cc)
+        testing/random_names_test.cc
+        wait_for_consistency_test.cc)
 
     # Export the list of unit tests so the Bazel BUILD file can pick it up.
     export_list_to_bazel("bigtable_client_unit_tests.bzl"
diff --git a/google/cloud/bigtable/bigtable_client_unit_tests.bzl b/google/cloud/bigtable/bigtable_client_unit_tests.bzl
@@ -58,7 +58,6 @@ bigtable_client_unit_tests = [
     "internal/legacy_row_reader_test.cc",
     "internal/logging_data_client_test.cc",
     "internal/prefix_range_end_test.cc",
-    "internal/wait_for_consistency_test.cc",
     "legacy_table_test.cc",
     "metadata_update_policy_test.cc",
     "mocks/mock_row_reader_test.cc",
@@ -84,4 +83,5 @@ bigtable_client_unit_tests = [
     "table_test.cc",
     "testing/cleanup_stale_resources_test.cc",
     "testing/random_names_test.cc",
+    "wait_for_consistency_test.cc",
 ]
diff --git a/google/cloud/bigtable/examples/table_admin_snippets.cc b/google/cloud/bigtable/examples/table_admin_snippets.cc
@@ -17,6 +17,7 @@
 #include "google/cloud/bigtable/table_admin.h"
 #include "google/cloud/bigtable/testing/cleanup_stale_resources.h"
 #include "google/cloud/bigtable/testing/random_names.h"
+#include "google/cloud/bigtable/wait_for_consistency.h"
 #include "google/cloud/internal/getenv.h"
 #include "google/cloud/log.h"
 #include <chrono>
@@ -622,38 +623,43 @@ void DropRowsByPrefix(
   (std::move(admin), argv.at(0), argv.at(1), argv.at(2), argv.at(3));
 }
 
-// TODO(#7732) - update this sample to use the helper method
 void WaitForConsistencyCheck(
-    google::cloud::bigtable_admin::BigtableTableAdminClient const&,
+    google::cloud::bigtable_admin::BigtableTableAdminClient admin,
     std::vector<std::string> const& argv) {
-  auto old_admin = google::cloud::bigtable::TableAdmin(
-      google::cloud::bigtable::MakeAdminClient(argv.at(0)), argv.at(1));
-
   //! [wait for consistency check]
   namespace cbt = ::google::cloud::bigtable;
+  namespace cbta = ::google::cloud::bigtable_admin;
+  using ::google::cloud::CompletionQueue;
   using ::google::cloud::future;
+  using ::google::cloud::Status;
   using ::google::cloud::StatusOr;
-  [](cbt::TableAdmin admin, std::string const& table_id) {
-    StatusOr<std::string> consistency_token =
-        admin.GenerateConsistencyToken(table_id);
+  [](cbta::BigtableTableAdminClient admin, std::string const& project_id,
+     std::string const& instance_id, std::string const& table_id) {
+    std::string table_name = cbt::TableName(project_id, instance_id, table_id);
+    StatusOr<google::bigtable::admin::v2::GenerateConsistencyTokenResponse>
+        consistency_token = admin.GenerateConsistencyToken(table_name);
     if (!consistency_token) {
       throw std::runtime_error(consistency_token.status().message());
     }
-    future<StatusOr<cbt::Consistency>> consistent_future =
-        admin.WaitForConsistency(table_id, *consistency_token);
-    future<void> fut = consistent_future.then(
-        [&consistency_token](future<StatusOr<cbt::Consistency>> f) {
-          auto is_consistent = f.get();
-          if (!is_consistent) {
-            throw std::runtime_error(is_consistent.status().message());
-          }
-          std::cout << "Table is consistent with token " << *consistency_token
-                    << "\n";
-        });
-    fut.get();  // simplify example by blocking until operation is done.
+    // Start a thread to perform the background work.
+    CompletionQueue cq;
+    std::thread cq_runner([&cq] { cq.Run(); });
+
+    std::string token = consistency_token->consistency_token();
+    future<Status> consistent_future =
+        cbta::AsyncWaitForConsistency(cq, admin, table_name, token);
+
+    // Simplify the example by blocking until the operation is done.
+    Status status = consistent_future.get();
+    if (!status.ok()) throw std::runtime_error(status.message());
+    std::cout << "Table is consistent with token " << token << "\n";
+
+    // Shutdown the work queue and join the background thread
+    cq.Shutdown();
+    cq_runner.join();
   }
   //! [wait for consistency check]
-  (std::move(old_admin), argv.at(2));
+  (std::move(admin), argv.at(0), argv.at(1), argv.at(2));
 }
 
 void CheckConsistency(
diff --git a/google/cloud/bigtable/google_cloud_cpp_bigtable.bzl b/google/cloud/bigtable/google_cloud_cpp_bigtable.bzl
@@ -95,7 +95,6 @@ google_cloud_cpp_bigtable_hdrs = [
     "internal/rpc_policy_parameters.h",
     "internal/rpc_policy_parameters.inc",
     "internal/unary_client_utils.h",
-    "internal/wait_for_consistency.h",
     "metadata_update_policy.h",
     "mutation_batcher.h",
     "mutation_branch.h",
@@ -117,6 +116,7 @@ google_cloud_cpp_bigtable_hdrs = [
     "table_config.h",
     "version.h",
     "version_info.h",
+    "wait_for_consistency.h",
 ]
 
 google_cloud_cpp_bigtable_srcs = [
@@ -179,7 +179,6 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/logging_data_client.cc",
     "internal/prefix_range_end.cc",
     "internal/readrowsparser.cc",
-    "internal/wait_for_consistency.cc",
     "metadata_update_policy.cc",
     "mutation_batcher.cc",
     "mutations.cc",
@@ -194,4 +193,5 @@ google_cloud_cpp_bigtable_srcs = [
     "table_admin.cc",
     "table_config.cc",
     "version.cc",
+    "wait_for_consistency.cc",
 ]
diff --git a/google/cloud/bigtable/table_admin.cc b/google/cloud/bigtable/table_admin.cc
@@ -13,7 +13,8 @@
 // limitations under the License.
 
 #include "google/cloud/bigtable/table_admin.h"
-#include "google/cloud/bigtable/internal/wait_for_consistency.h"
+#include "google/cloud/bigtable/admin/bigtable_table_admin_client.h"
+#include "google/cloud/bigtable/wait_for_consistency.h"
 #include "google/cloud/internal/retry_policy.h"
 #include "google/cloud/internal/time_utils.h"
 #include <google/protobuf/duration.pb.h>
@@ -226,9 +227,16 @@ Status TableAdmin::DropRowsByPrefix(std::string const& table_id,
 
 future<StatusOr<Consistency>> TableAdmin::WaitForConsistency(
     std::string const& table_id, std::string const& consistency_token) {
-  return bigtable_admin_internal::AsyncWaitForConsistency(
-             cq_, connection_, TableName(table_id), consistency_token,
-             policies_)
+  // We avoid lifetime issues due to ownership cycles, by holding the
+  // `BackgroundThreads` which run the `CompletionQueue` outside of the
+  // operation, in this class. If the `BackgroundThreads` running the
+  // `CompletionQueue` were instead owned by the Connection, we would have an
+  // ownership cycle. We have made this mistake before. See #7740 for more
+  // details.
+  auto client = bigtable_admin::BigtableTableAdminClient(connection_);
+  return bigtable_admin::AsyncWaitForConsistency(cq_, std::move(client),
+                                                 TableName(table_id),
+                                                 consistency_token, policies_)
       .then([](future<Status> f) -> StatusOr<Consistency> {
         auto s = f.get();
         if (!s.ok()) return s;
diff --git a/google/cloud/bigtable/wait_for_consistency.cc b/google/cloud/bigtable/wait_for_consistency.cc
@@ -12,31 +12,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "google/cloud/bigtable/internal/wait_for_consistency.h"
+#include "google/cloud/bigtable/wait_for_consistency.h"
 #include "google/cloud/bigtable/admin/bigtable_table_admin_options.h"
 #include "google/cloud/bigtable/admin/internal/bigtable_table_admin_option_defaults.h"
 #include <chrono>
 
 namespace google {
 namespace cloud {
-namespace bigtable_admin_internal {
+namespace bigtable_admin {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 namespace {
 
 // This class borrows heavily from `google::cloud::internal::AsyncRetryLoop`
 class AsyncWaitForConsistencyImpl
     : public std::enable_shared_from_this<AsyncWaitForConsistencyImpl> {
  public:
-  AsyncWaitForConsistencyImpl(
-      CompletionQueue cq,
-      std::shared_ptr<bigtable_admin::BigtableTableAdminConnection> connection,
-      std::string table_name, std::string consistency_token, Options options)
+  AsyncWaitForConsistencyImpl(CompletionQueue cq,
+                              BigtableTableAdminClient client,
+                              std::string table_name,
+                              std::string consistency_token, Options options)
       : cq_(std::move(cq)),
-        connection_(std::move(connection)),
-        options_(BigtableTableAdminDefaultOptions(std::move(options))),
-        polling_policy_(
-            options_
-                .get<bigtable_admin::BigtableTableAdminPollingPolicyOption>()) {
+        client_(std::move(client)),
+        options_(bigtable_admin_internal::BigtableTableAdminDefaultOptions(
+            std::move(options))),
+        polling_policy_(options_.get<BigtableTableAdminPollingPolicyOption>()) {
     request_.set_name(std::move(table_name));
     request_.set_consistency_token(std::move(consistency_token));
   }
@@ -73,13 +72,13 @@ class AsyncWaitForConsistencyImpl
   }
 
   void StartAttempt() {
-    internal::OptionsSpan span(options_);
     auto self = this->shared_from_this();
     auto state = StartOperation();
     if (state.cancelled) return;
-    SetPending(state.operation,
-               connection_->AsyncCheckConsistency(request_).then(
-                   [self](future<RespType> f) { self->OnAttempt(f.get()); }));
+    SetPending(
+        state.operation,
+        client_.AsyncCheckConsistency(request_, options_)
+            .then([self](future<RespType> f) { self->OnAttempt(f.get()); }));
   }
 
   void StartBackoff() {
@@ -152,7 +151,7 @@ class AsyncWaitForConsistencyImpl
 
   CompletionQueue cq_;
   bigtable::admin::v2::CheckConsistencyRequest request_;
-  std::shared_ptr<bigtable_admin::BigtableTableAdminConnection> connection_;
+  BigtableTableAdminClient client_;
   Options options_;
   std::shared_ptr<PollingPolicy> polling_policy_;
   promise<Status> result_;
@@ -169,17 +168,18 @@ class AsyncWaitForConsistencyImpl
 
 }  // namespace
 
-future<Status> AsyncWaitForConsistency(
-    CompletionQueue cq,
-    std::shared_ptr<bigtable_admin::BigtableTableAdminConnection> connection,
-    std::string table_name, std::string consistency_token, Options options) {
+future<Status> AsyncWaitForConsistency(CompletionQueue cq,
+                                       BigtableTableAdminClient client,
+                                       std::string table_name,
+                                       std::string consistency_token,
+                                       Options options) {
   auto loop = std::make_shared<AsyncWaitForConsistencyImpl>(
-      std::move(cq), std::move(connection), std::move(table_name),
+      std::move(cq), std::move(client), std::move(table_name),
       std::move(consistency_token), std::move(options));
   return loop->Start();
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
-}  // namespace bigtable_admin_internal
+}  // namespace bigtable_admin
 }  // namespace cloud
 }  // namespace google
diff --git a/google/cloud/bigtable/wait_for_consistency.h b/google/cloud/bigtable/wait_for_consistency.h
@@ -12,18 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_WAIT_FOR_CONSISTENCY_H
-#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_WAIT_FOR_CONSISTENCY_H
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_WAIT_FOR_CONSISTENCY_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_WAIT_FOR_CONSISTENCY_H
 
-#include "google/cloud/bigtable/admin/bigtable_table_admin_connection.h"
+#include "google/cloud/bigtable/admin/bigtable_table_admin_client.h"
 
 namespace google {
 namespace cloud {
-namespace bigtable_admin_internal {
+namespace bigtable_admin {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 /**
- * Checks consistency of a table with multiple calls using background threads
+ * Polls until a table is consistent, or until the polling policy has expired.
  *
  * @param cq the completion queue that will execute the asynchronous
  *    calls. The application must ensure that one or more threads are
@@ -34,19 +34,20 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  * @param consistency_token the consistency token of the table.
  * @param options (optional) configuration options. Users who wish to modify the
  *     default polling behavior can supply a custom polling policy with
- *     `BigtableTableAdminPollingPolicyOption`.
+ *     `BigtableTableAdminPollingPolicyOption`. Note that the client's polling
+ *     policy is not used for this operation.
  * @return the consistency status for the table. The status is OK if and only if
  *     the table is consistent.
  */
-future<Status> AsyncWaitForConsistency(
-    CompletionQueue cq,
-    std::shared_ptr<bigtable_admin::BigtableTableAdminConnection> connection,
-    std::string table_name, std::string consistency_token,
-    Options options = {});
+future<Status> AsyncWaitForConsistency(CompletionQueue cq,
+                                       BigtableTableAdminClient client,
+                                       std::string table_name,
+                                       std::string consistency_token,
+                                       Options options = {});
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
-}  // namespace bigtable_admin_internal
+}  // namespace bigtable_admin
 }  // namespace cloud
 }  // namespace google
 
-#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_WAIT_FOR_CONSISTENCY_H
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_WAIT_FOR_CONSISTENCY_H
diff --git a/google/cloud/bigtable/wait_for_consistency_test.cc b/google/cloud/bigtable/wait_for_consistency_test.cc
