diff --git a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc index dc840c820840e..f13530b023ad0 100644 --- a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc @@ -108,9 +108,30 @@ class AsyncStreamingReadWriteRpcTracking template Response UnaryHelper(std::shared_ptr>& pool, + OperationContext& oc, std::function fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); + oc.StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kUnary}); + std::shared_ptr stub = selection.channel->AcquireStub(); + Response result = fn(*stub); + selection.channel->ReleaseStub(); + return result; +} + +template +Response AsyncHelper(std::shared_ptr>& pool, + std::shared_ptr const& operation_context, + std::function fn) { + SelectedChannel selection = + pool->GetChannelRandomTwoLeastUsed(); + if (operation_context != nullptr) { + operation_context->StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kUnary}); + } std::shared_ptr stub = selection.channel->AcquireStub(); Response result = fn(*stub); selection.channel->ReleaseStub(); @@ -121,11 +142,17 @@ template std::unique_ptr> StreamingHelper( std::shared_ptr>& pool, + std::shared_ptr const& operation_context, std::function>(BigtableStub&)> fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); + if (operation_context != nullptr) { + operation_context->StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kStreaming}); + } std::shared_ptr stub = selection.channel->AcquireStub(); std::unique_ptr> result = fn(*stub); auto release_fn = [weak = selection.channel->MakeWeak()] { @@ -140,12 +167,18 @@ template std::unique_ptr> AsyncStreamingHelper( std::shared_ptr>& pool, + std::shared_ptr const& operation_context, std::function>( BigtableStub&)> fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); + if (operation_context != nullptr) { + operation_context->StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kStreaming}); + } std::shared_ptr stub = selection.channel->AcquireStub(); std::unique_ptr> result = fn(*stub); auto release_fn = [weak = selection.channel->MakeWeak()] { @@ -160,11 +193,17 @@ template std::unique_ptr> AsyncStreamingHelper( std::shared_ptr>& pool, + std::shared_ptr const& operation_context, std::function>(BigtableStub&)> fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); + if (operation_context != nullptr) { + operation_context->StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kStreaming}); + } std::shared_ptr stub = selection.channel->AcquireStub(); std::unique_ptr> result = fn(*stub); @@ -186,9 +225,9 @@ BigtableRandomTwoLeastUsed::ReadRows( google::bigtable::v2::ReadRowsRequest const& request, std::shared_ptr operation_context) { return StreamingHelper( - pool_, [&, context = std::move(context), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), + operation_context](BigtableStub& stub) mutable { return stub.ReadRows(std::move(context), options, request, std::move(operation_context)); }); @@ -201,9 +240,9 @@ BigtableRandomTwoLeastUsed::SampleRowKeys( google::bigtable::v2::SampleRowKeysRequest const& request, std::shared_ptr operation_context) { return StreamingHelper( - pool_, [&, context = std::move(context), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), + operation_context](BigtableStub& stub) mutable { return stub.SampleRowKeys(std::move(context), options, request, std::move(operation_context)); }); @@ -215,7 +254,7 @@ BigtableRandomTwoLeastUsed::MutateRow( google::bigtable::v2::MutateRowRequest const& request, OperationContext& operation_context) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.MutateRow(context, options, request, operation_context); }); } @@ -227,9 +266,9 @@ BigtableRandomTwoLeastUsed::MutateRows( google::bigtable::v2::MutateRowsRequest const& request, std::shared_ptr operation_context) { return StreamingHelper( - pool_, [&, context = std::move(context), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), + operation_context](BigtableStub& stub) mutable { return stub.MutateRows(std::move(context), options, request, std::move(operation_context)); }); @@ -241,7 +280,7 @@ BigtableRandomTwoLeastUsed::CheckAndMutateRow( google::bigtable::v2::CheckAndMutateRowRequest const& request, OperationContext& operation_context) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.CheckAndMutateRow(context, options, request, operation_context); }); @@ -253,7 +292,7 @@ BigtableRandomTwoLeastUsed::PingAndWarm( google::bigtable::v2::PingAndWarmRequest const& request, OperationContext& operation_context) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.PingAndWarm(context, options, request, operation_context); }); } @@ -265,7 +304,7 @@ BigtableRandomTwoLeastUsed::ReadModifyWriteRow( OperationContext& operation_context) { return UnaryHelper< StatusOr>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.ReadModifyWriteRow(context, options, request, operation_context); }); @@ -277,7 +316,7 @@ BigtableRandomTwoLeastUsed::PrepareQuery( google::bigtable::v2::PrepareQueryRequest const& request, OperationContext& operation_context) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.PrepareQuery(context, options, request, operation_context); }); } @@ -289,9 +328,9 @@ BigtableRandomTwoLeastUsed::ExecuteQuery( google::bigtable::v2::ExecuteQueryRequest const& request, std::shared_ptr operation_context) { return StreamingHelper( - pool_, [&, context = std::move(context), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), + operation_context](BigtableStub& stub) mutable { return stub.ExecuteQuery(std::move(context), options, request, std::move(operation_context)); }); @@ -306,9 +345,9 @@ BigtableRandomTwoLeastUsed::AsyncReadRows( google::bigtable::v2::ReadRowsRequest const& request, std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncReadRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); }); @@ -323,9 +362,9 @@ BigtableRandomTwoLeastUsed::AsyncSampleRowKeys( google::bigtable::v2::SampleRowKeysRequest const& request, std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncSampleRowKeys(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -339,10 +378,10 @@ BigtableRandomTwoLeastUsed::AsyncMutateRow( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, std::shared_ptr operation_context) { - return UnaryHelper>>( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + return AsyncHelper>>( + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); }); @@ -357,9 +396,9 @@ BigtableRandomTwoLeastUsed::AsyncMutateRows( google::bigtable::v2::MutateRowsRequest const& request, std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncMutateRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); }); @@ -372,11 +411,11 @@ BigtableRandomTwoLeastUsed::AsyncCheckAndMutateRow( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, std::shared_ptr operation_context) { - return UnaryHelper< + return AsyncHelper< future>>( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncCheckAndMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -390,11 +429,11 @@ BigtableRandomTwoLeastUsed::AsyncPingAndWarm( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, std::shared_ptr operation_context) { - return UnaryHelper< + return AsyncHelper< future>>( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncPingAndWarm(cq, std::move(context), std::move(options), request, std::move(operation_context)); }); @@ -407,11 +446,11 @@ BigtableRandomTwoLeastUsed::AsyncReadModifyWriteRow( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, std::shared_ptr operation_context) { - return UnaryHelper< + return AsyncHelper< future>>( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncReadModifyWriteRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -425,11 +464,11 @@ BigtableRandomTwoLeastUsed::AsyncPrepareQuery( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, std::shared_ptr operation_context) { - return UnaryHelper< + return AsyncHelper< future>>( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncPrepareQuery(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -442,7 +481,7 @@ BigtableRandomTwoLeastUsed::GetClientConfiguration( google::bigtable::v2::GetClientConfigurationRequest const& request, OperationContext& operation_context) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { + pool_, operation_context, [&](BigtableStub& stub) { return stub.GetClientConfiguration(context, options, request, operation_context); }); @@ -458,9 +497,9 @@ BigtableRandomTwoLeastUsed::AsyncOpenTable( std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncOpenTable(cq, std::move(context), std::move(options), std::move(operation_context)); }); @@ -476,9 +515,9 @@ BigtableRandomTwoLeastUsed::AsyncOpenAuthorizedView( std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncOpenAuthorizedView(cq, std::move(context), std::move(options), std::move(operation_context)); @@ -495,9 +534,9 @@ BigtableRandomTwoLeastUsed::AsyncOpenMaterializedView( std::shared_ptr operation_context) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), options = std::move(options), - operation_context = - std::move(operation_context)](BigtableStub& stub) mutable { + pool_, operation_context, + [&, context = std::move(context), options = std::move(options), + operation_context](BigtableStub& stub) mutable { return stub.AsyncOpenMaterializedView(cq, std::move(context), std::move(options), std::move(operation_context)); diff --git a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc index 331ef53f54aad..63b7669cd3e4f 100644 --- a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc +++ b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc @@ -50,7 +50,8 @@ class BigtableRandomTwoLeastUsedTest : public ::testing::Test { pool_ = DynamicChannelPool::Create( instance_name, cq_, channels, refresh_state, - stub_factory_fn_.AsStdFunction(), sizing_policy); + stub_factory_fn_.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); } ~BigtableRandomTwoLeastUsedTest() override { diff --git a/google/cloud/bigtable/internal/bigtable_stub_factory.cc b/google/cloud/bigtable/internal/bigtable_stub_factory.cc index d2b3e45c22083..8bea357c31eb7 100644 --- a/google/cloud/bigtable/internal/bigtable_stub_factory.cc +++ b/google/cloud/bigtable/internal/bigtable_stub_factory.cc @@ -178,7 +178,10 @@ std::shared_ptr CreateBigtableStubRandomTwoLeastUsed( std::move(children), std::move(refresh_state), std::move(refreshing_channel_stub_factory), options.get< - bigtable::experimental::DynamicChannelPoolSizingPolicyOption>())); + bigtable::experimental::DynamicChannelPoolSizingPolicyOption>(), + bigtable::internal::IsDirectPath(options) + ? TransportType::kDirectPath + : TransportType::kCloudPath)); } std::shared_ptr CreateDecoratedStubs( diff --git a/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc b/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc index c5be858f15f93..4b2c45e33304e 100644 --- a/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc +++ b/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc @@ -1,4 +1,3 @@ -#include "google/cloud/bigtable/internal/operation_context.h" // Copyright 2022 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,9 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" #include "google/cloud/bigtable/internal/bigtable_stub_factory.h" +#include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" #include "google/cloud/bigtable/internal/dynamic_channel_pool.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/options.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/common_options.h" diff --git a/google/cloud/bigtable/internal/dynamic_channel_pool.h b/google/cloud/bigtable/internal/dynamic_channel_pool.h index a04745ce72bff..a6eb3a49431f1 100644 --- a/google/cloud/bigtable/internal/dynamic_channel_pool.h +++ b/google/cloud/bigtable/internal/dynamic_channel_pool.h @@ -66,11 +66,12 @@ class DynamicChannelPool std::vector>> initial_channels, std::shared_ptr refresh_state, StubFactoryFn stub_factory_fn, - bigtable::experimental::DynamicChannelPoolSizingPolicy sizing_policy) { + bigtable::experimental::DynamicChannelPoolSizingPolicy sizing_policy, + TransportType transport_type) { auto pool = std::shared_ptr(new DynamicChannelPool( std::move(instance_name), std::move(cq), std::move(initial_channels), std::move(refresh_state), std::move(stub_factory_fn), - std::move(sizing_policy))); + std::move(sizing_policy), transport_type)); return pool; } @@ -105,6 +106,8 @@ class DynamicChannelPool return sizing_policy_; } + TransportType transport_type() const { return transport_type_; } + // Calls CheckPoolChannelHealth before picking a channel. // // Pick two random channels from channels_ and return the channel with the @@ -173,14 +176,16 @@ class DynamicChannelPool std::vector>> initial_wrapped_channels, std::shared_ptr refresh_state, StubFactoryFn stub_factory_fn, - bigtable::experimental::DynamicChannelPoolSizingPolicy sizing_policy) + bigtable::experimental::DynamicChannelPoolSizingPolicy sizing_policy, + TransportType transport_type) : instance_name_(std::move(instance_name)), 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_(static_cast(channels_.size())) { + next_channel_id_(static_cast(channels_.size())), + transport_type_(transport_type) { std::scoped_lock lk(mu_); SetSizeDecreaseCooldownTimer(lk); } @@ -453,6 +458,7 @@ class DynamicChannelPool future> pool_size_decrease_cooldown_timer_; std::uint32_t next_channel_id_; + TransportType const transport_type_ = TransportType::kCloudPath; }; GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END diff --git a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc index e3f3be2b56bfe..c9e9d771ca2b2 100644 --- a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc +++ b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc @@ -180,7 +180,7 @@ TEST_F(DynamicChannelPoolTest, SelectLeastUsedFromTwoChannels) { sizing_policy.minimum_channel_pool_size = 2; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn, sizing_policy); + stub_factory_fn, sizing_policy, TransportType::kCloudPath); auto selected = pool->GetChannelRandomTwoLeastUsed(); EXPECT_THAT(selected.outstanding_rpcs, Eq(5)); grpc::ClientContext context; @@ -247,7 +247,8 @@ TEST_F(DynamicChannelPoolTest, OneInitialChannel) { sizing_policy.minimum_channel_pool_size = 1; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); EXPECT_THAT(pool->size(), Eq(1)); auto selected = pool->GetChannelRandomTwoLeastUsed(); @@ -310,7 +311,8 @@ TEST_F(DynamicChannelPoolTest, EmptyInitialPool) { sizing_policy.minimum_channel_pool_size = 0; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); EXPECT_THAT(*pool, ::testing::IsEmpty()); @@ -356,7 +358,8 @@ TEST_F(DynamicChannelPoolTest, ScheduleAddChannelsPoolUndersized) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); { @@ -416,7 +419,8 @@ TEST_F(DynamicChannelPoolTest, ScheduleAddChannelsPoolAtMax) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); EXPECT_CALL(*mock_cq_impl_, RunAsync).Times(1); @@ -463,7 +467,8 @@ TEST_F(DynamicChannelPoolTest, ScheduleAddChannelsPoolBelowMax) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); EXPECT_CALL(*mock_cq_impl_, RunAsync).Times(1); @@ -520,7 +525,8 @@ TEST_F(DynamicChannelPoolTest, AddChannels) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); std::vector new_channel_ids = {0, 1}; wrapper.set_num_pending_channels(new_channel_ids.size()); @@ -557,7 +563,8 @@ TEST_F(DynamicChannelPoolTest, ScheduleRemoveChannelsAlreadyPending) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); promise p; @@ -597,7 +604,8 @@ TEST_F(DynamicChannelPoolTest, ScheduleRemoveChannelsNotAlreadyPending) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); EXPECT_CALL(*mock_cq_impl_, MakeRelativeTimer) @@ -643,7 +651,8 @@ TEST_F(DynamicChannelPoolTest, RemoveChannelsLoneChannelDrained) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); std::vector>> draining_channels; @@ -684,7 +693,8 @@ TEST_F(DynamicChannelPoolTest, RemoveChannelsSomeChannelsDrained) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); std::vector>> draining_channels; @@ -779,7 +789,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOneBad) { sizing_policy.minimum_channel_pool_size = 2; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); @@ -862,7 +873,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOtherOneBad) { sizing_policy.minimum_channel_pool_size = 2; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); @@ -951,7 +963,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsThreeChannelsOneBad) { sizing_policy.minimum_channel_pool_size = 2; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); @@ -1050,7 +1063,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsAllChannelsBad) { sizing_policy.minimum_channel_pool_size = 2; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); @@ -1112,7 +1126,8 @@ TEST_F(DynamicChannelPoolTest, CheckChannelPoolHealthNeedsIncrease) { sizing_policy.maximum_channel_pool_size = 1; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); // ScheduleAddChannels will NOT be called as the pool has max channels. @@ -1126,7 +1141,8 @@ TEST_F(DynamicChannelPoolTest, CheckChannelPoolHealthNeedsIncrease) { sizing_policy.maximum_channel_pool_size = 10; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); // ScheduleAddChannels will be called. @@ -1192,7 +1208,8 @@ TEST_F(DynamicChannelPoolTest, CheckChannelPoolHealthNeedsDecrease) { sizing_policy.minimum_average_outstanding_rpcs_per_channel = 5; auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, - stub_factory_fn.AsStdFunction(), sizing_policy); + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); DynamicChannelPoolTestWrapper wrapper(pool); // ScheduleAddChannels will NOT be called. @@ -1207,6 +1224,41 @@ TEST_F(DynamicChannelPoolTest, CheckChannelPoolHealthNeedsDecrease) { EXPECT_THAT(pool->size(), Eq(3)); } +TEST_F(DynamicChannelPoolTest, TransportType) { + auto instance_name = + bigtable::InstanceResource(Project("my-project"), "my-instance") + .FullName(); + auto refresh_state = std::make_shared( + fake_cq_impl_, std::chrono::milliseconds(1), + std::chrono::milliseconds(10)); + std::vector>> channels; + DynamicChannelPoolSizingPolicy sizing_policy; + MockFunction>>( + std::uint32_t, std::string const&, StubManager::Priming)> + stub_factory_fn; + + EXPECT_CALL(*mock_cq_impl_, MakeRelativeTimer) + .WillRepeatedly([&](std::chrono::nanoseconds ns) { + EXPECT_THAT(ns.count(), + Eq(std::chrono::nanoseconds( + sizing_policy.pool_size_decrease_cooldown_interval) + .count())); + return make_ready_future(StatusOr(std::chrono::system_clock::now())); + }); + + auto pool_cloud = DynamicChannelPool::Create( + instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kCloudPath); + EXPECT_THAT(pool_cloud->transport_type(), Eq(TransportType::kCloudPath)); + + auto pool_direct = DynamicChannelPool::Create( + instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, + stub_factory_fn.AsStdFunction(), sizing_policy, + TransportType::kDirectPath); + EXPECT_THAT(pool_direct->transport_type(), Eq(TransportType::kDirectPath)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace bigtable_internal diff --git a/google/cloud/bigtable/internal/operation_context_factory.cc b/google/cloud/bigtable/internal/operation_context_factory.cc index fb5ebb397b1d5..ed6d394b4b896 100644 --- a/google/cloud/bigtable/internal/operation_context_factory.cc +++ b/google/cloud/bigtable/internal/operation_context_factory.cc @@ -15,6 +15,7 @@ #include "google/cloud/bigtable/internal/operation_context_factory.h" #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS +#include "google/cloud/bigtable/internal/client_schema_metrics.h" #include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/bigtable/internal/table_schema_metrics.h" #include "google/cloud/bigtable/options.h" @@ -29,6 +30,7 @@ #include #include #include +#include #endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS namespace google { @@ -151,12 +153,50 @@ std::shared_ptr SimpleOperationContextFactory::ExecuteQuery( #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS +namespace { +ClientResourceLabels MakeClientLabels( + ClientResourceLabels const& base, + TableResourceLabels const& resource_labels, std::string_view app_profile) { + ClientResourceLabels labels = base; + labels.project_id = resource_labels.project_id; + labels.instance = resource_labels.instance; + labels.app_profile = std::string(app_profile); + return labels; +} + +std::shared_ptr MakeOperationContext( + TableResourceLabels const& resource_labels, std::string_view rpc_name, + std::string_view streaming, std::string_view client_uid, + std::string_view app_profile, + ClientResourceLabels const& client_resource_labels, + std::vector> const& metrics, + std::shared_ptr const& clock) { + TableDataLabels data_labels = {std::string(rpc_name), + std::string(streaming), + "cpp.Bigtable/" + version_string(), + std::string(client_uid), + std::string(app_profile), + "" /*=status*/}; + ClientResourceLabels client_labels = + MakeClientLabels(client_resource_labels, resource_labels, app_profile); + + return std::make_shared( + CloneMetrics(resource_labels, data_labels, client_labels, metrics), + clock); +} + +} // namespace + MetricsOperationContextFactory::MetricsOperationContextFactory( std::string client_uid, std::shared_ptr conn, std::shared_ptr clock, Options options) // NOLINT(performance-unnecessary-value-param) - : client_uid_(std::move(client_uid)), clock_(std::move(clock)) { + : client_uid_(std::move(client_uid)), + clock_(std::move(clock)), + client_resource_labels_(MakeClientResourceLabels( + /*project_id=*/"", /*instance=*/"", /*app_profile=*/"", options, + client_uid_, otel::MakeResourceDetector()->Detect())) { InitializeProvider(std::move(conn), std::move(options)); } @@ -170,7 +210,10 @@ MetricsOperationContextFactory::MetricsOperationContextFactory( MetricsOperationContextFactory::MetricsOperationContextFactory( std::string client_uid, std::shared_ptr const& metric) - : client_uid_(std::move(client_uid)) { + : client_uid_(std::move(client_uid)), + client_resource_labels_(MakeClientResourceLabels( + /*project_id=*/"", /*instance=*/"", /*app_profile=*/"", Options{}, + client_uid_, otel::MakeResourceDetector()->Detect())) { absl::call_once(read_row_metrics_.once, [this, metric]() { read_row_metrics_.metrics.push_back(metric); }); @@ -211,32 +254,68 @@ void MetricsOperationContextFactory::InitializeProvider( auto constexpr kTableLabel = "table"; auto constexpr kClusterLabel = "cluster"; auto constexpr kZoneLabel = "zone"; + auto constexpr kAppProfileLabel = "app_profile"; + auto constexpr kClientNameLabel = "client_name"; + auto constexpr kClientUidAttribute = "client_uid"; + auto constexpr kUuidLabel = "uuid"; + auto constexpr kClientProjectLabel = "client_project"; + auto constexpr kLocationLabel = "location"; + auto constexpr kCloudPlatformLabel = "cloud_platform"; + auto constexpr kHostIdLabel = "host_id"; + auto constexpr kHostnameLabel = "hostname"; auto dynamic_resource_fn = [=](opentelemetry::sdk::metrics::PointDataAttributes const& pda) { + auto const& attributes = pda.attributes.GetAttributes(); + auto get_attr = [&](std::string const& key) { + auto it = attributes.find(key); + if (it == attributes.end() || + !opentelemetry::nostd::holds_alternative( + it->second)) { + return std::string{}; + } + return opentelemetry::nostd::get(it->second); + }; + google::api::MonitoredResource resource; resource.set_type(kResourceType); auto& labels = *resource.mutable_labels(); - auto const& attributes = pda.attributes.GetAttributes(); - labels[kProjectLabel] = - std::get(attributes.find(kProjectLabel)->second); - labels[kInstanceLabel] = - std::get(attributes.find(kInstanceLabel)->second); - labels[kTableLabel] = - std::get(attributes.find(kTableLabel)->second); - labels[kClusterLabel] = - std::get(attributes.find(kClusterLabel)->second); - labels[kZoneLabel] = - std::get(attributes.find(kZoneLabel)->second); + labels[kProjectLabel] = get_attr(kProjectLabel); + labels[kInstanceLabel] = get_attr(kInstanceLabel); + + if (attributes.find(kTableLabel) != attributes.end()) { + labels[kTableLabel] = get_attr(kTableLabel); + labels[kClusterLabel] = get_attr(kClusterLabel); + labels[kZoneLabel] = get_attr(kZoneLabel); + return std::make_pair(labels[kProjectLabel], resource); + } + + labels[kAppProfileLabel] = get_attr(kAppProfileLabel); + labels[kClientNameLabel] = get_attr(kClientNameLabel); + labels[kUuidLabel] = get_attr(kClientUidAttribute); + std::string client_project = get_attr(kClientProjectLabel); + if (!client_project.empty()) { + labels[kClientProjectLabel] = std::move(client_project); + } + labels[kLocationLabel] = get_attr(kLocationLabel); + labels[kCloudPlatformLabel] = get_attr(kCloudPlatformLabel); + labels[kHostIdLabel] = get_attr(kHostIdLabel); + std::string hostname = get_attr(kHostnameLabel); + if (!hostname.empty()) { + labels[kHostnameLabel] = std::move(hostname); + } return std::make_pair(labels[kProjectLabel], resource); }; - std::set s{kProjectLabel, kInstanceLabel, kTableLabel, - kClusterLabel, kZoneLabel}; - auto resource_filter_fn = [resource_labels = - std::move(s)](std::string const& key) { - return internal::Contains(resource_labels, key); - }; + std::set resource_labels{ + kProjectLabel, kInstanceLabel, kTableLabel, kClusterLabel, + kZoneLabel, kAppProfileLabel, kClientNameLabel, kClientUidAttribute, + kUuidLabel, kClientProjectLabel, kLocationLabel, kCloudPlatformLabel, + kHostIdLabel, kHostnameLabel}; + auto resource_filter_fn = + [resource_labels = std::move(resource_labels)](std::string const& key) { + return internal::Contains(resource_labels, key); + }; auto reader_options = opentelemetry::sdk::metrics::PeriodicExportingMetricReaderOptions{}; @@ -291,20 +370,14 @@ std::shared_ptr MetricsOperationContextFactory::ReadRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_row_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, read_row_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"true", client_uid_, app_profile, client_resource_labels_, + read_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::ReadRows( @@ -321,20 +394,14 @@ std::shared_ptr MetricsOperationContextFactory::ReadRows( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_rows_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, read_rows_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"true", client_uid_, app_profile, client_resource_labels_, + read_rows_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::MutateRow( @@ -350,20 +417,14 @@ std::shared_ptr MetricsOperationContextFactory::MutateRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(mutate_row_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, mutate_row_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"false", client_uid_, app_profile, client_resource_labels_, + mutate_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::MutateRows( @@ -379,20 +440,14 @@ std::shared_ptr MetricsOperationContextFactory::MutateRows( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(mutate_rows_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, mutate_rows_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"true", client_uid_, app_profile, client_resource_labels_, + mutate_rows_metrics_.metrics, clock_); } std::shared_ptr @@ -409,21 +464,14 @@ MetricsOperationContextFactory::CheckAndMutateRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(check_and_mutate_row_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, - check_and_mutate_row_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"false", client_uid_, app_profile, client_resource_labels_, + check_and_mutate_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::SampleRowKeys( @@ -439,21 +487,14 @@ std::shared_ptr MetricsOperationContextFactory::SampleRowKeys( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(sample_row_keys_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, - sample_row_keys_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"true", client_uid_, app_profile, client_resource_labels_, + sample_row_keys_metrics_.metrics, clock_); } std::shared_ptr @@ -470,21 +511,14 @@ MetricsOperationContextFactory::ReadModifyWriteRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_modify_write_row_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromTableName(table_name); - TableDataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, - read_modify_write_row_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromTableName(table_name), kRpc, + /*streaming=*/"false", client_uid_, app_profile, client_resource_labels_, + read_modify_write_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::PrepareQuery( @@ -498,21 +532,14 @@ std::shared_ptr MetricsOperationContextFactory::PrepareQuery( v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(prepare_query_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromInstanceName(instance_name); - TableDataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; - - return std::make_shared( - CloneMetrics(resource_labels, data_labels, - prepare_query_metrics_.metrics), - clock_); + return MakeOperationContext( + TableResourceLabelsFromInstanceName(instance_name), kRpc, + /*streaming=*/"false", client_uid_, app_profile, client_resource_labels_, + prepare_query_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::ExecuteQuery( @@ -529,35 +556,43 @@ std::shared_ptr MetricsOperationContextFactory::ExecuteQuery( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(execute_query_metrics_.metrics, v); }); - auto resource_labels = TableResourceLabelsFromInstanceName(instance_name); - TableDataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + return MakeOperationContext( + TableResourceLabelsFromInstanceName(instance_name), kRpc, + /*streaming=*/"true", client_uid_, app_profile, client_resource_labels_, + execute_query_metrics_.metrics, clock_); +} - return std::make_shared( - CloneMetrics(resource_labels, data_labels, - execute_query_metrics_.metrics), - clock_); +std::vector> CloneMetrics( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + std::vector> const& metrics) { + return CloneMetrics(resource_labels, data_labels, ClientResourceLabels{}, + metrics); } std::vector> CloneMetrics( TableResourceLabels const& resource_labels, TableDataLabels const& data_labels, + ClientResourceLabels const& client_resource_labels, std::vector> const& metrics) { std::vector> v; v.reserve(metrics.size()); for (auto const& m : metrics) { // We should never add a nullptr Metric to the list. if (m == nullptr) continue; - if (m->schema() == MetricSchema::kTable) { - v.push_back(static_cast(m.get())->clone( - resource_labels, data_labels)); + switch (m->schema()) { + case MetricSchema::kTable: + v.push_back(static_cast(m.get())->clone( + resource_labels, data_labels)); + break; + case MetricSchema::kClient: + v.push_back(static_cast(m.get())->clone( + client_resource_labels)); + break; } } return v; diff --git a/google/cloud/bigtable/internal/operation_context_factory.h b/google/cloud/bigtable/internal/operation_context_factory.h index 50ea4c932f42f..7124f71dba95f 100644 --- a/google/cloud/bigtable/internal/operation_context_factory.h +++ b/google/cloud/bigtable/internal/operation_context_factory.h @@ -18,6 +18,7 @@ #include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/version.h" #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS +#include "google/cloud/bigtable/internal/client_schema_metrics.h" #include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/bigtable/internal/table_schema_metrics.h" #include "google/cloud/monitoring/v3/metric_connection.h" @@ -142,6 +143,7 @@ class MetricsOperationContextFactory : public OperationContextFactory { std::string client_uid_; std::shared_ptr clock_; std::shared_ptr provider_; + ClientResourceLabels client_resource_labels_; // These vectors are initialized exactly once and the initialization is // delayed until the first time the corresponding method is called. @@ -165,6 +167,12 @@ std::vector> CloneMetrics( TableDataLabels const& data_labels, std::vector> const& metrics); +std::vector> CloneMetrics( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + ClientResourceLabels const& client_resource_labels, + std::vector> const& metrics); + #endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END diff --git a/google/cloud/bigtable/internal/operation_context_factory_test.cc b/google/cloud/bigtable/internal/operation_context_factory_test.cc index d10a9d69596f4..300bff1661045 100644 --- a/google/cloud/bigtable/internal/operation_context_factory_test.cc +++ b/google/cloud/bigtable/internal/operation_context_factory_test.cc @@ -15,9 +15,13 @@ #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS #include "google/cloud/bigtable/internal/operation_context_factory.h" +#include "google/cloud/bigtable/internal/client_schema_metrics.h" #include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/bigtable/internal/table_schema_metrics.h" +#include "google/cloud/bigtable/options.h" +#include "google/cloud/monitoring/v3/metric_connection.h" #include +#include namespace google { namespace cloud { @@ -27,6 +31,7 @@ namespace { using ::testing::Eq; using ::testing::IsEmpty; +using ::testing::NotNull; using ::testing::SizeIs; class MockMetric : public TableSchemaMetric { @@ -61,7 +66,8 @@ TEST(MetricsOperationContextFactoryTest, ReadRow) { }); MetricsOperationContextFactory factory("my-client-uid", mock_metric); - auto operation_context = factory.ReadRow(table_full_name, app_profile); + std::shared_ptr operation_context = + factory.ReadRow(table_full_name, app_profile); } TEST(MetricsOperationContextFactoryTest, ReadRows) { @@ -79,7 +85,8 @@ TEST(MetricsOperationContextFactoryTest, ReadRows) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = factory.ReadRows(table_full_name, app_profile); + std::shared_ptr operation_context = + factory.ReadRows(table_full_name, app_profile); } TEST(MetricsOperationContextFactoryTest, MutateRow) { @@ -97,7 +104,8 @@ TEST(MetricsOperationContextFactoryTest, MutateRow) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = factory.MutateRow(table_full_name, app_profile); + std::shared_ptr operation_context = + factory.MutateRow(table_full_name, app_profile); } TEST(MetricsOperationContextFactoryTest, MutateRows) { @@ -115,7 +123,8 @@ TEST(MetricsOperationContextFactoryTest, MutateRows) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = factory.MutateRows(table_full_name, app_profile); + std::shared_ptr operation_context = + factory.MutateRows(table_full_name, app_profile); } TEST(MetricsOperationContextFactoryTest, CheckAndMutateRow) { @@ -133,7 +142,7 @@ TEST(MetricsOperationContextFactoryTest, CheckAndMutateRow) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = + std::shared_ptr operation_context = factory.CheckAndMutateRow(table_full_name, app_profile); } @@ -152,7 +161,8 @@ TEST(MetricsOperationContextFactoryTest, SampleRowKeys) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = factory.SampleRowKeys(table_full_name, app_profile); + std::shared_ptr operation_context = + factory.SampleRowKeys(table_full_name, app_profile); } TEST(MetricsOperationContextFactoryTest, ReadModifyWriteRow) { @@ -170,7 +180,7 @@ TEST(MetricsOperationContextFactoryTest, ReadModifyWriteRow) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = + std::shared_ptr operation_context = factory.ReadModifyWriteRow(table_full_name, app_profile); } @@ -188,7 +198,7 @@ TEST(MetricsOperationContextFactoryTest, PrepareQuery) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = + std::shared_ptr operation_context = factory.PrepareQuery(instance_full_name, app_profile); } @@ -206,10 +216,57 @@ TEST(MetricsOperationContextFactoryTest, ExecuteQuery) { }); MetricsOperationContextFactory factory({}, mock_metric); - auto operation_context = + std::shared_ptr operation_context = factory.ExecuteQuery(instance_full_name, app_profile); } +TEST(MetricsOperationContextFactoryTest, IncludesOutstandingRpcs) { + std::string app_profile = "my-app-profile"; + std::string table_full_name = + "projects/my-project/instances/my-instance/tables/my-table"; + auto options = + Options{}.set(std::chrono::seconds(60)); + MetricsOperationContextFactory factory( + "test-uid", + std::shared_ptr(nullptr), + std::move(options)); + std::shared_ptr operation_context = + factory.ReadRow(table_full_name, app_profile); + EXPECT_THAT(operation_context, NotNull()); + operation_context->StubSelection( + StubSelectionParams{10, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + TransportType::kDirectPath, RpcType::kUnary}); +} + +class MockClientMetric : public ClientSchemaMetric { + public: + MOCK_METHOD(std::unique_ptr, clone, + (ClientResourceLabels const&), (const, override)); +}; + +TEST(MetricsOperationContextFactoryTest, ClientResourceLabelsPopulated) { + std::string app_profile = "my-app-profile"; + std::string table_full_name = + "projects/my-project/instances/my-instance/tables/my-table"; + + auto mock_metric = std::make_shared(); + EXPECT_CALL(*mock_metric, clone(::testing::A())) + .WillOnce([&](ClientResourceLabels const& client_labels) { + EXPECT_THAT(client_labels.project_id, Eq("my-project")); + EXPECT_THAT(client_labels.instance, Eq("my-instance")); + EXPECT_THAT(client_labels.app_profile, Eq("my-app-profile")); + EXPECT_THAT(client_labels.client_name, + Eq("cpp.Bigtable/" + version_string())); + EXPECT_THAT(client_labels.client_uid, Eq("my-client-uid")); + return std::make_unique(); + }); + + MetricsOperationContextFactory factory("my-client-uid", mock_metric); + std::shared_ptr operation_context = + factory.ReadRow(table_full_name, app_profile); + EXPECT_THAT(operation_context, NotNull()); +} + class FakeTableMetric : public TableSchemaMetric { public: std::unique_ptr clone( @@ -218,17 +275,157 @@ class FakeTableMetric : public TableSchemaMetric { } }; +class FakeClientMetric : public ClientSchemaMetric { + public: + std::unique_ptr clone( + ClientResourceLabels const&) const override { + return std::make_unique(*this); + } +}; + TEST(MetricsOperationContextFactoryTest, CloneMetrics) { auto table_metric = std::make_shared(); + auto client_metric = std::make_shared(); + + TableResourceLabels resource_labels{"project", "instance", "table", "cluster", + "zone"}; + TableDataLabels data_labels{"method", "streaming", "client", + "uid", "profile", "status"}; + + std::vector> metrics = {table_metric, + client_metric}; + std::vector> cloned = + CloneMetrics(resource_labels, data_labels, metrics); + EXPECT_THAT(cloned, SizeIs(2)); +} + +TEST(MetricsOperationContextFactoryTest, CloneMetricsWithClientResourceLabels) { + auto table_metric = std::make_shared(); + auto client_metric = std::make_shared(); TableResourceLabels resource_labels{"project", "instance", "table", "cluster", "zone"}; TableDataLabels data_labels{"method", "streaming", "client", "uid", "profile", "status"}; + ClientResourceLabels client_labels{ + "project", "instance", "profile", "client", "uid", + "client-proj", "location", "gcp", "host", "hostname"}; + + std::vector> metrics = {table_metric, + client_metric}; + std::vector> cloned = + CloneMetrics(resource_labels, data_labels, client_labels, metrics); + EXPECT_THAT(cloned, SizeIs(2)); +} + +class MockMetricServiceConnection + : public monitoring_v3::MetricServiceConnection { + public: + ~MockMetricServiceConnection() override = default; + MOCK_METHOD(Status, CreateServiceTimeSeries, + (google::monitoring::v3::CreateTimeSeriesRequest const&), + (override)); + MOCK_METHOD(Status, CreateTimeSeries, + (google::monitoring::v3::CreateTimeSeriesRequest const&), + (override)); +}; + +TEST(MetricsOperationContextFactoryTest, + InitializeProviderExportsTableAndClientMetrics) { + auto mock_conn = std::make_shared(); + + EXPECT_CALL(*mock_conn, CreateServiceTimeSeries) + .WillRepeatedly( + [](google::monitoring::v3::CreateTimeSeriesRequest const& request) { + EXPECT_THAT(request.name(), Eq("projects/my-project")); + for (auto const& ts : request.time_series()) { + EXPECT_THAT(ts.resource().type(), Eq("bigtable_client_raw")); + if (ts.resource().labels().find("table") != + ts.resource().labels().end()) { + EXPECT_THAT(ts.resource().labels().at("project_id"), + Eq("my-project")); + EXPECT_THAT(ts.resource().labels().at("instance"), + Eq("my-instance")); + EXPECT_THAT(ts.resource().labels().at("table"), Eq("my-table")); + } else { + EXPECT_THAT(ts.resource().labels().at("project_id"), + Eq("my-project")); + EXPECT_THAT(ts.resource().labels().at("instance"), + Eq("my-instance")); + EXPECT_THAT(ts.resource().labels().at("app_profile"), + Eq("my-app-profile")); + EXPECT_THAT(ts.resource().labels().at("uuid"), Eq("test-uid")); + EXPECT_THAT(ts.resource().labels().at("client_name"), + Eq("cpp.Bigtable/" + version_string())); + } + // Verify resource filtering removed resource labels from metric + // labels: + EXPECT_TRUE(ts.metric().labels().find("project_id") == + ts.metric().labels().end()); + EXPECT_TRUE(ts.metric().labels().find("instance") == + ts.metric().labels().end()); + EXPECT_TRUE(ts.metric().labels().find("table") == + ts.metric().labels().end()); + EXPECT_TRUE(ts.metric().labels().find("app_profile") == + ts.metric().labels().end()); + EXPECT_TRUE(ts.metric().labels().find("uuid") == + ts.metric().labels().end()); + } + return Status(); + }); + + auto options = + Options{}.set(std::chrono::seconds(60)); + { + MetricsOperationContextFactory factory("test-uid", mock_conn, options); + std::shared_ptr op = factory.ReadRow( + "projects/my-project/instances/my-instance/tables/my-table", + "my-app-profile"); + ASSERT_THAT(op, NotNull()); + op->OnDone(Status()); + op->StubSelection( + StubSelectionParams{10, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + TransportType::kDirectPath, RpcType::kUnary}); + } +} + +TEST(MetricsOperationContextFactoryTest, InitializeProviderInstanceLevelRpc) { + auto mock_conn = std::make_shared(); + + EXPECT_CALL(*mock_conn, CreateServiceTimeSeries) + .WillRepeatedly( + [](google::monitoring::v3::CreateTimeSeriesRequest const& request) { + EXPECT_THAT(request.name(), Eq("projects/my-project")); + for (auto const& ts : request.time_series()) { + EXPECT_THAT(ts.resource().type(), Eq("bigtable_client_raw")); + EXPECT_THAT(ts.resource().labels().at("project_id"), + Eq("my-project")); + EXPECT_THAT(ts.resource().labels().at("instance"), + Eq("my-instance")); + } + return Status(); + }); + + auto options = + Options{}.set(std::chrono::seconds(60)); + { + MetricsOperationContextFactory factory("test-uid", mock_conn, options); + std::shared_ptr op = factory.ExecuteQuery( + "projects/my-project/instances/my-instance", "my-app-profile"); + ASSERT_THAT(op, NotNull()); + op->OnDone(Status()); + } +} + +TEST(MetricsOperationContextFactoryTest, InitializeProviderWithoutConnection) { + auto options = + Options{}.set(std::chrono::seconds(60)); + MetricsOperationContextFactory factory("test-uid", nullptr, options); - std::vector> metrics = {table_metric}; - auto cloned = CloneMetrics(resource_labels, data_labels, metrics); - EXPECT_THAT(cloned, SizeIs(1)); + std::shared_ptr operation_context = factory.ReadRow( + "projects/my-project/instances/my-instance/tables/my-table", + "my-app-profile"); + EXPECT_THAT(operation_context, NotNull()); } } // namespace