Merge branch 'main' into execute-query

Author: mpeddada1

google/cloud/bigtable/CMakeLists.txt

@@ -201,6 +201,8 @@ add_library(
     internal/legacy_row_reader.h
     internal/logging_data_client.cc
     internal/logging_data_client.h
+    internal/logging_result_set_reader.cc
+    internal/logging_result_set_reader.h
     internal/metrics.cc
     internal/metrics.h
     internal/mutate_rows_limiter.cc
@@ -484,6 +486,7 @@ if (BUILD_TESTING)
         internal/legacy_bulk_mutator_test.cc
         internal/legacy_row_reader_test.cc
         internal/logging_data_client_test.cc
+        internal/logging_result_set_reader_test.cc
         internal/metrics_test.cc
         internal/mutate_rows_limiter_test.cc
         internal/operation_context_factory_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

@@ -62,6 +62,7 @@ bigtable_client_unit_tests = [
     "internal/legacy_bulk_mutator_test.cc",
     "internal/legacy_row_reader_test.cc",
     "internal/logging_data_client_test.cc",
+    "internal/logging_result_set_reader_test.cc",
     "internal/metrics_test.cc",
     "internal/mutate_rows_limiter_test.cc",
     "internal/operation_context_factory_test.cc",

google/cloud/bigtable/google_cloud_cpp_bigtable.bzl

@@ -99,6 +99,7 @@ google_cloud_cpp_bigtable_hdrs = [
     "internal/legacy_async_row_sampler.h",
     "internal/legacy_row_reader.h",
     "internal/logging_data_client.h",
+    "internal/logging_result_set_reader.h",
     "internal/metrics.h",
     "internal/mutate_rows_limiter.h",
     "internal/operation_context.h",
@@ -217,6 +218,7 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/legacy_async_row_sampler.cc",
     "internal/legacy_row_reader.cc",
     "internal/logging_data_client.cc",
+    "internal/logging_result_set_reader.cc",
     "internal/metrics.cc",
     "internal/mutate_rows_limiter.cc",
     "internal/operation_context.cc",

google/cloud/bigtable/internal/logging_result_set_reader.cc

@@ -0,0 +1,52 @@
+// 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/logging_result_set_reader.h"
+#include "google/cloud/internal/debug_string.h"
+#include "google/cloud/log.h"
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+using ::google::cloud::internal::DebugString;
+
+void LoggingResultSetReader::TryCancel() { impl_->TryCancel(); }
+
+bool LoggingResultSetReader::Read(
+    absl::optional<std::string> const& resume_token,
+    UnownedPartialResultSet& result) {
+  if (resume_token) {
+    GCP_LOG(DEBUG) << __func__ << "() << resume_token=\""
+                   << DebugString(*resume_token, tracing_options_) << "\"";
+  } else {
+    GCP_LOG(DEBUG) << __func__ << "() << (unresumable)";
+  }
+  bool success = impl_->Read(resume_token, result);
+  if (!success) {
+    GCP_LOG(DEBUG) << __func__ << "() >> (failed)";
+  } else {
+    GCP_LOG(DEBUG) << __func__ << "() >> resumption="
+                   << (result.resumption ? "true" : "false");
+  }
+  return success;
+}
+
+Status LoggingResultSetReader::Finish() { return impl_->Finish(); }
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google

google/cloud/bigtable/internal/logging_result_set_reader.h

@@ -0,0 +1,56 @@
+// 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_LOGGING_RESULT_SET_READER_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_LOGGING_RESULT_SET_READER_H
+
+#include "google/cloud/bigtable/internal/partial_result_set_reader.h"
+#include "google/cloud/bigtable/version.h"
+#include "google/cloud/tracing_options.h"
+#include "absl/types/optional.h"
+#include <memory>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+/**
+ * A `PartialResultSetReader` decorator that logs the `Read()` `resume_token`
+ * parameter and the `PartialResultSet::resumption` return value. This is an
+ * extension to the standard `BigtableLogging` request/response logging.
+ */
+class LoggingResultSetReader : public PartialResultSetReader {
+ public:
+  LoggingResultSetReader(std::unique_ptr<PartialResultSetReader> impl,
+                         TracingOptions tracing_options)
+      : impl_(std::move(impl)), tracing_options_(std::move(tracing_options)) {}
+  ~LoggingResultSetReader() override = default;
+
+  void TryCancel() override;
+  bool Read(absl::optional<std::string> const& resume_token,
+            UnownedPartialResultSet& result) override;
+  Status Finish() override;
+
+ private:
+  std::unique_ptr<PartialResultSetReader> impl_;
+  TracingOptions tracing_options_;
+};
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_LOGGING_RESULT_SET_READER_H

google/cloud/bigtable/internal/logging_result_set_reader_test.cc

@@ -0,0 +1,98 @@
+// 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/logging_result_set_reader.h"
+#include "google/cloud/bigtable/testing/mock_partial_result_set_reader.h"
+#include "google/cloud/log.h"
+#include "google/cloud/testing_util/scoped_log.h"
+#include "google/cloud/testing_util/status_matchers.h"
+#include "google/cloud/tracing_options.h"
+#include <gmock/gmock.h>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+namespace {
+
+using ::testing::_;
+using ::testing::AllOf;
+using ::testing::Contains;
+using ::testing::HasSubstr;
+using ::testing::IsEmpty;
+using ::testing::StartsWith;
+
+class LoggingResultSetReaderTest : public ::testing::Test {
+ protected:
+  testing_util::ScopedLog log_;
+};
+
+TEST_F(LoggingResultSetReaderTest, TryCancel) {
+  auto mock = std::make_unique<bigtable_testing::MockPartialResultSetReader>();
+  EXPECT_CALL(*mock, TryCancel()).Times(1);
+  LoggingResultSetReader reader(std::move(mock), TracingOptions{});
+  reader.TryCancel();
+
+  EXPECT_THAT(log_.ExtractLines(), IsEmpty());
+}
+
+TEST_F(LoggingResultSetReaderTest, Read) {
+  auto mock = std::make_unique<bigtable_testing::MockPartialResultSetReader>();
+  EXPECT_CALL(*mock, Read(_, _))
+      .WillOnce([](absl::optional<std::string> const&,
+                   UnownedPartialResultSet& result) {
+        result.resumption = false;
+        result.result.set_resume_token("test-token");
+        return true;
+      })
+      .WillOnce([] { return false; });
+  LoggingResultSetReader reader(std::move(mock), TracingOptions{});
+  google::bigtable::v2::PartialResultSet partial_result_set;
+  auto result =
+      UnownedPartialResultSet::FromPartialResultSet(partial_result_set);
+  ASSERT_TRUE(reader.Read("", result));
+  EXPECT_EQ("test-token", result.result.resume_token());
+
+  auto log_lines = log_.ExtractLines();
+  EXPECT_THAT(log_lines, AllOf(Contains(StartsWith("Read()"))));
+  EXPECT_THAT(log_lines, Contains(HasSubstr("resume_token=\"\"")));
+  EXPECT_THAT(log_lines, Contains(HasSubstr("resumption=false")));
+
+  ASSERT_FALSE(reader.Read("test-token", result));
+
+  log_lines = log_.ExtractLines();
+  EXPECT_THAT(log_lines, AllOf(Contains(StartsWith("Read()"))));
+  EXPECT_THAT(log_lines, Contains(HasSubstr("resume_token=\"test-token\"")));
+  EXPECT_THAT(log_lines, Contains(HasSubstr("(failed)")));
+}
+
+TEST_F(LoggingResultSetReaderTest, Finish) {
+  Status const expected_status = Status(StatusCode::kOutOfRange, "weird");
+  auto mock = std::make_unique<bigtable_testing::MockPartialResultSetReader>();
+  EXPECT_CALL(*mock, Finish()).WillOnce([expected_status] {
+    return expected_status;  // NOLINT(performance-no-automatic-move)
+  });
+  LoggingResultSetReader reader(std::move(mock), TracingOptions{});
+  auto status = reader.Finish();
+  EXPECT_EQ(expected_status, status);
+
+  auto log_lines = log_.ExtractLines();
+  EXPECT_THAT(log_lines, IsEmpty());
+}
+
+}  // namespace
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google

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_)