make refresh_state a member of channel pool

Author: scotthart

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc

@@ -74,12 +74,12 @@ class AsyncStreamingReadRpcTracking
 }  // namespace
 
 BigtableRandomTwoLeastUsed::BigtableRandomTwoLeastUsed(
-    CompletionQueue cq,
+    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(cq), std::move(children), std::move(refresh_state),
           std::move(refreshing_channel_stub_factory_fn))) {}
 
 std::unique_ptr<google::cloud::internal::StreamingReadRpc<

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h

@@ -30,7 +30,7 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 class BigtableRandomTwoLeastUsed : public BigtableStub {
  public:
   BigtableRandomTwoLeastUsed(
-      CompletionQueue cq,
+      CompletionQueue cq, std::shared_ptr<ConnectionRefreshState> refresh_state,
       DynamicChannelPool<BigtableStub>::StubFactoryFn
           refreshing_channel_stub_factory_fn,
       std::vector<std::shared_ptr<BigtableStub>> children);

google/cloud/bigtable/internal/bigtable_stub_factory.cc

@@ -79,7 +79,8 @@ 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) {
+        refreshing_channel_stub_factory,
+    std::shared_ptr<ConnectionRefreshState> refresh_state) {
   std::vector<std::shared_ptr<BigtableStub>> children(
       (std::max)(1, options.get<GrpcNumChannelsOption>()));
   int id = 0;
@@ -88,7 +89,8 @@ std::shared_ptr<BigtableStub> CreateBigtableStubRandomTwoLeastUsed(
                   return refreshing_channel_stub_factory(id++);
                 });
   return std::make_shared<BigtableRandomTwoLeastUsed>(
-      std::move(cq), refreshing_channel_stub_factory, std::move(children));
+      std::move(cq), std::move(refresh_state), refreshing_channel_stub_factory,
+      std::move(children));
 }
 
 std::shared_ptr<BigtableStub> CreateDecoratedStubs(
@@ -99,28 +101,35 @@ std::shared_ptr<BigtableStub> CreateDecoratedStubs(
   auto refresh = std::make_shared<ConnectionRefreshState>(
       cq_impl, options.get<bigtable::MinConnectionRefreshOption>(),
       options.get<bigtable::MaxConnectionRefreshOption>());
-  auto refreshing_channel_stub_factory = [stub_factory, cq_impl, refresh, &auth,
-                                          options](int id) {
-    auto channel = CreateGrpcChannel(*auth, options, id);
-    if (refresh->enabled()) ScheduleChannelRefresh(cq_impl, refresh, channel);
-    return stub_factory(std::move(channel));
-  };
 
   std::shared_ptr<BigtableStub> stub;
   if (options.has<ChannelSelectionStrategyOption>() &&
       options.get<ChannelSelectionStrategyOption>() ==
           ChannelSelectionStrategy::kRandomTwoLeastUsed) {
+    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));
+        cq, options, std::move(refreshing_channel_stub_factory),
+        std::move(refresh));
   } else {
+    auto refreshing_channel_stub_factory = [stub_factory, cq_impl, refresh,
+                                            &auth, options](int id) {
+      auto channel = CreateGrpcChannel(*auth, options, id);
+      if (refresh->enabled()) ScheduleChannelRefresh(cq_impl, refresh, channel);
+      return stub_factory(std::move(channel));
+    };
     stub = CreateBigtableStubRoundRobin(
         options, std::move(refreshing_channel_stub_factory));
+    if (refresh->enabled()) {
+      stub = std::make_shared<BigtableChannelRefresh>(std::move(stub),
+                                                      std::move(refresh));
+    }
   }
 
-  if (refresh->enabled()) {
-    stub = std::make_shared<BigtableChannelRefresh>(std::move(stub),
-                                                    std::move(refresh));
-  }
   if (auth->RequiresConfigureContext()) {
     stub = std::make_shared<BigtableAuth>(std::move(auth), std::move(stub));
   }

google/cloud/bigtable/internal/bigtable_stub_factory.h

@@ -16,6 +16,7 @@
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_STUB_FACTORY_H
 
 #include "google/cloud/bigtable/internal/bigtable_stub.h"
+#include "google/cloud/bigtable/internal/connection_refresh_state.h"
 #include "google/cloud/completion_queue.h"
 #include "google/cloud/internal/unified_grpc_credentials.h"
 #include "google/cloud/options.h"
@@ -44,7 +45,8 @@ 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);
+        refreshing_channel_stub_factory,
+    std::shared_ptr<ConnectionRefreshState> refresh_state);
 
 /// Used in testing to create decorated mocks.
 std::shared_ptr<BigtableStub> CreateDecoratedStubs(

google/cloud/bigtable/internal/dynamic_channel_pool.h

@@ -15,6 +15,7 @@
 #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_DYNAMIC_CHANNEL_POOL_H
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_DYNAMIC_CHANNEL_POOL_H
 
+#include "google/cloud/bigtable/internal/connection_refresh_state.h"
 #include "google/cloud/completion_queue.h"
 #include "google/cloud/internal/random.h"
 #include "google/cloud/version.h"
@@ -98,24 +99,28 @@ class DynamicChannelPool
 
   static std::shared_ptr<DynamicChannelPool> Create(
       CompletionQueue cq, std::size_t initial_size,
-      StubFactoryFn stub_factory_fn, SizingPolicy sizing_policy = {}) {
+      StubFactoryFn stub_factory_fn,
+      std::shared_ptr<ConnectionRefreshState> refresh_state,
+      SizingPolicy sizing_policy = {}) {
     std::cout << __PRETTY_FUNCTION__ << std::endl;
     std::vector<std::shared_ptr<StubUsageWrapper<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(stub_factory_fn), std::move(sizing_policy)));
+        std::move(refresh_state), std::move(stub_factory_fn),
+        std::move(sizing_policy)));
   }
 
   static std::shared_ptr<DynamicChannelPool> Create(
       CompletionQueue cq, std::vector<std::shared_ptr<T>> initial_channels,
+      std::shared_ptr<ConnectionRefreshState> refresh_state,
       StubFactoryFn stub_factory_fn, SizingPolicy sizing_policy = {}) {
     std::cout << __PRETTY_FUNCTION__ << std::endl;
     auto pool = std::shared_ptr<DynamicChannelPool>(new DynamicChannelPool(
-        std::move(cq), std::move(initial_channels), std::move(stub_factory_fn),
-        std::move(sizing_policy)));
+        std::move(cq), std::move(initial_channels), std::move(refresh_state),
+        std::move(stub_factory_fn), std::move(sizing_policy)));
     return pool;
   }
 
@@ -166,17 +171,21 @@ class DynamicChannelPool
   DynamicChannelPool(CompletionQueue cq,
                      std::vector<std::shared_ptr<StubUsageWrapper<T>>>
                          initial_wrapped_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_(std::move(initial_wrapped_channels)),
         sizing_policy_(std::move(sizing_policy)),
         next_channel_id_(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)),
@@ -299,6 +308,7 @@ class DynamicChannelPool
   mutable std::mutex mu_;
   CompletionQueue cq_;
   google::cloud::internal::DefaultPRNG rng_;
+  std::shared_ptr<ConnectionRefreshState> refresh_state_;
   StubFactoryFn stub_factory_fn_;
   std::vector<std::shared_ptr<StubUsageWrapper<T>>> channels_;
   SizingPolicy sizing_policy_;