impl(bigtable): update DefaultRowReader to support new OperationContext (#15308)

Author: scotthart

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -59,6 +59,23 @@ inline bool enable_server_retries(Options const& options) {
   return options.get<EnableServerRetriesOption>();
 }
 
+// This function allows for ReadRow and ReadRowsFull to provide an instance of
+// an OperationContext specific to that operation.
+bigtable::RowReader ReadRowsHelper(
+    std::shared_ptr<BigtableStub>& stub,
+    internal::ImmutableOptions const& current,
+    bigtable::ReadRowsParams
+        params,  // NOLINT(performance-unnecessary-value-param)
+    std::shared_ptr<OperationContext>
+        operation_context) {  // NOLINT(performance-unnecessary-value-param)
+  auto impl = std::make_shared<DefaultRowReader>(
+      stub, std::move(params.app_profile_id), std::move(params.table_name),
+      std::move(params.row_set), params.rows_limit, std::move(params.filter),
+      params.reverse, retry_policy(*current), backoff_policy(*current),
+      enable_server_retries(*current), std::move(operation_context));
+  return MakeRowReader(std::move(impl));
+}
+
 }  // namespace
 
 bigtable::Row TransformReadModifyWriteRowResponse(
@@ -99,7 +116,8 @@ DataConnectionImpl::DataConnectionImpl(
     std::string client_uid =
         internal::Sample(gen, 16, "abcdefghijklmnopqrstuvwxyz0123456789");
     operation_context_factory_ =
-        std::make_unique<MetricsOperationContextFactory>(std::move(client_uid));
+        std::make_unique<MetricsOperationContextFactory>(std::move(client_uid),
+                                                         options_);
   } else {
     operation_context_factory_ =
         std::make_unique<SimpleOperationContextFactory>();
@@ -231,12 +249,9 @@ DataConnectionImpl::AsyncBulkApply(std::string const& table_name,
 bigtable::RowReader DataConnectionImpl::ReadRowsFull(
     bigtable::ReadRowsParams params) {
   auto current = google::cloud::internal::SaveCurrentOptions();
-  auto impl = std::make_shared<DefaultRowReader>(
-      stub_, std::move(params.app_profile_id), std::move(params.table_name),
-      std::move(params.row_set), params.rows_limit, std::move(params.filter),
-      params.reverse, retry_policy(*current), backoff_policy(*current),
-      enable_server_retries(*current));
-  return MakeRowReader(std::move(impl));
+  auto operation_context = std::make_shared<OperationContext>();
+  return ReadRowsHelper(stub_, current, std::move(params),
+                        std::move(operation_context));
 }
 
 StatusOr<std::pair<bool, bigtable::Row>> DataConnectionImpl::ReadRow(
@@ -245,9 +260,16 @@ StatusOr<std::pair<bool, bigtable::Row>> DataConnectionImpl::ReadRow(
   auto current = google::cloud::internal::SaveCurrentOptions();
   bigtable::RowSet row_set(std::move(row_key));
   std::int64_t const rows_limit = 1;
-  auto reader = ReadRowsFull(bigtable::ReadRowsParams{
-      table_name, app_profile_id(*current), std::move(row_set), rows_limit,
-      std::move(filter)});
+  // TODO(sdhart): ensure OperationContext::OnDone is called correctly.
+  // TODO(sdhart): add ReadRow tests once we call
+  //  OperationContextFactory::ReadRow to create the operation_context.
+  auto operation_context = std::make_shared<OperationContext>();
+  auto reader =
+      ReadRowsHelper(stub_, current,
+                     bigtable::ReadRowsParams{
+                         table_name, app_profile_id(*current),
+                         std::move(row_set), rows_limit, std::move(filter)},
+                     std::move(operation_context));
 
   auto it = reader.begin();
   if (it == reader.end()) return std::make_pair(false, bigtable::Row("", {}));

google/cloud/bigtable/internal/default_row_reader.cc

@@ -29,7 +29,7 @@ DefaultRowReader::DefaultRowReader(
     bigtable::Filter filter, bool reverse,
     std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy, bool enable_server_retries,
-    Sleeper sleeper)
+    std::shared_ptr<OperationContext> operation_context, Sleeper sleeper)
     : stub_(std::move(stub)),
       app_profile_id_(std::move(app_profile_id)),
       table_name_(std::move(table_name)),
@@ -40,7 +40,8 @@ DefaultRowReader::DefaultRowReader(
       retry_policy_(std::move(retry_policy)),
       backoff_policy_(std::move(backoff_policy)),
       enable_server_retries_(enable_server_retries),
-      sleeper_(std::move(sleeper)) {}
+      sleeper_(std::move(sleeper)),
+      operation_context_(std::move(operation_context)) {}
 
 void DefaultRowReader::MakeRequest() {
   response_ = {};
@@ -62,10 +63,11 @@ void DefaultRowReader::MakeRequest() {
   }
 
   auto const& options = internal::CurrentOptions();
-  context_ = std::make_shared<grpc::ClientContext>();
-  internal::ConfigureContext(*context_, options);
-  operation_context_.PreCall(*context_);
-  stream_ = stub_->ReadRows(context_, options, request);
+  client_context_ = std::make_shared<grpc::ClientContext>();
+  internal::ConfigureContext(*client_context_, options);
+  operation_context_->PreCall(*client_context_);
+  called_post_call_ = false;
+  stream_ = stub_->ReadRows(client_context_, options, request);
   stream_is_open_ = true;
 
   parser_ = bigtable::internal::ReadRowsParserFactory().Create(reverse_);
@@ -78,9 +80,9 @@ bool DefaultRowReader::NextChunk() {
     auto v = stream_->Read();
     if (absl::holds_alternative<Status>(v)) {
       last_status_ = absl::get<Status>(std::move(v));
+      operation_context_->PostCall(*client_context_, last_status_);
+      called_post_call_ = true;
       response_ = {};
-      operation_context_.PostCall(*context_, {});
-      context_.reset();
       return false;
     }
     response_ = absl::get<google::bigtable::v2::ReadRowsResponse>(std::move(v));
@@ -92,6 +94,11 @@ bool DefaultRowReader::NextChunk() {
 }
 
 absl::variant<Status, bigtable::Row> DefaultRowReader::Advance() {
+  // We only want to call ElementRequest if an RPC has previously been
+  // made.
+  if (stream_is_open_) {
+    operation_context_->ElementRequest(*client_context_);
+  }
   if (operation_cancelled_) {
     return internal::CancelledError(
         "call cancelled",
@@ -100,6 +107,7 @@ absl::variant<Status, bigtable::Row> DefaultRowReader::Advance() {
   while (true) {
     auto variant = AdvanceOrFail();
     if (absl::holds_alternative<bigtable::Row>(variant)) {
+      operation_context_->ElementDelivery(*client_context_);
       return absl::get<bigtable::Row>(std::move(variant));
     }
 
@@ -143,16 +151,19 @@ absl::variant<Status, bigtable::Row> DefaultRowReader::Advance() {
 }
 
 absl::variant<Status, bigtable::Row> DefaultRowReader::AdvanceOrFail() {
-  grpc::Status status;
-  if (!stream_) {
-    MakeRequest();
-  }
+  grpc::Status grpc_status;
+  if (!stream_) MakeRequest();
   while (!parser_->HasNext()) {
     if (NextChunk()) {
       parser_->HandleChunk(
           std::move(*(response_.mutable_chunks(processed_chunks_count_))),
-          status);
-      if (!status.ok()) return MakeStatusFromRpcError(status);
+          grpc_status);
+      if (!grpc_status.ok()) {
+        auto status = MakeStatusFromRpcError(grpc_status);
+        operation_context_->PostCall(*client_context_, status);
+        called_post_call_ = true;
+        return status;
+      }
       continue;
     }
 
@@ -161,21 +172,30 @@ absl::variant<Status, bigtable::Row> DefaultRowReader::AdvanceOrFail() {
     // fails during cleanup.
     stream_is_open_ = false;
     if (!last_status_.ok()) return last_status_;
-    parser_->HandleEndOfStream(status);
-    return MakeStatusFromRpcError(status);
+    parser_->HandleEndOfStream(grpc_status);
+    return MakeStatusFromRpcError(grpc_status);
   }
 
   // We have a complete row in the parser.
-  bigtable::Row parsed_row = parser_->Next(status);
+  bigtable::Row parsed_row = parser_->Next(grpc_status);
 
-  if (!status.ok()) return MakeStatusFromRpcError(status);
+  if (!grpc_status.ok()) return MakeStatusFromRpcError(grpc_status);
 
   ++rows_count_;
   last_read_row_key_ = parsed_row.row_key();
   return parsed_row;
 }
 
 void DefaultRowReader::Cancel() {
+  // As the destructor also calls Cancel, we want to call OnDone exactly once.
+  // If parser_ == nullptr, then we know no RPC was ever attempted.
+  if (!operation_cancelled_ && parser_) {
+    if (!called_post_call_) {
+      called_post_call_ = true;
+      operation_context_->PostCall(*client_context_, last_status_);
+    }
+    operation_context_->OnDone(last_status_);
+  }
   operation_cancelled_ = true;
   if (!stream_is_open_) return;
   stream_->Cancel();
@@ -185,7 +205,7 @@ void DefaultRowReader::Cancel() {
       stream_->Read())) {
   }
 
-  context_.reset();
+  client_context_.reset();
   stream_is_open_ = false;
 }
 

google/cloud/bigtable/internal/default_row_reader.h

@@ -54,6 +54,7 @@ class DefaultRowReader : public RowReaderImpl {
       bigtable::Filter filter, bool reverse,
       std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
       std::unique_ptr<BackoffPolicy> backoff_policy, bool enable_server_retries,
+      std::shared_ptr<OperationContext> operation_context,
       Sleeper sleeper = [](auto d) { std::this_thread::sleep_for(d); });
 
   ~DefaultRowReader() override;
@@ -101,15 +102,17 @@ class DefaultRowReader : public RowReaderImpl {
   std::unique_ptr<BackoffPolicy> backoff_policy_;
   bool enable_server_retries_;
   Sleeper sleeper_;
-  std::shared_ptr<grpc::ClientContext> context_;
-  OperationContext operation_context_;
+  std::shared_ptr<grpc::ClientContext> client_context_;
+  std::shared_ptr<OperationContext> operation_context_;
 
   std::unique_ptr<bigtable::internal::ReadRowsParser> parser_;
   std::unique_ptr<
       internal::StreamingReadRpc<google::bigtable::v2::ReadRowsResponse>>
       stream_;
   bool stream_is_open_ = false;
   bool operation_cancelled_ = false;
+  // TODO(#15314): Refactor state machine to not require this flag.
+  bool called_post_call_ = false;
 
   /// The end of stream Status.
   Status last_status_;