impl(bigtable): improve query execution refresh and retry (#15788)

Author: scotthart

google/cloud/bigtable/ci/run_conformance_tests_proxy_bazel.sh

@@ -47,36 +47,14 @@ go test -v \
   -proxy_addr=:9999
 exit_status=$?
 
-# Run all the ExecuteQuery tests that either work or we plan to skip such as
-# CloseClient
+# Run all the ExecuteQuery tests that either work or we plan to skip due to
+# unimplemented features or issues with the tests themselves.
 go test -v \
-  -run "TestExecuteQuery|TestExecuteQuery_PlanRefresh$|TestExecuteQuery_PlanRefresh_WithMetadataChange|TestExecuteQuery_PlanRefresh_Retries|TestExecuteQuery_FailsOnSuccesfulStreamWithNoToken" \
-  -skip "CloseClient|FailsOnInvalidType|FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField|TestExecuteQuery_PlanRefresh_AfterResumeTokenCausesError|TestExecuteQuery_RetryTest_WithPlanRefresh|TestExecuteQuery_PlanRefresh_RespectsDeadline|TestExecuteQuery_PlanRefresh_RecoversAfterPermanentError" \
+  -run "TestExecuteQuery" \
+  -skip "CloseClient|TestExecuteQuery_FailsOnInvalidType|TestExecuteQuery_PlanRefresh_RespectsDeadline" \
   -proxy_addr=:9999
 exit_status=$?
 
-# These next four go test commands group the currently failing ExecuteQuery
-# tests into groups that exercise similar functionality and should be worked on
-# together.
-
-# Metadata tests b/461232934
-#go test -v \
-#  -run "FailsOnInvalidType" \
-#  -proxy_addr=:9999
-#exit_status=$?
-
-# Response/Metadata mismatches b/461233335
-go test -v \
-  -run "FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField" \
-  -proxy_addr=:9999
-exit_status=$?
-
-# QueryPlan refresh tests b/461233613
-#go test -v \
-#  -run "RetryTest_WithPlanRefresh|PlanRefresh|PlanRefresh_RecoversAfterPermanentError" \
-#  -proxy_addr=:9999
-#exit_status=$?
-
 # Remove the entire module cache, including unpacked source code of versioned
 # dependencies.
 go clean -modcache

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -39,7 +39,6 @@
 #include "google/cloud/internal/random.h"
 #include "google/cloud/internal/retry_loop.h"
 #include "google/cloud/internal/streaming_read_rpc.h"
-#include "absl/functional/bind_front.h"
 #include <memory>
 #include <string>
 
@@ -59,9 +58,16 @@ inline std::unique_ptr<bigtable::DataRetryPolicy> retry_policy(
 }
 
 inline std::unique_ptr<bigtable::DataRetryPolicy>
-query_plan_refresh_retry_policy(Options const& options) {
+execute_query_plan_refresh_retry_policy(Options const& options) {
   return options
-      .get<bigtable::experimental::QueryPlanRefreshRetryPolicyOption>()
+      .get<bigtable::experimental::ExecuteQueryPlanRefreshRetryPolicyOption>()
+      ->clone();
+}
+
+inline std::unique_ptr<bigtable::DataRetryPolicy>
+query_plan_refresh_function_retry_policy(Options const& options) {
+  return options
+      .get<bigtable::experimental::QueryPlanRefreshFunctionRetryPolicyOption>()
       ->clone();
 }
 
@@ -111,10 +117,6 @@ bool IsStatusMetadataIndicatingRetryPolicyExhausted(Status const& status) {
                                                 "retry-policy-exhausted"));
 }
 
-bool IsStatusIndicatingInternalError(Status const& status) {
-  return status.code() == StatusCode::kInternal;
-}
-
 class DefaultPartialResultSetReader
     : public bigtable_internal::PartialResultSetReader {
  public:
@@ -760,7 +762,7 @@ StatusOr<bigtable::PreparedQuery> DataConnectionImpl::PrepareQuery(
   auto const* func = __func__;
   auto refresh_fn = [this, request, func]() mutable {
     auto current = google::cloud::internal::SaveCurrentOptions();
-    auto retry = retry_policy(*current);
+    auto retry = query_plan_refresh_function_retry_policy(*current);
     auto backoff = backoff_policy(*current);
     auto operation_context = operation_context_factory_->PrepareQuery(
         request.instance_name(), app_profile_id(*current));
@@ -843,7 +845,7 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
 
         auto refresh_fn = [this, request, func]() mutable {
           auto current = google::cloud::internal::SaveCurrentOptions();
-          auto retry = retry_policy(*current);
+          auto retry = query_plan_refresh_function_retry_policy(*current);
           auto backoff = backoff_policy(*current);
           auto operation_context = operation_context_factory_->PrepareQuery(
               request.instance_name(), app_profile_id(*current));
@@ -882,6 +884,139 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
       });
 }
 
+class QueryPlanRefreshingPartialResultSource
+    : public bigtable::ResultSourceInterface {
+ public:
+  using SourceFactoryFn =
+      std::function<StatusOr<std::unique_ptr<bigtable::ResultSourceInterface>>(
+          google::bigtable::v2::ExecuteQueryRequest& request,
+          google::bigtable::v2::ResultSetMetadata metadata,
+          std::unique_ptr<bigtable::DataRetryPolicy> retry_policy_prototype,
+          std::unique_ptr<BackoffPolicy> backoff_policy_prototype,
+          std::shared_ptr<OperationContext> const& operation_context)>;
+  static std::unique_ptr<QueryPlanRefreshingPartialResultSource> Create(
+      google::bigtable::v2::ExecuteQueryRequest request,
+      SourceFactoryFn source_factory, std::shared_ptr<QueryPlan> query_plan,
+      std::shared_ptr<OperationContext> operation_context,
+      std::unique_ptr<BackoffPolicy> backoff_policy,
+      std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
+      std::unique_ptr<bigtable::DataRetryPolicy>
+          query_plan_refresh_retry_policy,
+      internal::ImmutableOptions options) {
+    return std::unique_ptr<QueryPlanRefreshingPartialResultSource>(
+        new QueryPlanRefreshingPartialResultSource(
+            std::move(request), std::move(source_factory),
+            std::move(query_plan), std::move(operation_context),
+            std::move(backoff_policy), std::move(retry_policy),
+            std::move(query_plan_refresh_retry_policy), std::move(options)));
+  }
+
+  ~QueryPlanRefreshingPartialResultSource() override = default;
+
+  StatusOr<bigtable::QueryRow> NextRow() override {
+    StatusOr<bigtable::QueryRow> row;
+    do {
+      if (!query_plan_valid_) source_ = absl::nullopt;
+      if (!source_.has_value() || !source_->ok()) {
+        UpdateSource();
+      }
+      row = (**source_)->NextRow();
+      if (ExecuteQueryPlanRefreshRetry::IsQueryPlanExpired(row.status())) {
+        query_plan_valid_ = false;
+        query_plan_->Invalidate(row.status(),
+                                query_plan_data_->prepared_query());
+      }
+    } while (!query_plan_refresh_retry_policy_->IsExhausted() &&
+             ExecuteQueryPlanRefreshRetry::IsQueryPlanExpired(row.status()));
+    return row;
+  }
+
+  absl::optional<google::bigtable::v2::ResultSetMetadata> Metadata() override {
+    if (!source_.has_value()) return absl::nullopt;
+    return (**source_)->Metadata();
+  }
+
+ private:
+  QueryPlanRefreshingPartialResultSource(
+      google::bigtable::v2::ExecuteQueryRequest request,
+      SourceFactoryFn source_factory, std::shared_ptr<QueryPlan> query_plan,
+      std::shared_ptr<OperationContext> operation_context,
+      std::unique_ptr<BackoffPolicy> backoff_policy,
+      std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
+      std::unique_ptr<bigtable::DataRetryPolicy>
+          query_plan_refresh_retry_policy,
+      internal::ImmutableOptions options)
+      : request_(std::move(request)),
+        source_factory_(std::move(source_factory)),
+        query_plan_(std::move(query_plan)),
+        operation_context_(std::move(operation_context)),
+        backoff_policy_(std::move(backoff_policy)),
+        retry_policy_(std::move(retry_policy)),
+        query_plan_refresh_retry_policy_(
+            std::move(query_plan_refresh_retry_policy)),
+        query_plan_backoff_policy_(backoff_policy_->clone()),
+        options_(std::move(options)) {}
+
+  void UpdateSource() {
+    internal::OptionsSpan options_span(options_);
+    while (!query_plan_refresh_retry_policy_->IsExhausted()) {
+      query_plan_data_ = query_plan_->response();
+      if (query_plan_data_.ok()) {
+        query_plan_valid_ = true;
+        request_.set_prepared_query(query_plan_data_->prepared_query());
+        source_ = source_factory_(
+            request_, query_plan_data_->metadata(),
+            options_->get<bigtable::DataRetryPolicyOption>()->clone(),
+            options_->get<bigtable::DataBackoffPolicyOption>()->clone(),
+            operation_context_);
+        if (source_->ok()) return;
+
+        last_status_ = source_->status();
+        if (ExecuteQueryPlanRefreshRetry::IsQueryPlanExpired(
+                source_->status())) {
+          query_plan_valid_ = false;
+          query_plan_->Invalidate(source_->status(),
+                                  query_plan_data_->prepared_query());
+        }
+        if (IsStatusMetadataIndicatingRetryPolicyExhausted(source_->status())) {
+          source_ = std::make_unique<StatusOnlyResultSetSource>(last_status_);
+          return;
+        }
+      } else {
+        last_status_ = query_plan_data_.status();
+      }
+
+      auto delay = internal::Backoff(
+          last_status_, __func__, *query_plan_refresh_retry_policy_,
+          *query_plan_backoff_policy_, Idempotency::kIdempotent,
+          false /* enable_server_retries */);
+      if (!delay) break;
+      std::this_thread::sleep_for(*delay);
+    }
+    auto retry_loop_error = internal::RetryLoopError(
+        last_status_, __func__,
+        query_plan_refresh_retry_policy_->IsExhausted());
+    source_ = std::make_unique<StatusOnlyResultSetSource>(retry_loop_error);
+    last_status_ = retry_loop_error;
+  }
+
+  google::bigtable::v2::ExecuteQueryRequest request_;
+  SourceFactoryFn source_factory_;
+  std::shared_ptr<QueryPlan> query_plan_;
+  std::shared_ptr<OperationContext> operation_context_;
+  std::unique_ptr<BackoffPolicy> backoff_policy_;
+  std::unique_ptr<bigtable::DataRetryPolicy> retry_policy_;
+  std::unique_ptr<bigtable::DataRetryPolicy> query_plan_refresh_retry_policy_;
+  std::unique_ptr<BackoffPolicy> query_plan_backoff_policy_;
+  internal::ImmutableOptions options_;
+
+  absl::optional<StatusOr<std::unique_ptr<bigtable::ResultSourceInterface>>>
+      source_;
+  StatusOr<google::bigtable::v2::PrepareQueryResponse> query_plan_data_;
+  bool query_plan_valid_ = false;
+  Status last_status_;
+};
+
 bigtable::RowStream DataConnectionImpl::ExecuteQuery(
     bigtable::ExecuteQueryParams params) {
   auto current = google::cloud::internal::SaveCurrentOptions();
@@ -892,7 +1027,7 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
   auto const tracing_enabled = RpcStreamTracingEnabled();
   auto const& tracing_options = RpcTracingOptions();
 
-  auto retry_resume_fn =
+  auto source_fn =
       [stub = stub_, tracing_enabled, tracing_options](
           google::bigtable::v2::ExecuteQueryRequest& request,
           google::bigtable::v2::ResultSetMetadata metadata,
@@ -930,56 +1065,15 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
   auto operation_context = operation_context_factory_->ExecuteQuery(
       request.instance_name(), app_profile_id(*current));
 
-  auto query_plan = params.bound_query.query_plan_;
-  auto query_plan_retry_policy = query_plan_refresh_retry_policy(*current);
-  auto query_plan_backoff_policy = backoff_policy(*current);
-  Status last_status;
-
-  // TODO(sdhart): OperationContext needs to be plumbed through the QueryPlan
-  // refresh_fn so that it's shared with the ExecuteQuery attempts.
-  while (!query_plan_retry_policy->IsExhausted()) {
-    // Snapshot query_plan data.
-    // This access could cause a query plan refresh to occur.
-    StatusOr<google::bigtable::v2::PrepareQueryResponse> query_plan_data =
-        query_plan->response();
-
-    if (query_plan_data.ok()) {
-      request.set_prepared_query(query_plan_data->prepared_query());
-      auto source = retry_resume_fn(
-          request, query_plan_data->metadata(), retry_policy(*current),
-          backoff_policy(*current), operation_context);
-      if (source.ok()) {
-        return bigtable::RowStream(*std::move(source));
-      }
-      last_status = source.status();
-
-      if (IsStatusIndicatingInternalError(source.status())) {
-        return bigtable::RowStream(std::make_unique<StatusOnlyResultSetSource>(
-            std::move(last_status)));
-      }
+  auto query_plan_refreshing_source =
+      QueryPlanRefreshingPartialResultSource::Create(
+          std::move(request), std::move(source_fn),
+          std::move(params.bound_query.query_plan_),
+          std::move(operation_context), backoff_policy(*current),
+          retry_policy(*current),
+          execute_query_plan_refresh_retry_policy(*current), current);
 
-      if (QueryPlanRefreshRetry::IsQueryPlanExpired(source.status())) {
-        query_plan->Invalidate(source.status(),
-                               query_plan_data->prepared_query());
-      }
-      if (IsStatusMetadataIndicatingRetryPolicyExhausted(source.status())) {
-        return bigtable::RowStream(std::make_unique<StatusOnlyResultSetSource>(
-            std::move(last_status)));
-      }
-    } else {
-      last_status = query_plan_data.status();
-    }
-
-    auto delay =
-        internal::Backoff(last_status, __func__, *query_plan_retry_policy,
-                          *query_plan_backoff_policy, Idempotency::kIdempotent,
-                          false /* enable_server_retries */);
-    if (!delay) break;
-    std::this_thread::sleep_for(*delay);
-  }
-  return bigtable::RowStream(
-      std::make_unique<StatusOnlyResultSetSource>(internal::RetryLoopError(
-          last_status, __func__, query_plan_retry_policy->IsExhausted())));
+  return bigtable::RowStream(std::move(query_plan_refreshing_source));
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/bigtable/internal/data_connection_impl_test.cc

@@ -212,10 +212,10 @@ DataLimitedErrorCountRetryPolicy TestRetryPolicy() {
   return DataLimitedErrorCountRetryPolicy(kNumRetries);
 }
 
-bigtable::experimental::QueryPlanRefreshLimitedErrorCountRetryPolicy
+bigtable::experimental::ExecuteQueryPlanRefreshLimitedErrorCountRetryPolicy
 TestQueryPlanRefreshRetryPolicy() {
-  return bigtable::experimental::QueryPlanRefreshLimitedErrorCountRetryPolicy(
-      kNumRetries);
+  return bigtable::experimental::
+      ExecuteQueryPlanRefreshLimitedErrorCountRetryPolicy(kNumRetries);
 }
 
 ExponentialBackoffPolicy TestBackoffPolicy() {
@@ -255,7 +255,8 @@ Options CallOptionsWithoutClientContextSetup() {
       Options{}
           .set<bigtable::AppProfileIdOption>(kAppProfile)
           .set<DataRetryPolicyOption>(TestRetryPolicy().clone())
-          .set<bigtable::experimental::QueryPlanRefreshRetryPolicyOption>(
+          .set<
+              bigtable::experimental::ExecuteQueryPlanRefreshRetryPolicyOption>(
               TestQueryPlanRefreshRetryPolicy().clone())
           .set<DataBackoffPolicyOption>(TestBackoffPolicy().clone()));
 }

google/cloud/bigtable/internal/defaults.cc

@@ -251,12 +251,21 @@ Options DefaultDataOptions(Options opts) {
             kBigtableLimits.maximum_retry_period)
             .clone());
   }
-  if (!opts.has<bigtable::experimental::QueryPlanRefreshRetryPolicyOption>()) {
-    opts.set<bigtable::experimental::QueryPlanRefreshRetryPolicyOption>(
-        bigtable::experimental::QueryPlanRefreshLimitedTimeRetryPolicy(
+  if (!opts.has<
+          bigtable::experimental::ExecuteQueryPlanRefreshRetryPolicyOption>()) {
+    opts.set<bigtable::experimental::ExecuteQueryPlanRefreshRetryPolicyOption>(
+        bigtable::experimental::ExecuteQueryPlanRefreshLimitedTimeRetryPolicy(
             kBigtableLimits.maximum_retry_period)
             .clone());
   }
+  if (!opts.has<bigtable::experimental::
+                    QueryPlanRefreshFunctionRetryPolicyOption>()) {
+    opts.set<bigtable::experimental::QueryPlanRefreshFunctionRetryPolicyOption>(
+        bigtable::experimental::QueryPlanRefreshFunctionLimitedTimeRetryPolicy(
+            kBigtableLimits.maximum_retry_period)
+            .clone());
+  }
+
   if (!opts.has<bigtable::DataBackoffPolicyOption>()) {
     opts.set<bigtable::DataBackoffPolicyOption>(
         ExponentialBackoffPolicy(kBigtableLimits.initial_delay / 2,

google/cloud/bigtable/internal/partial_result_set_source.cc

@@ -53,9 +53,7 @@ PartialResultSetSource::Create(
   auto status = source->ReadFromStream();
 
   // Any error during parsing will be returned.
-  if (!status.ok()) {
-    return status;
-  }
+  if (!status.ok()) return status;
 
   return {std::move(source)};
 }
@@ -142,15 +140,23 @@ Status PartialResultSetSource::ReadFromStream() {
   if (reader_->Read(resume_token_, result_set)) {
     return ProcessDataFromStream(result_set.result);
   }
+
+  last_status_ = reader_->Finish();
   state_ = State::kFinished;
-  // buffered_rows_ and read_buffer_ are expected to be empty because the last
-  // successful read would have had a sentinel resume_token, causing
-  // ProcessDataFromStream to commit them.
-  if (!buffered_rows_.empty() || !read_buffer_.empty()) {
+
+  if (ExecuteQueryPlanRefreshRetry::IsQueryPlanExpired(last_status_)) {
+    if (resume_token_ && !resume_token_->empty()) {
+      return internal::InternalError(
+          "Query plan expired during a retry attempt", GCP_ERROR_INFO());
+    }
+  } else if (!buffered_rows_.empty() || !read_buffer_.empty()) {
+    // buffered_rows_ and read_buffer_ are expected to be empty because the last
+    // successful read would have had a sentinel resume_token, causing
+    // ProcessDataFromStream to commit them.
     return internal::InternalError("Stream ended with uncommitted rows.",
                                    GCP_ERROR_INFO());
   }
-  last_status_ = reader_->Finish();
+
   return last_status_;
 }