impl(bigtable): update AsyncRowReader to support new OperationContext (#15302)

Author: scotthart

google/cloud/bigtable/internal/async_row_reader.cc

@@ -45,19 +45,24 @@ void AsyncRowReader::MakeRequest() {
   parser_ = bigtable::internal::ReadRowsParserFactory().Create(reverse_);
 
   internal::ScopedCallContext scope(call_context_);
-  context_ = std::make_shared<grpc::ClientContext>();
-  internal::ConfigureContext(*context_, *call_context_.options);
-  operation_context_->PreCall(*context_);
+  client_context_ = std::make_shared<grpc::ClientContext>();
+  internal::ConfigureContext(*client_context_, *call_context_.options);
+  operation_context_->PreCall(*client_context_);
 
   auto self = this->shared_from_this();
   PerformAsyncStreamingRead(
-      stub_->AsyncReadRows(cq_, context_, options_, request),
+      stub_->AsyncReadRows(cq_, client_context_, options_, request),
       [self](v2::ReadRowsResponse r) {
         return self->OnDataReceived(std::move(r));
       },
       [self](Status s) { self->OnStreamFinished(std::move(s)); });
 }
 
+void AsyncRowReader::UserWantsRows() {
+  operation_context_->ElementRequest(*client_context_);
+  TryGiveRowToUser();
+}
+
 void AsyncRowReader::TryGiveRowToUser() {
   // The user is likely to ask for more rows immediately after receiving a
   // row, which means that this function will be called recursively. The depth
@@ -102,6 +107,7 @@ void AsyncRowReader::TryGiveRowToUser() {
 
   auto self = this->shared_from_this();
   bool const break_recursion = recursion_level_ >= 100;
+  operation_context_->ElementDelivery(*client_context_);
   on_row_(std::move(row)).then([self, break_recursion](future<bool> fut) {
     bool should_cancel;
 #if GOOGLE_CLOUD_CPP_HAVE_EXCEPTIONS
@@ -165,6 +171,7 @@ future<bool> AsyncRowReader::OnDataReceived(
 }
 
 void AsyncRowReader::OnStreamFinished(Status status) {
+  operation_context_->PostCall(*client_context_, status);
   // assert(!continue_reading_);
   if (status_.ok()) {
     status_ = std::move(status);
@@ -206,6 +213,7 @@ void AsyncRowReader::OnStreamFinished(Status status) {
   if (status_.ok()) {
     // We've successfully finished the scan.
     whole_op_finished_ = true;
+    operation_context_->OnDone(status_);
     TryGiveRowToUser();
     return;
   }
@@ -217,11 +225,11 @@ void AsyncRowReader::OnStreamFinished(Status status) {
     // Can't retry.
     status_ = std::move(delay).status();
     whole_op_finished_ = true;
+    operation_context_->OnDone(status_);
     TryGiveRowToUser();
     return;
   }
-  operation_context_->PostCall(*context_, {});
-  context_.reset();
+  client_context_.reset();
   auto self = this->shared_from_this();
   internal::TracedAsyncBackoff(cq_, *call_context_.options, *delay,
                                "Async Backoff")

google/cloud/bigtable/internal/async_row_reader.h

@@ -64,13 +64,14 @@ class AsyncRowReader : public std::enable_shared_from_this<AsyncRowReader> {
                      bigtable::Filter filter, bool reverse,
                      std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
                      std::unique_ptr<BackoffPolicy> backoff_policy,
-                     bool enable_server_retries) {
+                     bool enable_server_retries,
+                     std::shared_ptr<OperationContext> operation_context) {
     auto reader = std::shared_ptr<AsyncRowReader>(new AsyncRowReader(
         std::move(cq), std::move(stub), std::move(app_profile_id),
         std::move(table_name), std::move(on_row), std::move(on_finish),
         std::move(row_set), rows_limit, std::move(filter), reverse,
         std::move(retry_policy), std::move(backoff_policy),
-        enable_server_retries));
+        enable_server_retries, std::move(operation_context)));
     reader->MakeRequest();
   }
 
@@ -82,7 +83,8 @@ class AsyncRowReader : public std::enable_shared_from_this<AsyncRowReader> {
                  bigtable::Filter filter, bool reverse,
                  std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
                  std::unique_ptr<BackoffPolicy> backoff_policy,
-                 bool enable_server_retries)
+                 bool enable_server_retries,
+                 std::shared_ptr<OperationContext> operation_context)
       : cq_(std::move(cq)),
         stub_(std::move(stub)),
         app_profile_id_(std::move(app_profile_id)),
@@ -97,15 +99,16 @@ class AsyncRowReader : public std::enable_shared_from_this<AsyncRowReader> {
         backoff_policy_(std::move(backoff_policy)),
         enable_server_retries_(enable_server_retries),
         options_(internal::SaveCurrentOptions()),
-        call_context_(options_) {}
+        call_context_(options_),
+        operation_context_(std::move(operation_context)) {}
 
   void MakeRequest();
 
   /**
    * Called when the user asks for more rows via satisfying the future returned
    * from the row callback.
    */
-  void UserWantsRows() { TryGiveRowToUser(); }
+  void UserWantsRows();
 
   /**
    * Attempt to call a user callback.
@@ -175,9 +178,8 @@ class AsyncRowReader : public std::enable_shared_from_this<AsyncRowReader> {
   int recursion_level_ = 0;
   internal::ImmutableOptions options_;
   internal::CallContext call_context_;
-  std::shared_ptr<grpc::ClientContext> context_;
-  std::shared_ptr<OperationContext> operation_context_ =
-      std::make_shared<OperationContext>();
+  std::shared_ptr<grpc::ClientContext> client_context_;
+  std::shared_ptr<OperationContext> operation_context_;
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END