working integration tests with new stub and pool

Author: scotthart

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc

@@ -89,6 +89,7 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<
 BigtableRandomTwoLeastUsed::ReadRows(
     std::shared_ptr<grpc::ClientContext> context, Options const& options,
     google::bigtable::v2::ReadRowsRequest const& request) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto child = Child();
   auto stub = child->AcquireStub();
   auto result = stub->ReadRows(std::move(context), options, request);
@@ -136,6 +137,7 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<
 BigtableRandomTwoLeastUsed::MutateRows(
     std::shared_ptr<grpc::ClientContext> context, Options const& options,
     google::bigtable::v2::MutateRowsRequest const& request) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto child = Child();
   auto stub = child->AcquireStub();
   auto result = stub->MutateRows(std::move(context), options, request);
@@ -334,6 +336,7 @@ BigtableRandomTwoLeastUsed::AsyncPrepareQuery(
 
 std::shared_ptr<internal::StubWrapper<BigtableStub>>
 BigtableRandomTwoLeastUsed::Child() {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   return pool_->GetChannelRandomTwoLeastUsed();
   //  std::unique_lock<std::mutex> lk(mu_);
   //  std::vector<std::size_t> indices(pool_->size(lk) - 1);

google/cloud/bigtable/internal/bigtable_stub.cc

@@ -36,6 +36,7 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<
 DefaultBigtableStub::ReadRows(
     std::shared_ptr<grpc::ClientContext> context, Options const&,
     google::bigtable::v2::ReadRowsRequest const& request) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto stream = grpc_stub_->ReadRows(context.get(), request);
   return std::make_unique<google::cloud::internal::StreamingReadRpcImpl<
       google::bigtable::v2::ReadRowsResponse>>(std::move(context),
@@ -70,6 +71,7 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<
 DefaultBigtableStub::MutateRows(
     std::shared_ptr<grpc::ClientContext> context, Options const&,
     google::bigtable::v2::MutateRowsRequest const& request) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto stream = grpc_stub_->MutateRows(context.get(), request);
   return std::make_unique<google::cloud::internal::StreamingReadRpcImpl<
       google::bigtable::v2::MutateRowsResponse>>(std::move(context),

google/cloud/bigtable/internal/bigtable_stub_factory.cc

@@ -17,6 +17,7 @@
 #include "google/cloud/bigtable/internal/bigtable_channel_refresh.h"
 #include "google/cloud/bigtable/internal/bigtable_logging_decorator.h"
 #include "google/cloud/bigtable/internal/bigtable_metadata_decorator.h"
+#include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h"
 #include "google/cloud/bigtable/internal/bigtable_round_robin_decorator.h"
 #include "google/cloud/bigtable/internal/bigtable_tracing_stub.h"
 #include "google/cloud/bigtable/internal/connection_refresh_state.h"
@@ -73,6 +74,18 @@ std::shared_ptr<BigtableStub> CreateBigtableStubRoundRobin(
   return std::make_shared<BigtableRoundRobin>(std::move(children));
 }
 
+std::shared_ptr<BigtableStub> CreateBigtableStubRandomTwoLeastUsed(
+    Options const& options, CompletionQueue cq,
+    std::function<std::shared_ptr<BigtableStub>(int)> child_factory) {
+  std::vector<std::shared_ptr<BigtableStub>> children(
+      (std::max)(1, options.get<GrpcNumChannelsOption>()));
+  int id = 0;
+  std::generate(children.begin(), children.end(),
+                [&id, &child_factory] { return child_factory(id++); });
+  return std::make_shared<BigtableRandomTwoLeastUsed>(cq, child_factory,
+                                                      std::move(children));
+}
+
 std::shared_ptr<BigtableStub> CreateDecoratedStubs(
     std::shared_ptr<internal::GrpcAuthenticationStrategy> auth,
     CompletionQueue const& cq, Options const& options,
@@ -87,7 +100,10 @@ std::shared_ptr<BigtableStub> CreateDecoratedStubs(
     if (refresh->enabled()) ScheduleChannelRefresh(cq_impl, refresh, channel);
     return base_factory(std::move(channel));
   };
-  auto stub = CreateBigtableStubRoundRobin(options, std::move(child_factory));
+  //  auto stub = CreateBigtableStubRoundRobin(options,
+  //  std::move(child_factory));
+  auto stub = CreateBigtableStubRandomTwoLeastUsed(options, cq,
+                                                   std::move(child_factory));
   if (refresh->enabled()) {
     stub = std::make_shared<BigtableChannelRefresh>(std::move(stub),
                                                     std::move(refresh));

google/cloud/bigtable/internal/bulk_mutator.cc

@@ -214,6 +214,7 @@ grpc::Status BulkMutator::MakeOneRequest(bigtable::DataClient& client,
 Status BulkMutator::MakeOneRequest(BigtableStub& stub,
                                    MutateRowsLimiter& limiter,
                                    Options const& options) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   // Send the request to the server.
   auto const& mutations = state_.BeforeStart();
 
@@ -226,8 +227,10 @@ Status BulkMutator::MakeOneRequest(BigtableStub& stub,
   // Potentially throttle the request
   limiter.Acquire();
 
+  std::cout << __PRETTY_FUNCTION__ << ": pre-stub" << std::endl;
   // Read the stream of responses.
   auto stream = stub.MutateRows(client_context, options, mutations);
+  std::cout << __PRETTY_FUNCTION__ << ": post-stub" << std::endl;
   absl::optional<Status> status;
   while (true) {
     btproto::MutateRowsResponse response;

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -102,6 +102,7 @@ bigtable::RowReader ReadRowsHelper(
         params,  // NOLINT(performance-unnecessary-value-param)
     std::shared_ptr<OperationContext>
         operation_context) {  // NOLINT(performance-unnecessary-value-param)
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto impl = std::make_shared<DefaultRowReader>(
       stub, std::move(params.app_profile_id), std::move(params.table_name),
       std::move(params.row_set), params.rows_limit, std::move(params.filter),
@@ -338,6 +339,7 @@ future<Status> DataConnectionImpl::AsyncApply(std::string const& table_name,
 
 std::vector<bigtable::FailedMutation> DataConnectionImpl::BulkApply(
     std::string const& table_name, bigtable::BulkMutation mut) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto current = google::cloud::internal::SaveCurrentOptions();
   if (mut.empty()) return {};
   auto operation_context = operation_context_factory_->MutateRows(
@@ -350,6 +352,7 @@ std::vector<bigtable::FailedMutation> DataConnectionImpl::BulkApply(
   std::unique_ptr<bigtable::DataRetryPolicy> retry;
   std::unique_ptr<BackoffPolicy> backoff;
   Status status;
+  std::cout << __PRETTY_FUNCTION__ << ": pre-loop" << std::endl;
   while (true) {
     status = mutator.MakeOneRequest(*stub_, *limiter_, *current);
     if (!mutator.HasPendingMutations()) break;
@@ -361,6 +364,7 @@ std::vector<bigtable::FailedMutation> DataConnectionImpl::BulkApply(
     if (!delay) break;
     std::this_thread::sleep_for(*delay);
   }
+  std::cout << __PRETTY_FUNCTION__ << ": post-loop" << std::endl;
   operation_context->OnDone(status);
   return std::move(mutator).OnRetryDone();
 }
@@ -380,6 +384,7 @@ DataConnectionImpl::AsyncBulkApply(std::string const& table_name,
 
 bigtable::RowReader DataConnectionImpl::ReadRowsFull(
     bigtable::ReadRowsParams params) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto current = google::cloud::internal::SaveCurrentOptions();
   auto operation_context = operation_context_factory_->ReadRows(
       params.table_name, params.app_profile_id);
@@ -660,6 +665,7 @@ void DataConnectionImpl::AsyncReadRows(
     std::function<future<bool>(bigtable::Row)> on_row,
     std::function<void(Status)> on_finish, bigtable::RowSet row_set,
     std::int64_t rows_limit, bigtable::Filter filter) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto current = google::cloud::internal::SaveCurrentOptions();
   auto operation_context = operation_context_factory_->ReadRows(
       table_name, app_profile_id(*current));

google/cloud/bigtable/table.cc

@@ -222,6 +222,7 @@ RowReader Table::ReadRows(RowSet row_set, Filter filter, Options opts) {
 
 RowReader Table::ReadRows(RowSet row_set, std::int64_t rows_limit,
                           Filter filter, Options opts) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   if (connection_) {
     OptionsSpan span(MergeOptions(std::move(opts), options_));
     return connection_->ReadRows(table_name_, std::move(row_set), rows_limit,

google/cloud/bigtable/testing/table_integration_test.cc

@@ -122,6 +122,7 @@ void TableAdminTestEnvironment::TearDown() {
 }
 
 void TableIntegrationTest::SetUp() {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   data_connection_ = MakeDataConnection();
   data_client_ = bigtable::MakeDataClient(TableTestEnvironment::project_id(),
                                           TableTestEnvironment::instance_id());

google/cloud/bigtable/tests/data_integration_test.cc

@@ -206,6 +206,7 @@ TEST_P(DataIntegrationTest, TableReadRowNotExistTest) {
 }
 
 TEST_P(DataIntegrationTest, TableReadRowsAllRows) {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
   auto table = GetTable(GetParam());
   std::string const row_key1 = "row-key-1";
   std::string const row_key2 = "row-key-2";

google/cloud/internal/channel_pool.h

@@ -74,27 +74,30 @@ template <typename T>
 class DynamicChannelPool
     : public std::enable_shared_from_this<DynamicChannelPool<T>> {
  public:
-  using StubFactoryFn = std::function<std::shared_ptr<T>()>;
+  using StubFactoryFn = std::function<std::shared_ptr<T>(int id)>;
   struct SizingPolicy {
     // To avoid channel churn, the pool will not add or remove channels more
     // frequently that this period.
-    std::chrono::milliseconds resize_cooldown_interval;
+    std::chrono::milliseconds resize_cooldown_interval =
+        std::chrono::milliseconds(60 * 1000);
 
     // If the average number of outstanding RPCs is below this threshold,
     // the pool size will be decreased.
-    std::size_t minimum_average_outstanding_rpcs_per_channel;
+    std::size_t minimum_average_outstanding_rpcs_per_channel = 20;
     // If the average number of outstanding RPCs is above this threshold,
     // the pool size will be increased.
-    std::size_t maximum_average_outstanding_rpcs_per_channel;
+    std::size_t maximum_average_outstanding_rpcs_per_channel = 80;
 
     // When channels are removed from the pool, we have to wait until all
     // outstanding RPCs on that channel are completed before destroying it.
-    std::chrono::milliseconds decommissioned_channel_polling_interval;
+    std::chrono::milliseconds decommissioned_channel_polling_interval =
+        std::chrono::milliseconds(30 * 1000);
   };
 
   static std::shared_ptr<DynamicChannelPool> Create(
       CompletionQueue cq, std::size_t initial_size, StubFactoryFn factory_fn,
       SizingPolicy sizing_policy = {}) {
+    std::cout << __PRETTY_FUNCTION__ << std::endl;
     std::vector<std::shared_ptr<StubWrapper<T>>> initial_wrapped_channels;
     for (std::size_t i = 0; i < initial_size; ++i) {
       initial_wrapped_channels.emplace_back(factory_fn());
@@ -103,9 +106,11 @@ class DynamicChannelPool
         std::move(cq), std::move(initial_wrapped_channels),
         std::move(factory_fn), std::move(sizing_policy)));
   }
+
   static std::shared_ptr<DynamicChannelPool> Create(
       CompletionQueue cq, std::vector<std::shared_ptr<T>> initial_channels,
       StubFactoryFn 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(factory_fn),
         std::move(sizing_policy)));
@@ -132,16 +137,20 @@ class DynamicChannelPool
   //  }
 
   std::shared_ptr<StubWrapper<T>> GetChannelRandomTwoLeastUsed() {
+    std::cout << __PRETTY_FUNCTION__ << std::endl;
     std::unique_lock<std::mutex> lk(mu_);
 
+    std::cout << __PRETTY_FUNCTION__ << ": channels_size()=" << channels_.size()
+              << std::endl;
     // TODO: check if resize is needed.
 
     std::vector<std::size_t> indices(channels_.size());
     // TODO(sdhart): Maybe use iota on iterators instead of indices
     std::iota(indices.begin(), indices.end(), 0);
     std::shuffle(indices.begin(), indices.end(), rng_);
-    auto channel_1 = channels_[indices[0]];
-    auto channel_2 = channels_[indices[1]];
+
+    std::shared_ptr<StubWrapper<T>> channel_1 = channels_[indices[0]];
+    std::shared_ptr<StubWrapper<T>> channel_2 = channels_[indices[1]];
 
     return channel_1->outstanding_rpcs(lk) < channel_2->outstanding_rpcs(lk)
                ? channel_1
@@ -156,18 +165,26 @@ class DynamicChannelPool
       : cq_(std::move(cq)),
         factory_fn_(std::move(factory_fn)),
         channels_(std::move(initial_wrapped_channels)),
-        sizing_policy_(std::move(sizing_policy)) {}
+        sizing_policy_(std::move(sizing_policy)),
+        next_channel_id_(channels_.size()) {}
 
   DynamicChannelPool(CompletionQueue cq,
                      std::vector<std::shared_ptr<T>> initial_channels,
                      StubFactoryFn factory_fn, SizingPolicy sizing_policy)
       : cq_(std::move(cq)),
         factory_fn_(std::move(factory_fn)),
-        channels_(initial_channels.size()),
-        sizing_policy_(std::move(sizing_policy)) {
+        channels_(),
+        sizing_policy_(std::move(sizing_policy)),
+        next_channel_id_(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<StubWrapper<T>>(channel));
+      channels_.push_back(std::make_shared<StubWrapper<T>>(std::move(channel)));
     }
+    //    for (auto i = 0; i < channels_.size(); ++i) {
+    //        std::cout << __PRETTY_FUNCTION__ << ": channels_[" << i <<
+    //        "].get()=" << channels_[i].get() << std::endl;
+    //    }
   }
 
   void ScheduleAddChannel() {}
@@ -234,6 +251,7 @@ class DynamicChannelPool
   SizingPolicy sizing_policy_;
   std::vector<std::shared_ptr<StubWrapper<T>>> decommissioned_channels_;
   future<StatusOr<std::chrono::system_clock::time_point>> decommission_timer_;
+  int next_channel_id_;
 };
 
 }  // namespace internal