connection refresh callback sets status

Author: scotthart

ci/cloudbuild/builds/integration-production.sh

@@ -39,8 +39,8 @@ excluded_rules=(
   "-//google/cloud/storagecontrol:v2_samples_storage_control_anywhere_cache_samples"
 )
 
+#   --test_filter="*ReadRowsAllRows*" --test_timeout=30 \
 io::log_h2 "Running the integration tests against prod"
 mapfile -t integration_args < <(integration::bazel_args)
 io::run bazel test "${args[@]}" "${integration_args[@]}" --test_output=all \
-  --test_filter="*ReadRowsAllRows*" --test_timeout=30 \
   //google/cloud/bigtable/tests:data_integration_test

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc

@@ -73,15 +73,6 @@ class AsyncStreamingReadRpcTracking
 
 }  // namespace
 
-BigtableRandomTwoLeastUsed::BigtableRandomTwoLeastUsed(
-    CompletionQueue cq, std::shared_ptr<ConnectionRefreshState> refresh_state,
-    DynamicChannelPool<BigtableStub>::StubFactoryFn
-        refreshing_channel_stub_factory_fn,
-    std::vector<std::shared_ptr<BigtableStub>> children)
-    : pool_(DynamicChannelPool<BigtableStub>::Create(
-          std::move(cq), std::move(children), std::move(refresh_state),
-          std::move(refreshing_channel_stub_factory_fn))) {}
-
 std::unique_ptr<google::cloud::internal::StreamingReadRpc<
     google::bigtable::v2::ReadRowsResponse>>
 BigtableRandomTwoLeastUsed::ReadRows(

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h

@@ -29,11 +29,10 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 class BigtableRandomTwoLeastUsed : public BigtableStub {
  public:
-  BigtableRandomTwoLeastUsed(
-      CompletionQueue cq, std::shared_ptr<ConnectionRefreshState> refresh_state,
-      DynamicChannelPool<BigtableStub>::StubFactoryFn
-          refreshing_channel_stub_factory_fn,
-      std::vector<std::shared_ptr<BigtableStub>> children);
+  explicit BigtableRandomTwoLeastUsed(
+      std::shared_ptr<DynamicChannelPool<BigtableStub>> pool)
+      : pool_(std::move(pool)) {}
+
   ~BigtableRandomTwoLeastUsed() override = default;
 
   std::unique_ptr<google::cloud::internal::StreamingReadRpc<

google/cloud/bigtable/internal/bigtable_stub_factory.cc

@@ -77,23 +77,47 @@ std::shared_ptr<BigtableStub> CreateBigtableStubRoundRobin(
 }
 
 std::shared_ptr<BigtableStub> CreateBigtableStubRandomTwoLeastUsed(
-    CompletionQueue cq, Options const& options,
-    std::function<std::shared_ptr<BigtableStub>(int)>
-        refreshing_channel_stub_factory,
+    std::shared_ptr<internal::GrpcAuthenticationStrategy> auth,
+    std::shared_ptr<internal::CompletionQueueImpl> cq_impl,
+    Options const& options, BaseBigtableStubFactory stub_factory,
     std::shared_ptr<ConnectionRefreshState> refresh_state) {
   std::cout << __PRETTY_FUNCTION__ << std::endl;
-  std::vector<std::shared_ptr<BigtableStub>> children(
-      (std::max)(1, options.get<GrpcNumChannelsOption>()));
-  int id = 0;
+
+  auto refreshing_channel_stub_factory =
+      [stub_factory = std::move(stub_factory), cq_impl, refresh_state,
+       auth = std::move(auth), options](std::uint32_t id)
+      -> std::shared_ptr<ChannelUsageWrapper<BigtableStub>> {
+    auto wrapper = std::make_shared<ChannelUsageWrapper<BigtableStub>>();
+    auto connection_status_fn = [weak = wrapper->MakeWeak()](Status const& s) {
+      if (auto self = weak.lock()) {
+        self->SetLastRefreshStatus(s);
+      }
+      if (!s.ok()) {
+        GCP_LOG(WARNING) << "Failed to refresh connection. Error: " << s;
+      }
+    };
+    auto channel = CreateGrpcChannel(*auth, options, id);
+    ScheduleChannelRefresh(cq_impl, refresh_state, channel,
+                           std::move(connection_status_fn));
+    wrapper->set_channel(stub_factory(std::move(channel)));
+    return wrapper;
+  };
+
+  std::vector<std::shared_ptr<ChannelUsageWrapper<BigtableStub>>> children(
+      std::max(1, options.get<GrpcNumChannelsOption>()));
+  std::uint32_t id = 0;
   std::generate(children.begin(), children.end(),
                 [&id, &refreshing_channel_stub_factory] {
                   return refreshing_channel_stub_factory(id++);
                 });
-  std::cout << __PRETTY_FUNCTION__ << ": children.size()=" << children.size()
-            << std::endl;
+
   return std::make_shared<BigtableRandomTwoLeastUsed>(
-      std::move(cq), std::move(refresh_state), refreshing_channel_stub_factory,
-      std::move(children));
+      DynamicChannelPool<BigtableStub>::Create(
+          CompletionQueue(std::move(cq_impl)), std::move(children),
+          std::move(refresh_state),
+          std::move(refreshing_channel_stub_factory)));
+  //      CompletionQueue(cq_impl), std::move(refresh_state),
+  //      std::move(refreshing_channel_stub_factory), std::move(children));
 }
 
 std::shared_ptr<BigtableStub> CreateDecoratedStubs(
@@ -109,14 +133,9 @@ std::shared_ptr<BigtableStub> CreateDecoratedStubs(
   if (options.has<bigtable::experimental::ChannelPoolTypeOption>() &&
       options.get<bigtable::experimental::ChannelPoolTypeOption>() ==
           bigtable::experimental::ChannelPoolType::kDynamic) {
-    auto refreshing_channel_stub_factory = [stub_factory, cq_impl, refresh,
-                                            &auth, options](int id) {
-      auto channel = CreateGrpcChannel(*auth, options, id);
-      ScheduleChannelRefresh(cq_impl, refresh, channel);
-      return stub_factory(std::move(channel));
-    };
     stub = CreateBigtableStubRandomTwoLeastUsed(
-        cq, options, std::move(refreshing_channel_stub_factory),
+        auth, std::move(cq_impl), options, stub_factory,
+        //        std::move(refreshing_channel_stub_factory),
         std::move(refresh));
   } else {
     auto refreshing_channel_stub_factory = [stub_factory, cq_impl, refresh,

google/cloud/bigtable/internal/bigtable_stub_factory.h

@@ -43,9 +43,11 @@ std::shared_ptr<BigtableStub> CreateBigtableStubRoundRobin(
                                 refreshing_channel_stub_factory);
 
 std::shared_ptr<BigtableStub> CreateBigtableStubRandomTwoLeastUsed(
-    CompletionQueue cq, Options const& options,
-    std::function<std::shared_ptr<BigtableStub>(int)>
-        refreshing_channel_stub_factory,
+    std::shared_ptr<internal::GrpcAuthenticationStrategy> auth,
+    std::shared_ptr<internal::CompletionQueueImpl> cq_impl,
+    Options const& options, BaseBigtableStubFactory stub_factory,
+    //    std::function<std::shared_ptr<BigtableStub>(int)>
+    //        refreshing_channel_stub_factory,
     std::shared_ptr<ConnectionRefreshState> refresh_state);
 
 /// Used in testing to create decorated mocks.

google/cloud/bigtable/internal/connection_refresh_state.cc

@@ -95,8 +95,7 @@ void ScheduleChannelRefresh(
                        connection_status_fn = std::move(connection_status_fn)](
                           future<Status> fut) {
                   auto conn_status = fut.get();
-                  if (connection_status_fn) connection_status_fn(conn_status);
-
+                  connection_status_fn(conn_status);
                   auto channel = weak_channel.lock();
                   if (!channel) return;
                   auto cq_impl = weak_cq_impl.lock();

google/cloud/bigtable/internal/dynamic_channel_pool.h

@@ -34,6 +34,7 @@ template <typename T>
 class ChannelUsageWrapper
     : public std::enable_shared_from_this<ChannelUsageWrapper<T>> {
  public:
+  ChannelUsageWrapper() = default;
   explicit ChannelUsageWrapper(std::shared_ptr<T> stub)
       : stub_(std::move(stub)) {}
 
@@ -55,6 +56,11 @@ class ChannelUsageWrapper
     last_refresh_status_ = std::move(s);
   }
 
+  ChannelUsageWrapper& set_channel(std::shared_ptr<T> channel) {
+    stub_ = std::move(channel);
+    return *this;
+  }
+
   std::weak_ptr<ChannelUsageWrapper<T>> MakeWeak() {
     return this->shared_from_this();
   }
@@ -81,7 +87,8 @@ template <typename T>
 class DynamicChannelPool
     : public std::enable_shared_from_this<DynamicChannelPool<T>> {
  public:
-  using StubFactoryFn = std::function<std::shared_ptr<T>(int id)>;
+  using StubFactoryFn =
+      std::function<std::shared_ptr<ChannelUsageWrapper<T>>(std::uint32_t id)>;
   struct SizingPolicy {
     // To avoid channel churn, the pool will not add or remove channels more
     // frequently that this period.
@@ -117,26 +124,8 @@ class DynamicChannelPool
   };
 
   static std::shared_ptr<DynamicChannelPool> Create(
-      CompletionQueue cq, std::size_t initial_size,
-      StubFactoryFn stub_factory_fn,
-      std::shared_ptr<ConnectionRefreshState> refresh_state,
-      SizingPolicy sizing_policy = {}) {
-    std::cout << __PRETTY_FUNCTION__ << ": enter" << std::endl;
-    std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>
-        initial_wrapped_channels;
-    for (std::size_t i = 0; i < initial_size; ++i) {
-      initial_wrapped_channels.emplace_back(stub_factory_fn());
-    }
-    auto pool = std::shared_ptr<DynamicChannelPool>(new DynamicChannelPool(
-        std::move(cq), std::move(initial_wrapped_channels),
-        std::move(refresh_state), std::move(stub_factory_fn),
-        std::move(sizing_policy)));
-    std::cout << __PRETTY_FUNCTION__ << ": return pool" << std::endl;
-    return pool;
-  }
-
-  static std::shared_ptr<DynamicChannelPool> Create(
-      CompletionQueue cq, std::vector<std::shared_ptr<T>> initial_channels,
+      CompletionQueue cq,
+      std::vector<std::shared_ptr<ChannelUsageWrapper<T>>> initial_channels,
       std::shared_ptr<ConnectionRefreshState> refresh_state,
       StubFactoryFn stub_factory_fn, SizingPolicy sizing_policy = {}) {
     std::cout << __PRETTY_FUNCTION__ << ": enter" << std::endl;
@@ -193,26 +182,22 @@ class DynamicChannelPool
     std::iota(iterators.begin(), iterators.end(), channels_.begin());
     std::shuffle(iterators.begin(), iterators.end(), rng_);
 
-    //    std::cout << __PRETTY_FUNCTION__ << ": shuffled iterators" <<
-    //    std::endl;
-    // std::vector<
-    //       typename
-    //       std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator>
-    //       iterators;
-    typename std::vector<typename std::vector<
-        std::shared_ptr<ChannelUsageWrapper<T>>>::iterator>::iterator
-        shuffle_iter = iterators.begin();
-    typename std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator
-        channel_1 = *shuffle_iter;
+    //    typename std::vector<typename std::vector<
+    //        std::shared_ptr<ChannelUsageWrapper<T>>>::iterator>::iterator
+    auto shuffle_iter = iterators.begin();
+    //    typename
+    //    std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator
+    auto channel_1 = *shuffle_iter;
     std::shared_ptr<ChannelUsageWrapper<T>> c = *channel_1;
     //    std::cout << __PRETTY_FUNCTION__
     //              << ": check channel 1=" << c.get() << std::endl;
     auto channel_1_rpcs = shuffle_iter != iterators.end()
                               ? (*channel_1)->outstanding_rpcs()
                               : Status{StatusCode::kNotFound, ""};
     ++shuffle_iter;
-    typename std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator
-        channel_2 = *shuffle_iter;
+    //    typename
+    //    std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator
+    auto channel_2 = *shuffle_iter;
     // We want to snapshot these outstanding_rpcs values.
     //    std::cout << __PRETTY_FUNCTION__
     //              << ": check channel 2=" << (channel_2)->get() << std::endl;
@@ -264,7 +249,7 @@ class DynamicChannelPool
       return *channel_2;
     }
 
-    // TODO: we have no usable channels in the entire pool; this is bad.
+    // TODO(sdhart): we have no usable channels in the entire pool; this is bad.
     std::cout << __PRETTY_FUNCTION__ << ": NO USABLE CHANNELS" << std::endl;
     return nullptr;
   }
@@ -280,26 +265,7 @@ class DynamicChannelPool
         stub_factory_fn_(std::move(stub_factory_fn)),
         channels_(std::move(initial_wrapped_channels)),
         sizing_policy_(std::move(sizing_policy)),
-        next_channel_id_(channels_.size()) {
-    sizing_policy_.minimum_channel_pool_size = channels_.size();
-  }
-
-  DynamicChannelPool(CompletionQueue cq,
-                     std::vector<std::shared_ptr<T>> initial_channels,
-                     std::shared_ptr<ConnectionRefreshState> refresh_state,
-                     StubFactoryFn stub_factory_fn, SizingPolicy sizing_policy)
-      : cq_(std::move(cq)),
-        refresh_state_(std::move(refresh_state)),
-        stub_factory_fn_(std::move(stub_factory_fn)),
-        channels_(),
-        sizing_policy_(std::move(sizing_policy)),
-        next_channel_id_(static_cast<int>(initial_channels.size())) {
-    std::cout << __PRETTY_FUNCTION__ << ": wrap initial_channels" << std::endl;
-    channels_.reserve(initial_channels.size());
-    for (auto& channel : initial_channels) {
-      channels_.push_back(
-          std::make_shared<ChannelUsageWrapper<T>>(std::move(channel)));
-    }
+        next_channel_id_(static_cast<std::uint32_t>(channels_.size())) {
     sizing_policy_.minimum_channel_pool_size = channels_.size();
   }
 
@@ -348,8 +314,8 @@ class DynamicChannelPool
     std::vector<std::shared_ptr<ChannelUsageWrapper<T>>> new_stubs;
     new_stubs.reserve(new_channel_ids.size());
     for (auto const& id : new_channel_ids) {
-      new_stubs.push_back(
-          std::make_shared<ChannelUsageWrapper<T>>(stub_factory_fn_(id)));
+      new_stubs.push_back(stub_factory_fn_(id));
+      //          std::make_shared<ChannelUsageWrapper<T>>(stub_factory_fn_(id)));
     }
     std::unique_lock<std::mutex> lk(mu_);
     channels_.insert(channels_.end(),
@@ -391,7 +357,8 @@ class DynamicChannelPool
       }
       draining_channels_.pop_back();
     }
-    // TODO: If iterators becomes a member variable perhaps add logic to call
+    // TODO(sdhart): If iterators becomes a member variable perhaps add logic to
+    // call
     //  shrink_to_fit on iterators_ if there's a large
     // difference between iterators_.capacity and channels_.size
   }
@@ -446,7 +413,7 @@ class DynamicChannelPool
   future<void> remove_channel_poll_timer_;
   absl::optional<future<StatusOr<std::chrono::system_clock::time_point>>>
       pool_resize_cooldown_timer_ = absl::nullopt;
-  int next_channel_id_;
+  std::uint32_t next_channel_id_;
   //  std::vector<
   //      typename
   //      std::vector<std::shared_ptr<ChannelUsageWrapper<T>>>::iterator>

google/cloud/bigtable/internal/dynamic_channel_pool_test.cc

@@ -34,13 +34,15 @@ TEST(DynamicChannelPoolTest, GetChannelRandomTwoLeastUsed) {
       fake_cq_impl, std::chrono::milliseconds(1),
       std::chrono::milliseconds(10));
 
-  auto stub_factory_fn = [](int) -> std::shared_ptr<BigtableStub> {
-    return std::make_shared<MockBigtableStub>();
+  auto stub_factory_fn =
+      [](int) -> std::shared_ptr<ChannelUsageWrapper<BigtableStub>> {
+    auto mock = std::make_shared<MockBigtableStub>();
+    return std::make_shared<ChannelUsageWrapper<BigtableStub>>(mock);
   };
 
   DynamicChannelPool<BigtableStub>::SizingPolicy sizing_policy;
 
-  std::vector<std::shared_ptr<BigtableStub>> channels(10);
+  std::vector<std::shared_ptr<ChannelUsageWrapper<BigtableStub>>> channels(10);
   int id = 0;
   std::generate(channels.begin(), channels.end(),
                 [&]() { return stub_factory_fn(id++); });