impl(bigtable): add ScheduleStubRefresh function using AsyncPingAndWarm

google/cloud/bigtable/internal/connection_refresh_state.cc

@@ -30,6 +30,12 @@ namespace {
  */
 auto constexpr kConnectionReadyTimeout = std::chrono::seconds(10);
 
+void LogFailedConnectionRefresh(Status const& conn_status) {
+  if (!conn_status.ok()) {
+    GCP_LOG(WARNING) << "Failed to refresh connection. Error: " << conn_status;
+  }
+}
+
 }  // namespace
 
 ConnectionRefreshState::ConnectionRefreshState(
@@ -81,10 +87,7 @@ void ScheduleChannelRefresh(
                   std::chrono::system_clock::now() + kConnectionReadyTimeout)
                 .then([weak_channel, weak_cq_impl, state](future<Status> fut) {
                   auto conn_status = fut.get();
-                  if (!conn_status.ok()) {
-                    GCP_LOG(WARNING) << "Failed to refresh connection. Error: "
-                                     << conn_status;
-                  }
+                  LogFailedConnectionRefresh(conn_status);
                   auto channel = weak_channel.lock();
                   if (!channel) return;
                   auto cq_impl = weak_cq_impl.lock();
@@ -95,6 +98,65 @@ void ScheduleChannelRefresh(
   state->timers().RegisterTimer(std::move(timer_future));
 }
 
+void ScheduleStubRefresh(
+    std::shared_ptr<internal::CompletionQueueImpl> const& cq_impl,
+    std::shared_ptr<ConnectionRefreshState> const& state,
+    std::shared_ptr<BigtableStub> const& stub, std::string const& instance_name,
+    std::function<void(Status const&)> connection_status_fn) {
+  if (!connection_status_fn) {
+    connection_status_fn = LogFailedConnectionRefresh;
+  }
+  // The timers will only hold weak pointers to the channel or to the
+  // completion queue, so if either of them are destroyed, the timer chain
+  // will simply not continue.
+  std::weak_ptr<BigtableStub> weak_stub(stub);
+  std::weak_ptr<internal::CompletionQueueImpl> weak_cq_impl(cq_impl);
+  auto cq = CompletionQueue(cq_impl);
+  using TimerFuture = future<StatusOr<std::chrono::system_clock::time_point>>;
+  auto timer_future =
+      cq.MakeRelativeTimer(state->RandomizedRefreshDelay())
+          .then([weak_stub, weak_cq_impl, state, instance_name,
+                 connection_status_fn =
+                     std::move(connection_status_fn)](TimerFuture fut) mutable {
+            if (!fut.get()) {
+              // Timer cancelled.
+              return;
+            }
+            auto stub = weak_stub.lock();
+            if (!stub) return;
+            auto cq_impl = weak_cq_impl.lock();
+            if (!cq_impl) return;
+            auto cq = CompletionQueue(cq_impl);
+
+            auto client_context = std::make_shared<grpc::ClientContext>();
+            google::cloud::internal::ImmutableOptions options;
+            google::bigtable::v2::PingAndWarmRequest request;
+            request.set_name(instance_name);
+            // Use the client_context to set a deadline similar to
+            // AsyncWaitConnectionReady.
+            client_context->set_deadline(std::chrono::system_clock::now() +
+                                         kConnectionReadyTimeout);
+            stub->AsyncPingAndWarm(cq, client_context, std::move(options),
+                                   request)
+                .then(
+                    [weak_stub, weak_cq_impl, state, instance_name,
+                     connection_status_fn = std::move(connection_status_fn)](
+                        future<
+                            StatusOr<google::bigtable::v2::PingAndWarmResponse>>
+                            fut) mutable {
+                      auto response = fut.get();
+                      connection_status_fn(response.status());
+                      auto stub = weak_stub.lock();
+                      if (!stub) return;
+                      auto cq_impl = weak_cq_impl.lock();
+                      if (!cq_impl) return;
+                      ScheduleStubRefresh(cq_impl, state, stub, instance_name,
+                                          std::move(connection_status_fn));
+                    });
+          });
+  state->timers().RegisterTimer(std::move(timer_future));
+}
+
 void OutstandingTimers::RegisterTimer(future<void> fut) {
   std::unique_lock<std::mutex> lk(mu_);
   if (shutdown_) {

google/cloud/bigtable/internal/connection_refresh_state.h

@@ -15,6 +15,7 @@
 #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_CONNECTION_REFRESH_STATE_H
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_CONNECTION_REFRESH_STATE_H
 
+#include "google/cloud/bigtable/internal/bigtable_stub.h"
 #include "google/cloud/bigtable/version.h"
 #include "google/cloud/completion_queue.h"
 #include "google/cloud/future.h"
@@ -81,13 +82,23 @@ class ConnectionRefreshState {
 };
 
 /**
- * Schedule a chain of timers to refresh the connection.
+ * Schedule a chain of timers to refresh the grpc::Channel using
+ * AsyncWaitConnectionReady which leverages grpc::Channel::GetState.
  */
 void ScheduleChannelRefresh(
     std::shared_ptr<internal::CompletionQueueImpl> const& cq,
     std::shared_ptr<ConnectionRefreshState> const& state,
     std::shared_ptr<grpc::Channel> const& channel);
 
+/**
+ * Schedule a chain of timers to refresh the BigtableStub using PingAndWarm.
+ */
+void ScheduleStubRefresh(
+    std::shared_ptr<internal::CompletionQueueImpl> const& cq,
+    std::shared_ptr<ConnectionRefreshState> const& state,
+    std::shared_ptr<BigtableStub> const& stub, std::string const& instance_name,
+    std::function<void(Status const&)> connection_status_fn = {});
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud

google/cloud/bigtable/internal/connection_refresh_state_test.cc

@@ -12,14 +12,21 @@
 // See the License for the specific language governing permissions and
 
 #include "google/cloud/bigtable/internal/connection_refresh_state.h"
+#include "google/cloud/bigtable/testing/mock_bigtable_stub.h"
 #include "google/cloud/testing_util/status_matchers.h"
 #include <gmock/gmock.h>
+#include <thread>
 
 namespace google {
 namespace cloud {
 namespace bigtable_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+using ::google::cloud::bigtable::testing::MockBigtableStub;
+using ::google::cloud::testing_util::IsOk;
+using ::testing::Eq;
+using ::testing::MockFunction;
+
 using TimerFuture = future<StatusOr<std::chrono::system_clock::time_point>>;
 
 class OutstandingTimersTest : public ::testing::Test {
@@ -118,6 +125,71 @@ TEST(ConnectionRefreshState, Disabled) {
   EXPECT_FALSE(state.enabled());
 }
 
+class ScheduleStubRefreshTest : public ::testing::Test {
+ public:
+  ScheduleStubRefreshTest() : thread1_([this] { cq_.Run(); }) {}
+  ~ScheduleStubRefreshTest() override {
+    cq_.Shutdown();
+    thread1_.join();
+  }
+
+ protected:
+  CompletionQueue cq_;
+  std::thread thread1_;
+};
+
+TEST_F(ScheduleStubRefreshTest, RefreshedUsingAsyncPingAndWarm) {
+  auto cq_impl = internal::GetCompletionQueueImpl(cq_);
+  auto refresh_state = std::make_shared<ConnectionRefreshState>(
+      cq_impl, std::chrono::milliseconds(1), std::chrono::milliseconds(2));
+  std::string instance_name = "projects/my-project/instances/my-instance";
+
+  // These promises are used to coordinate thread execution to ensure the test
+  // does not finish before the CompletionQueue thread executes all the tasks
+  // on the queue.
+  promise<void> p;
+  promise<StatusOr<google::bigtable::v2::PingAndWarmResponse>> p2;
+
+  auto mock_stub = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock_stub, AsyncPingAndWarm)
+      .WillRepeatedly(
+          [&](CompletionQueue&, std::shared_ptr<grpc::ClientContext> const&,
+              internal::ImmutableOptions const&,
+              google::bigtable::v2::PingAndWarmRequest const& request)
+              -> future<StatusOr<google::bigtable::v2::PingAndWarmResponse>> {
+            EXPECT_THAT(request.name(), Eq(instance_name));
+            return p2.get_future();
+          });
+
+  MockFunction<void(Status const&)> mock_fn;
+  EXPECT_CALL(mock_fn, Call).WillOnce([&p](Status const& s) -> void {
+    EXPECT_THAT(s, IsOk());
+    p.set_value();
+  });
+
+  ScheduleStubRefresh(cq_impl, refresh_state, mock_stub, instance_name,
+                      mock_fn.AsStdFunction());
+  p2.set_value(google::bigtable::v2::PingAndWarmResponse{});
+  p.get_future().get();
+}
+
+TEST_F(ScheduleStubRefreshTest, RefreshTimerCancelled) {
+  auto cq_impl = internal::GetCompletionQueueImpl(cq_);
+  auto refresh_state = std::make_shared<ConnectionRefreshState>(
+      cq_impl, std::chrono::seconds(60), std::chrono::seconds(120));
+  std::string instance_name = "projects/my-project/instances/my-instance";
+
+  auto mock_stub = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock_stub, AsyncPingAndWarm).Times(0);
+
+  MockFunction<void(Status const&)> mock_fn;
+  EXPECT_CALL(mock_fn, Call).Times(0);
+
+  ScheduleStubRefresh(cq_impl, refresh_state, mock_stub, instance_name,
+                      mock_fn.AsStdFunction());
+  cq_impl->CancelAll();
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud