impl(bigtable): improve QueryPlan and add more tests (#15688)

Author: scotthart

google/cloud/bigtable/internal/query_plan.cc

@@ -26,8 +26,9 @@ auto constexpr kRefreshDeadlineOffsetMs = 1000;
 }  // namespace
 
 std::shared_ptr<QueryPlan> QueryPlan::Create(
-    CompletionQueue cq, google::bigtable::v2::PrepareQueryResponse response,
-    RefreshFn fn, std::shared_ptr<Clock> clock) {
+    CompletionQueue cq,
+    StatusOr<google::bigtable::v2::PrepareQueryResponse> response, RefreshFn fn,
+    std::shared_ptr<Clock> clock) {
   auto plan = std::shared_ptr<QueryPlan>(new QueryPlan(
       std::move(cq), std::move(clock), std::move(fn), std::move(response)));
   plan->Initialize();
@@ -36,7 +37,7 @@ std::shared_ptr<QueryPlan> QueryPlan::Create(
 
 void QueryPlan::Initialize() {
   std::unique_lock<std::mutex> lock(mu_);
-  ScheduleRefresh(lock);
+  if (state_ == RefreshState::kDone) ScheduleRefresh(lock);
 }
 
 // ScheduleRefresh should only be called after updating response_.
@@ -71,7 +72,7 @@ void QueryPlan::ExpiredRefresh() {
       state_ = RefreshState::kBegin;
     }
   }
-  RefreshQueryPlan(RefreshMode::kExpired);
+  RefreshQueryPlan();
 }
 
 void QueryPlan::Invalidate(Status status,
@@ -82,26 +83,17 @@ void QueryPlan::Invalidate(Status status,
     // query plan, so we track what the previous plan id was.
     if (!IsRefreshing(lock) && old_query_plan_id_ != invalid_query_plan_id) {
       old_query_plan_id_ = invalid_query_plan_id;
+      response_ = std::move(status);
       state_ = RefreshState::kBegin;
     }
   }
-  RefreshQueryPlan(RefreshMode::kInvalidated, std::move(status));
 }
 
-void QueryPlan::RefreshQueryPlan(RefreshMode mode, Status error) {
+void QueryPlan::RefreshQueryPlan() {
   {
     std::unique_lock<std::mutex> lock_1(mu_);
-#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_QUERY_PLAN_REFRESH_ASSERT
-    assert(waiting_threads_ >= 0);
-#endif
-    ++waiting_threads_;
     cond_.wait(lock_1, [this] { return state_ != RefreshState::kPending; });
-    --waiting_threads_;
-#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_QUERY_PLAN_REFRESH_ASSERT
-    assert(waiting_threads_ >= 0);
-#endif
     if (state_ == RefreshState::kDone) return;
-    if (mode == RefreshMode::kInvalidated) response_ = std::move(error);
     state_ = RefreshState::kPending;
   }
   auto response = refresh_fn_().get();
@@ -114,7 +106,7 @@ void QueryPlan::RefreshQueryPlan(RefreshMode mode, Status error) {
       done = true;
       // If we have to refresh an invalidated query plan, cancel any existing
       // timer before starting a new one.
-      refresh_timer_.cancel();
+      if (refresh_timer_.valid()) refresh_timer_.cancel();
       ScheduleRefresh(lock_2);
     } else {
       // If there are no waiting threads that could call the refresh_fn, then
@@ -124,13 +116,7 @@ void QueryPlan::RefreshQueryPlan(RefreshMode mode, Status error) {
       // If there are waiting threads, then we want to try again to get a
       // refreshed query plan, but we want to avoid a stampede of refresh RPCs
       // so we only notify one of the waiting threads.
-#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_QUERY_PLAN_REFRESH_ASSERT
-      assert(waiting_threads_ >= 0);
-#endif
-      if (waiting_threads_ == 0) {
-        state_ = RefreshState::kDone;
-        done = true;
-      }
+      state_ = RefreshState::kBegin;
     }
   }
   if (done) {
@@ -147,7 +133,7 @@ StatusOr<google::bigtable::v2::PrepareQueryResponse> QueryPlan::response() {
       return response_;
     }
     lock.unlock();
-    RefreshQueryPlan(RefreshMode::kAlreadyRefreshing);
+    RefreshQueryPlan();
     lock.lock();
   }
 

google/cloud/bigtable/internal/query_plan.h

@@ -39,7 +39,8 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
 
   // Calls the constructor and then Initialize.
   static std::shared_ptr<QueryPlan> Create(
-      CompletionQueue cq, google::bigtable::v2::PrepareQueryResponse response,
+      CompletionQueue cq,
+      StatusOr<google::bigtable::v2::PrepareQueryResponse> response,
       RefreshFn fn, std::shared_ptr<Clock> clock = std::make_shared<Clock>());
 
   // Invalidates the current QueryPlan and triggers a refresh.
@@ -57,8 +58,9 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
 
  private:
   QueryPlan(CompletionQueue cq, std::shared_ptr<Clock> clock, RefreshFn fn,
-            google::bigtable::v2::PrepareQueryResponse response)
-      : cq_(std::move(cq)),
+            StatusOr<google::bigtable::v2::PrepareQueryResponse> response)
+      : state_(response.ok() ? RefreshState::kDone : RefreshState::kBegin),
+        cq_(std::move(cq)),
         clock_(std::move(clock)),
         refresh_fn_(std::move(fn)),
         response_(std::move(response)) {}
@@ -73,12 +75,9 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
   // capturing a std::weak_ptr to this that calls RefreshQueryPlan.
   void ScheduleRefresh(std::unique_lock<std::mutex> const&);
 
-  enum class RefreshMode { kExpired, kInvalidated, kAlreadyRefreshing };
   // Performs the synchronization around calling RefreshFn and updating
   // response_.
-  //  void RefreshQueryPlan();
-
-  void RefreshQueryPlan(RefreshMode mode, Status error = {});
+  void RefreshQueryPlan();
 
   void ExpiredRefresh();
 
@@ -93,17 +92,14 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
     kPending,  // waiting for an active thread to refresh response_
     kDone,     // response_ has been refreshed
   };
-  RefreshState state_ = RefreshState::kDone;
+  RefreshState state_;
 
   CompletionQueue cq_;
   std::shared_ptr<Clock> clock_;
   RefreshFn refresh_fn_;
   future<void> refresh_timer_;
   mutable std::mutex mu_;
   std::condition_variable cond_;
-  // waiting_threads_ is only a snapshot, but it helps us reduce the number of
-  // RPCs in flight to refresh the same query plan.
-  int waiting_threads_ = 0;
   std::string old_query_plan_id_;
   StatusOr<google::bigtable::v2::PrepareQueryResponse>
       response_;  // GUARDED_BY(mu_)