impl(bigtable): update AsyncRowSampler to support new OperationContext (#15304)

Author: scotthart

google/cloud/bigtable/internal/async_row_sampler.cc

@@ -31,11 +31,12 @@ future<StatusOr<std::vector<bigtable::RowKeySample>>> AsyncRowSampler::Create(
     CompletionQueue cq, std::shared_ptr<BigtableStub> stub,
     std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy, bool enable_server_retries,
-    std::string const& app_profile_id, std::string const& table_name) {
-  std::shared_ptr<AsyncRowSampler> sampler(
-      new AsyncRowSampler(std::move(cq), std::move(stub),
-                          std::move(retry_policy), std::move(backoff_policy),
-                          enable_server_retries, app_profile_id, table_name));
+    std::string const& app_profile_id, std::string const& table_name,
+    std::shared_ptr<OperationContext> operation_context) {
+  std::shared_ptr<AsyncRowSampler> sampler(new AsyncRowSampler(
+      std::move(cq), std::move(stub), std::move(retry_policy),
+      std::move(backoff_policy), enable_server_retries, app_profile_id,
+      table_name, std::move(operation_context)));
   sampler->StartIteration();
   return sampler->promise_.get_future();
 }
@@ -44,7 +45,8 @@ AsyncRowSampler::AsyncRowSampler(
     CompletionQueue cq, std::shared_ptr<BigtableStub> stub,
     std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy, bool enable_server_retries,
-    std::string const& app_profile_id, std::string const& table_name)
+    std::string const& app_profile_id, std::string const& table_name,
+    std::shared_ptr<OperationContext> operation_context)
     : cq_(std::move(cq)),
       stub_(std::move(stub)),
       retry_policy_(std::move(retry_policy)),
@@ -54,21 +56,22 @@ AsyncRowSampler::AsyncRowSampler(
       table_name_(std::move(table_name)),
       promise_([this] { keep_reading_ = false; }),
       options_(internal::SaveCurrentOptions()),
-      call_context_(options_) {}
+      call_context_(options_),
+      operation_context_(std::move(operation_context)) {}
 
 void AsyncRowSampler::StartIteration() {
   v2::SampleRowKeysRequest request;
   request.set_app_profile_id(app_profile_id_);
   request.set_table_name(table_name_);
 
   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<v2::SampleRowKeysResponse>(
-      stub_->AsyncSampleRowKeys(cq_, context_, options_, request),
+      stub_->AsyncSampleRowKeys(cq_, client_context_, options_, request),
       [self](v2::SampleRowKeysResponse response) {
         return self->OnRead(std::move(response));
       },
@@ -84,20 +87,23 @@ future<bool> AsyncRowSampler::OnRead(v2::SampleRowKeysResponse response) {
 }
 
 void AsyncRowSampler::OnFinish(Status const& status) {
+  operation_context_->PostCall(*client_context_, status);
+
   if (status.ok()) {
     promise_.set_value(std::move(samples_));
+    operation_context_->OnDone(status);
     return;
   }
   auto delay = internal::Backoff(status, "AsyncSampleRows", *retry_policy_,
                                  *backoff_policy_, Idempotency::kIdempotent,
                                  enable_server_retries_);
   if (!delay) {
     promise_.set_value(std::move(delay).status());
+    operation_context_->OnDone(status);
     return;
   }
 
-  operation_context_->PostCall(*context_, {});
-  context_.reset();
+  client_context_.reset();
   samples_.clear();
   auto self = this->shared_from_this();
   internal::TracedAsyncBackoff(cq_, *call_context_.options, *delay,

google/cloud/bigtable/internal/async_row_sampler.h

@@ -43,14 +43,16 @@ class AsyncRowSampler : public std::enable_shared_from_this<AsyncRowSampler> {
       CompletionQueue cq, std::shared_ptr<BigtableStub> stub,
       std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
       std::unique_ptr<BackoffPolicy> backoff_policy, bool enable_server_retries,
-      std::string const& app_profile_id, std::string const& table_name);
+      std::string const& app_profile_id, std::string const& table_name,
+      std::shared_ptr<OperationContext> operation_context);
 
  private:
   AsyncRowSampler(CompletionQueue cq, std::shared_ptr<BigtableStub> stub,
                   std::unique_ptr<bigtable::DataRetryPolicy> retry_policy,
                   std::unique_ptr<BackoffPolicy> backoff_policy,
                   bool enable_server_retries, std::string const& app_profile_id,
-                  std::string const& table_name);
+                  std::string const& table_name,
+                  std::shared_ptr<OperationContext> operation_context);
 
   void StartIteration();
   future<bool> OnRead(google::bigtable::v2::SampleRowKeysResponse response);
@@ -69,9 +71,8 @@ class AsyncRowSampler : public std::enable_shared_from_this<AsyncRowSampler> {
   promise<StatusOr<std::vector<bigtable::RowKeySample>>> promise_;
   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