From dbd838662f054678499e42715c9e17ecb4d45601 Mon Sep 17 00:00:00 2001 From: Zijie Li Date: Mon, 11 May 2026 15:02:35 -0400 Subject: [PATCH 1/2] feat(bigtable): add AttemptLatency2 metric and populate peer info labels This introduces the `AttemptLatency2` metric for DirectPath to record attempt latencies with the fields extracted from the decoded `PeerInfo` trailing metadata, populating `peer_info_labels_` and forwarding them to `IntoLabelMap`. Also added `AttemptLatency2Test` to test the newly populated peer info labels. Refactored `SetClusterZone` to use the new helper function `CreateServerMetadata`. --- google/cloud/bigtable/internal/metrics.cc | 103 ++++- google/cloud/bigtable/internal/metrics.h | 42 +- .../cloud/bigtable/internal/metrics_test.cc | 396 +++++++++++++++++- .../internal/operation_context_factory.cc | 9 + 4 files changed, 526 insertions(+), 24 deletions(-) diff --git a/google/cloud/bigtable/internal/metrics.cc b/google/cloud/bigtable/internal/metrics.cc index e2b09ef9ac2f5..af64a56cc7ffb 100644 --- a/google/cloud/bigtable/internal/metrics.cc +++ b/google/cloud/bigtable/internal/metrics.cc @@ -17,6 +17,7 @@ #include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/bigtable/version.h" #include "absl/strings/charconv.h" +#include "absl/strings/escaping.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" @@ -38,13 +39,23 @@ auto constexpr kMeterInstrumentationScopeVersion = "v1"; // to the map should be more performant than performing a set_difference every // time. LabelMap IntoLabelMap(ResourceLabels const& r, DataLabels const& d, - std::set const& filtered_data_labels) { + std::set const& filtered_data_labels, + std::optional const& peer_info_labels) { LabelMap labels = { {"project_id", r.project_id}, {"instance", r.instance}, {"table", r.table}, {"cluster", r.cluster.empty() ? "" : r.cluster}, {"zone", r.zone.empty() ? "global" : r.zone}}; + + if (peer_info_labels) { + labels.insert({ + {"transport_type", peer_info_labels->transport_type}, + {"transport_region", peer_info_labels->transport_region}, + {"transport_subzone", peer_info_labels->transport_subzone}, + }); + } + std::map data = {{ {"method", d.method}, {"streaming", d.streaming}, @@ -74,6 +85,7 @@ LabelMap IntoLabelMap(ResourceLabels const& r, DataLabels const& d, std::set_difference(data.begin(), data.end(), filtered_data_labels.begin(), filtered_data_labels.end(), std::inserter(labels, labels.begin()), Compare()); + return labels; } @@ -103,6 +115,47 @@ GetResponseParamsFromTrailingMetadata( return absl::nullopt; } +absl::optional GetPeerInfoFromTrailingMetadata( + grpc::ClientContext const& client_context) { + auto metadata = client_context.GetServerTrailingMetadata(); + // Base64 encoded peer info header key defined by the server. + auto iter = metadata.find("bigtable-peer-info"); + if (iter == metadata.end()) return absl::nullopt; + std::string decoded; + if (!absl::Base64Unescape( + absl::string_view{iter->second.data(), iter->second.size()}, + &decoded)) { + return absl::nullopt; + } + google::bigtable::v2::PeerInfo p; + if (p.ParseFromString(decoded)) return p; + return absl::nullopt; +} + +std::string TransportTypeToString( + google::bigtable::v2::PeerInfo::TransportType type) { + switch (type) { + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN: + return "transport_type_unknown"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_EXTERNAL: + return "transport_type_external"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_CLOUD_PATH: + return "transport_type_cloud_path"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS: + return "transport_type_direct_access"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_UNKNOWN: + return "transport_type_session_unknown"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_EXTERNAL: + return "transport_type_session_external"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_CLOUD_PATH: + return "transport_type_session_cloud_path"; + case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_DIRECT_ACCESS: + return "transport_type_session_direct_access"; + default: + return "transport_type_unknown"; + } +} + absl::optional GetServerLatencyFromInitialMetadata( grpc::ClientContext const& client_context) { auto const& initial_metadata = client_context.GetServerInitialMetadata(); @@ -225,6 +278,54 @@ std::unique_ptr AttemptLatency::clone(ResourceLabels resource_labels, return m; } +AttemptLatency2::AttemptLatency2( + std::string const& instrumentation_scope, + opentelemetry::nostd::shared_ptr< + opentelemetry::metrics::MeterProvider> const& provider) + : attempt_latencies2_(provider + ->GetMeter(instrumentation_scope, + kMeterInstrumentationScopeVersion) + ->CreateDoubleHistogram("attempt_latencies2")) {} + +void AttemptLatency2::PreCall(opentelemetry::context::Context const&, + PreCallParams const& p) { + attempt_start_ = std::move(p.attempt_start); +} + +void AttemptLatency2::PostCall(opentelemetry::context::Context const& context, + grpc::ClientContext const& client_context, + PostCallParams const& p) { + auto response_params = GetResponseParamsFromTrailingMetadata(client_context); + if (response_params) { + resource_labels_.cluster = response_params->cluster_id(); + resource_labels_.zone = response_params->zone_id(); + } + auto peer_info = GetPeerInfoFromTrailingMetadata(client_context); + peer_info_labels_.transport_type = TransportTypeToString( + peer_info ? peer_info->transport_type() + : google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN); + if (peer_info) { + peer_info_labels_.transport_region = + peer_info->application_frontend_region(); + peer_info_labels_.transport_subzone = + peer_info->application_frontend_subzone(); + } + + data_labels_.status = StatusCodeToString(p.attempt_status.code()); + auto attempt_elapsed = std::chrono::duration_cast( + p.attempt_end - attempt_start_); + auto m = IntoLabelMap(resource_labels_, data_labels_, {}, peer_info_labels_); + attempt_latencies2_->Record(attempt_elapsed.count(), std::move(m), context); +} + +std::unique_ptr AttemptLatency2::clone(ResourceLabels resource_labels, + DataLabels data_labels) const { + auto m = std::make_unique(*this); + m->resource_labels_ = std::move(resource_labels); + m->data_labels_ = std::move(data_labels); + return m; +} + RetryCount::RetryCount( std::string const& instrumentation_scope, opentelemetry::nostd::shared_ptr< diff --git a/google/cloud/bigtable/internal/metrics.h b/google/cloud/bigtable/internal/metrics.h index 743e58ac50b54..34f9bd70ebbaa 100644 --- a/google/cloud/bigtable/internal/metrics.h +++ b/google/cloud/bigtable/internal/metrics.h @@ -20,6 +20,7 @@ #include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/version.h" #include "google/cloud/status.h" +#include "google/bigtable/v2/peer_info.pb.h" #include "google/bigtable/v2/response_params.pb.h" #include #include @@ -52,9 +53,19 @@ struct DataLabels { std::string status; }; +// Labels populated from the peer info metadata. +struct PeerInfoLabels { + std::string transport_type; + std::string transport_region; + std::string transport_subzone; +}; + using LabelMap = std::unordered_map; -LabelMap IntoLabelMap(ResourceLabels const& r, DataLabels const& d, - std::set const& filtered_data_labels = {}); +// `peer_info_labels` is optional because only AttemptLatency2 populates it. +LabelMap IntoLabelMap( + ResourceLabels const& r, DataLabels const& d, + std::set const& filtered_data_labels = {}, + std::optional const& peer_info_labels = std::nullopt); bool HasServerTiming(grpc::ClientContext const& client_context); bool IsConnectivityError(google::cloud::Status const& status, @@ -62,6 +73,10 @@ bool IsConnectivityError(google::cloud::Status const& status, absl::optional GetResponseParamsFromTrailingMetadata( grpc::ClientContext const& client_context); +absl::optional GetPeerInfoFromTrailingMetadata( + grpc::ClientContext const& client_context); +std::string TransportTypeToString( + google::bigtable::v2::PeerInfo::TransportType type); absl::optional GetServerLatencyFromInitialMetadata( grpc::ClientContext const& client_context); @@ -154,6 +169,29 @@ class AttemptLatency : public Metric { OperationContext::Clock::time_point attempt_start_; }; +// Similar to AttemptLatency and also populates the peer info. +class AttemptLatency2 : public Metric { + public: + AttemptLatency2(std::string const& instrumentation_scope, + opentelemetry::nostd::shared_ptr< + opentelemetry::metrics::MeterProvider> const& provider); + void PreCall(opentelemetry::context::Context const&, + PreCallParams const& p) override; + void PostCall(opentelemetry::context::Context const& context, + grpc::ClientContext const& client_context, + PostCallParams const& p) override; + std::unique_ptr clone(ResourceLabels resource_labels, + DataLabels data_labels) const override; + + private: + ResourceLabels resource_labels_; + DataLabels data_labels_; + PeerInfoLabels peer_info_labels_; + opentelemetry::nostd::shared_ptr> + attempt_latencies2_; + OperationContext::Clock::time_point attempt_start_; +}; + class RetryCount : public Metric { public: RetryCount(std::string const& instrumentation_scope, diff --git a/google/cloud/bigtable/internal/metrics_test.cc b/google/cloud/bigtable/internal/metrics_test.cc index 41a8bb17b5b58..2e2530d6c81a4 100644 --- a/google/cloud/bigtable/internal/metrics_test.cc +++ b/google/cloud/bigtable/internal/metrics_test.cc @@ -21,6 +21,8 @@ #include "google/cloud/bigtable/version.h" #include "google/cloud/testing_util/fake_clock.h" #include "google/cloud/testing_util/validate_metadata.h" +#include "absl/strings/escaping.h" +#include "google/bigtable/v2/peer_info.pb.h" #include #include @@ -285,6 +287,60 @@ TEST(GetResponseParamsFromMetadata, EmptyHeader) { EXPECT_FALSE(result); } +TEST(GetPeerInfoFromMetadata, NonEmptyHeader) { + google::bigtable::v2::PeerInfo expected_peer_info; + expected_peer_info.set_transport_type( + google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); + expected_peer_info.set_application_frontend_region("my-region"); + expected_peer_info.set_application_frontend_subzone("my-subzone"); + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.trailers.emplace( + "bigtable-peer-info", + absl::Base64Escape(expected_peer_info.SerializeAsString())); + SetServerMetadata(client_context, server_metadata); + + auto result = GetPeerInfoFromTrailingMetadata(client_context); + ASSERT_TRUE(result); + EXPECT_THAT(result->transport_type(), + Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS)); + EXPECT_THAT(result->application_frontend_region(), Eq("my-region")); + EXPECT_THAT(result->application_frontend_subzone(), Eq("my-subzone")); +} + +TEST(GetPeerInfoFromMetadata, EmptyHeader) { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromTrailingMetadata(client_context), + Eq(absl::nullopt)); +} + +TEST(GetPeerInfoFromMetadata, EmptyString) { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.trailers.emplace("bigtable-peer-info", ""); + SetServerMetadata(client_context, server_metadata); + + auto result = GetPeerInfoFromTrailingMetadata(client_context); + ASSERT_TRUE(result); + EXPECT_THAT(result->transport_type(), + Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN)); + EXPECT_THAT(result->application_frontend_region(), Eq("")); + EXPECT_THAT(result->application_frontend_subzone(), Eq("")); +} + +TEST(GetPeerInfoFromMetadata, InvalidBase64) { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.trailers.emplace("bigtable-peer-info", "invalid-base64!"); + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromTrailingMetadata(client_context), + Eq(absl::nullopt)); +} + std::unordered_map MakeAttributesMap( opentelemetry::common::KeyValueIterable const& attributes) { std::unordered_map m; @@ -301,15 +357,44 @@ std::unordered_map MakeAttributesMap( return m; } -void SetClusterZone(grpc::ClientContext& client_context, - std::string const& cluster, std::string const& zone) { - google::bigtable::v2::ResponseParams expected_response_params; - expected_response_params.set_cluster_id(cluster); - expected_response_params.set_zone_id(zone); +struct CreateServerMetadataOptions { + bool include_peer_info = true; + bool include_response_params = true; +}; + +// Helper function to create server metadata by populating the trailers based on +// the options. +void CreateServerMetadata(grpc::ClientContext& client_context, + CreateServerMetadataOptions const& options = + CreateServerMetadataOptions{}) { RpcMetadata server_metadata; - server_metadata.trailers.emplace( - "x-goog-ext-425905942-bin", expected_response_params.SerializeAsString()); - SetServerMetadata(client_context, server_metadata); + if (options.include_peer_info) { + google::bigtable::v2::PeerInfo peer_info; + peer_info.set_transport_type( + google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); + peer_info.set_application_frontend_region("my-region"); + peer_info.set_application_frontend_subzone("my-subzone"); + server_metadata.trailers.emplace( + "bigtable-peer-info", + absl::Base64Escape(peer_info.SerializeAsString())); + } + if (options.include_response_params) { + google::bigtable::v2::ResponseParams expected_response_params; + expected_response_params.set_cluster_id("my-cluster"); + expected_response_params.set_zone_id("my-zone"); + server_metadata.trailers.emplace( + "x-goog-ext-425905942-bin", + expected_response_params.SerializeAsString()); + } + + if (!server_metadata.trailers.empty()) { + SetServerMetadata(client_context, server_metadata); + } +} + +void SetClusterZone(grpc::ClientContext& client_context) { + CreateServerMetadata(client_context, + CreateServerMetadataOptions{false, true}); } TEST(OperationLatencyTest, FirstAttemptSuccess) { @@ -371,7 +456,7 @@ TEST(OperationLatencyTest, FirstAttemptSuccess) { auto clone = operation_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -444,7 +529,7 @@ TEST(OperationLatencyTest, ThirdAttemptSuccess) { auto clone = operation_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -604,7 +689,7 @@ TEST(AttemptLatencyTest, NoRetry) { auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -707,7 +792,7 @@ TEST(AttemptLatencyTest, ThreeAttempts) { auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -803,6 +888,275 @@ TEST(AttemptLatencyTest, UsesDefaultClusterAndZone) { clone->OnDone(otel_context, {clock->Now(), Status{StatusCode::kOk, "ok"}}); } +TEST(AttemptLatency2Test, NoRetry) { + auto mock_histogram = std::make_unique>(); + EXPECT_CALL( + *mock_histogram, + Record(A(), A(), + A())) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attributes, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(1.234)); + EXPECT_THAT( + MakeAttributesMap(attributes), + UnorderedElementsAre( + Pair("project_id", "my-project-id"), + Pair("instance", "my-instance"), Pair("cluster", "my-cluster"), + Pair("table", "my-table"), Pair("zone", "my-zone"), + Pair("method", "my-method"), Pair("streaming", "my-streaming"), + Pair("status", "OK"), Pair("client_name", "my-client-name"), + Pair("client_uid", "my-client-uid"), + Pair("app_profile", "my-app-profile"), + Pair("transport_type", "transport_type_direct_access"), + Pair("transport_region", "my-region"), + Pair("transport_subzone", "my-subzone"))); + }); + + opentelemetry::nostd::shared_ptr mock_meter = + std::make_shared(); + EXPECT_CALL(*mock_meter, CreateDoubleHistogram) + .WillOnce([mock = std::move(mock_histogram)]( + opentelemetry::nostd::string_view name, + opentelemetry::nostd::string_view, + opentelemetry::nostd::string_view) mutable { + EXPECT_THAT(name, Eq("attempt_latencies2")); + return std::move(mock); + }); + + opentelemetry::nostd::shared_ptr mock_provider = + std::make_shared(); + EXPECT_CALL(*mock_provider, GetMeter) +#if OPENTELEMETRY_ABI_VERSION_NO >= 2 + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view, + opentelemetry::common::KeyValueIterable const*) mutable { +#else + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view) mutable { +#endif + EXPECT_THAT(scope, Eq("my-instrument-scope")); + EXPECT_THAT(scope_version, Eq("v1")); + return mock_meter; + }); + + AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); + ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", + ""}; + DataLabels data_labels{"my-method", "my-streaming", "my-client-name", + "my-client-uid", "my-app-profile", ""}; + auto clone = attempt_latency.clone(resource_labels, data_labels); + + grpc::ClientContext client_context; + CreateServerMetadata(client_context); + + auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); + auto clock = std::make_shared(); + + clock->SetTime(std::chrono::steady_clock::now()); + clone->PreCall(otel_context, {clock->Now(), true}); + clock->AdvanceTime(std::chrono::microseconds(1234)); + clone->PostCall(otel_context, client_context, + {clock->Now(), Status{StatusCode::kOk, "ok"}}); +} + +TEST(AttemptLatency2Test, ThreeAttempts) { + auto mock_histogram = std::make_unique>(); + EXPECT_CALL( + *mock_histogram, + Record(A(), A(), + A())) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attributes, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(1.0)); + EXPECT_THAT( + MakeAttributesMap(attributes), + UnorderedElementsAre( + Pair("project_id", "my-project-id"), + Pair("instance", "my-instance"), Pair("cluster", "my-cluster"), + Pair("table", "my-table"), Pair("zone", "my-zone"), + Pair("method", "my-method"), Pair("streaming", "my-streaming"), + Pair("status", "OK"), Pair("client_name", "my-client-name"), + Pair("client_uid", "my-client-uid"), + Pair("app_profile", "my-app-profile"), + Pair("transport_type", "transport_type_direct_access"), + Pair("transport_region", "my-region"), + Pair("transport_subzone", "my-subzone"))); + }) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attributes, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(2.0)); + EXPECT_THAT( + MakeAttributesMap(attributes), + UnorderedElementsAre( + Pair("project_id", "my-project-id"), + Pair("instance", "my-instance"), Pair("cluster", "my-cluster"), + Pair("table", "my-table"), Pair("zone", "my-zone"), + Pair("method", "my-method"), Pair("streaming", "my-streaming"), + Pair("status", "OK"), Pair("client_name", "my-client-name"), + Pair("client_uid", "my-client-uid"), + Pair("app_profile", "my-app-profile"), + Pair("transport_type", "transport_type_direct_access"), + Pair("transport_region", "my-region"), + Pair("transport_subzone", "my-subzone"))); + }) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attributes, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(3.0)); + EXPECT_THAT( + MakeAttributesMap(attributes), + UnorderedElementsAre( + Pair("project_id", "my-project-id"), + Pair("instance", "my-instance"), Pair("cluster", "my-cluster"), + Pair("table", "my-table"), Pair("zone", "my-zone"), + Pair("method", "my-method"), Pair("streaming", "my-streaming"), + Pair("status", "OK"), Pair("client_name", "my-client-name"), + Pair("client_uid", "my-client-uid"), + Pair("app_profile", "my-app-profile"), + Pair("transport_type", "transport_type_direct_access"), + Pair("transport_region", "my-region"), + Pair("transport_subzone", "my-subzone"))); + }); + + opentelemetry::nostd::shared_ptr mock_meter = + std::make_shared(); + EXPECT_CALL(*mock_meter, CreateDoubleHistogram) + .WillOnce([mock = std::move(mock_histogram)]( + opentelemetry::nostd::string_view name, + opentelemetry::nostd::string_view, + opentelemetry::nostd::string_view) mutable { + EXPECT_THAT(name, Eq("attempt_latencies2")); + return std::move(mock); + }); + + opentelemetry::nostd::shared_ptr mock_provider = + std::make_shared(); + EXPECT_CALL(*mock_provider, GetMeter) +#if OPENTELEMETRY_ABI_VERSION_NO >= 2 + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view, + opentelemetry::common::KeyValueIterable const*) mutable { +#else + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view) mutable { +#endif + EXPECT_THAT(scope, Eq("my-instrument-scope")); + EXPECT_THAT(scope_version, Eq("v1")); + return mock_meter; + }); + + AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); + ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", + ""}; + DataLabels data_labels{"my-method", "my-streaming", "my-client-name", + "my-client-uid", "my-app-profile", ""}; + auto clone = attempt_latency.clone(resource_labels, data_labels); + + grpc::ClientContext client_context; + CreateServerMetadata(client_context); + + auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); + auto clock = std::make_shared(); + + clock->SetTime(std::chrono::steady_clock::now()); + clone->PreCall(otel_context, {clock->Now(), true}); + clock->AdvanceTime(std::chrono::milliseconds(1)); + clone->PostCall(otel_context, client_context, + {clock->Now(), Status{StatusCode::kOk, "ok"}}); + clock->AdvanceTime(std::chrono::milliseconds(100)); + clone->PreCall(otel_context, {clock->Now(), false}); + clock->AdvanceTime(std::chrono::milliseconds(2)); + clone->PostCall(otel_context, client_context, + {clock->Now(), Status{StatusCode::kOk, "ok"}}); + clone->PreCall(otel_context, {clock->Now(), false}); + clock->AdvanceTime(std::chrono::milliseconds(3)); + clone->PostCall(otel_context, client_context, + {clock->Now(), Status{StatusCode::kOk, "ok"}}); +} + +TEST(AttemptLatency2Test, UsesDefaultClusterAndZone) { + auto mock_histogram = std::make_unique>(); + EXPECT_CALL( + *mock_histogram, + Record(A(), A(), + A())) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attributes, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(1.234)); + EXPECT_THAT( + MakeAttributesMap(attributes), + UnorderedElementsAre( + Pair("project_id", "my-project-id"), + Pair("instance", "my-instance"), + Pair("cluster", ""), Pair("table", "my-table"), + Pair("zone", "global"), Pair("method", "my-method"), + Pair("streaming", "my-streaming"), Pair("status", "OK"), + Pair("client_name", "my-client-name"), + Pair("client_uid", "my-client-uid"), + Pair("app_profile", "my-app-profile"), + Pair("transport_type", "transport_type_direct_access"), + Pair("transport_region", "my-region"), + Pair("transport_subzone", "my-subzone"))); + }); + + opentelemetry::nostd::shared_ptr mock_meter = + std::make_shared(); + EXPECT_CALL(*mock_meter, CreateDoubleHistogram) + .WillOnce([mock = std::move(mock_histogram)]( + opentelemetry::nostd::string_view name, + opentelemetry::nostd::string_view, + opentelemetry::nostd::string_view) mutable { + EXPECT_THAT(name, Eq("attempt_latencies2")); + return std::move(mock); + }); + + opentelemetry::nostd::shared_ptr mock_provider = + std::make_shared(); + EXPECT_CALL(*mock_provider, GetMeter) +#if OPENTELEMETRY_ABI_VERSION_NO >= 2 + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view, + opentelemetry::common::KeyValueIterable const*) mutable { +#else + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view) mutable { +#endif + EXPECT_THAT(scope, Eq("my-instrument-scope")); + EXPECT_THAT(scope_version, Eq("v1")); + return mock_meter; + }); + + AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); + ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", + ""}; + DataLabels data_labels{"my-method", "my-streaming", "my-client-name", + "my-client-uid", "my-app-profile", ""}; + auto clone = attempt_latency.clone(resource_labels, data_labels); + + grpc::ClientContext client_context; + CreateServerMetadata(client_context, + CreateServerMetadataOptions{true, false}); + + auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); + auto clock = std::make_shared(); + + clock->SetTime(std::chrono::steady_clock::now()); + clone->PreCall(otel_context, {clock->Now(), true}); + clock->AdvanceTime(std::chrono::microseconds(1234)); + clone->PostCall(otel_context, client_context, + {clock->Now(), Status{StatusCode::kOk, "ok"}}); +} + TEST(RetryCountTest, NoRetry) { auto mock_counter = std::make_unique>(); EXPECT_CALL(*mock_counter, @@ -862,7 +1216,7 @@ TEST(RetryCountTest, NoRetry) { auto clone = retry_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -935,7 +1289,7 @@ TEST(RetryCountTest, ThreeAttempts) { auto clone = retry_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -1091,7 +1445,7 @@ TEST(FirstResponseLatency, Success) { auto clone = first_response_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -1156,7 +1510,7 @@ TEST(FirstResponseLatency, NoDataReceived) { auto clone = first_response_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -1578,7 +1932,7 @@ TEST(ServerLatency, NoServerTiming) { auto clone = server_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); clone->PostCall(otel_context, client_context, @@ -1801,7 +2155,7 @@ TEST(ConnectivityErrorCount, OkAndMissingServerTiming) { auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); clone->PostCall(otel_context, client_context, @@ -1870,7 +2224,7 @@ TEST(ConnectivityErrorCount, DeadlineExceededAndMissingServerTiming) { auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); clone->PostCall( @@ -1938,7 +2292,7 @@ TEST(ApplicationBlockingLatency, Success) { auto clone = application_blocking_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); @@ -2025,7 +2379,7 @@ TEST(ApplicationBlockingLatency, StreamingData) { auto clone = application_blocking_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; - SetClusterZone(client_context, "my-cluster", "my-zone"); + SetClusterZone(client_context); auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto clock = std::make_shared(); diff --git a/google/cloud/bigtable/internal/operation_context_factory.cc b/google/cloud/bigtable/internal/operation_context_factory.cc index 39f45fdc3069c..6337dcac4113b 100644 --- a/google/cloud/bigtable/internal/operation_context_factory.cc +++ b/google/cloud/bigtable/internal/operation_context_factory.cc @@ -276,6 +276,7 @@ std::shared_ptr MetricsOperationContextFactory::ReadRow( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -303,6 +304,7 @@ std::shared_ptr MetricsOperationContextFactory::ReadRows( std::vector> v; 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_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back( @@ -331,6 +333,7 @@ std::shared_ptr MetricsOperationContextFactory::MutateRow( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -358,6 +361,7 @@ std::shared_ptr MetricsOperationContextFactory::MutateRows( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -386,6 +390,7 @@ MetricsOperationContextFactory::CheckAndMutateRow( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -414,6 +419,7 @@ std::shared_ptr MetricsOperationContextFactory::SampleRowKeys( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -442,6 +448,7 @@ MetricsOperationContextFactory::ReadModifyWriteRow( std::vector> v; 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_)); v.emplace_back( std::make_shared(kRpc, provider_)); @@ -470,6 +477,7 @@ std::shared_ptr MetricsOperationContextFactory::PrepareQuery( std::vector> v; 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_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); @@ -495,6 +503,7 @@ std::shared_ptr MetricsOperationContextFactory::ExecuteQuery( std::vector> v; 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_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back( From 5145ea77335565b50cc3f8dfd32496296689d404 Mon Sep 17 00:00:00 2001 From: Zijie Li Date: Mon, 11 May 2026 15:02:35 -0400 Subject: [PATCH 2/2] feat(bigtable): add AttemptLatency2 metric and populate peer info labels This introduces the `AttemptLatency2` metric for DirectPath to record attempt latencies with the fields extracted from the decoded `PeerInfo` trailing metadata, populating `peer_info_labels_` and forwarding them to `IntoLabelMap`. Also added `AttemptLatency2Test` to test the newly populated peer info labels. Refactored `SetClusterZone` to use the new helper function `CreateServerMetadata`. --- google/cloud/bigtable/internal/metrics.cc | 66 +++---- google/cloud/bigtable/internal/metrics.h | 7 +- .../cloud/bigtable/internal/metrics_test.cc | 169 ++++++++++++------ 3 files changed, 145 insertions(+), 97 deletions(-) diff --git a/google/cloud/bigtable/internal/metrics.cc b/google/cloud/bigtable/internal/metrics.cc index af64a56cc7ffb..86d6a2ec778f7 100644 --- a/google/cloud/bigtable/internal/metrics.cc +++ b/google/cloud/bigtable/internal/metrics.cc @@ -115,45 +115,31 @@ GetResponseParamsFromTrailingMetadata( return absl::nullopt; } -absl::optional GetPeerInfoFromTrailingMetadata( +std::optional GetPeerInfoFromServerMetadata( grpc::ClientContext const& client_context) { - auto metadata = client_context.GetServerTrailingMetadata(); - // Base64 encoded peer info header key defined by the server. - auto iter = metadata.find("bigtable-peer-info"); - if (iter == metadata.end()) return absl::nullopt; + // The peer info is sent in the initial metadata and encoded in WebSafeBase64. std::string decoded; - if (!absl::Base64Unescape( - absl::string_view{iter->second.data(), iter->second.size()}, + auto const& init_metadata = client_context.GetServerInitialMetadata(); + auto iter_init = init_metadata.find("bigtable-peer-info"); + if (iter_init == init_metadata.end() || + !absl::WebSafeBase64Unescape( + absl::string_view{iter_init->second.data(), iter_init->second.size()}, &decoded)) { - return absl::nullopt; + // Find it in trailing metadata if not found in initial metadata or failed + // to decode. + auto const& trailing_metadata = client_context.GetServerTrailingMetadata(); + auto iter_trailing = trailing_metadata.find("bigtable-peer-info"); + if (iter_trailing == trailing_metadata.end() || + !absl::WebSafeBase64Unescape( + absl::string_view{iter_trailing->second.data(), + iter_trailing->second.size()}, + &decoded)) { + return std::nullopt; + } } google::bigtable::v2::PeerInfo p; if (p.ParseFromString(decoded)) return p; - return absl::nullopt; -} - -std::string TransportTypeToString( - google::bigtable::v2::PeerInfo::TransportType type) { - switch (type) { - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN: - return "transport_type_unknown"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_EXTERNAL: - return "transport_type_external"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_CLOUD_PATH: - return "transport_type_cloud_path"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS: - return "transport_type_direct_access"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_UNKNOWN: - return "transport_type_session_unknown"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_EXTERNAL: - return "transport_type_session_external"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_CLOUD_PATH: - return "transport_type_session_cloud_path"; - case google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_SESSION_DIRECT_ACCESS: - return "transport_type_session_direct_access"; - default: - return "transport_type_unknown"; - } + return std::nullopt; } absl::optional GetServerLatencyFromInitialMetadata( @@ -252,7 +238,7 @@ AttemptLatency::AttemptLatency( void AttemptLatency::PreCall(opentelemetry::context::Context const&, PreCallParams const& p) { - attempt_start_ = std::move(p.attempt_start); + attempt_start_ = p.attempt_start; } void AttemptLatency::PostCall(opentelemetry::context::Context const& context, @@ -289,7 +275,7 @@ AttemptLatency2::AttemptLatency2( void AttemptLatency2::PreCall(opentelemetry::context::Context const&, PreCallParams const& p) { - attempt_start_ = std::move(p.attempt_start); + attempt_start_ = p.attempt_start; } void AttemptLatency2::PostCall(opentelemetry::context::Context const& context, @@ -300,10 +286,12 @@ void AttemptLatency2::PostCall(opentelemetry::context::Context const& context, resource_labels_.cluster = response_params->cluster_id(); resource_labels_.zone = response_params->zone_id(); } - auto peer_info = GetPeerInfoFromTrailingMetadata(client_context); - peer_info_labels_.transport_type = TransportTypeToString( - peer_info ? peer_info->transport_type() - : google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN); + + auto peer_info = GetPeerInfoFromServerMetadata(client_context); + peer_info_labels_.transport_type = + absl::AsciiStrToLower(google::bigtable::v2::PeerInfo::TransportType_Name( + peer_info ? peer_info->transport_type() + : google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN)); if (peer_info) { peer_info_labels_.transport_region = peer_info->application_frontend_region(); diff --git a/google/cloud/bigtable/internal/metrics.h b/google/cloud/bigtable/internal/metrics.h index 34f9bd70ebbaa..ddaf055351d48 100644 --- a/google/cloud/bigtable/internal/metrics.h +++ b/google/cloud/bigtable/internal/metrics.h @@ -73,11 +73,10 @@ bool IsConnectivityError(google::cloud::Status const& status, absl::optional GetResponseParamsFromTrailingMetadata( grpc::ClientContext const& client_context); -absl::optional GetPeerInfoFromTrailingMetadata( +// Retrieve the peer info from server headers or trailers. Returns nullopt if +// not found or decoding or parsing fails. +std::optional GetPeerInfoFromServerMetadata( grpc::ClientContext const& client_context); -std::string TransportTypeToString( - google::bigtable::v2::PeerInfo::TransportType type); - absl::optional GetServerLatencyFromInitialMetadata( grpc::ClientContext const& client_context); diff --git a/google/cloud/bigtable/internal/metrics_test.cc b/google/cloud/bigtable/internal/metrics_test.cc index 2e2530d6c81a4..d4d924493d737 100644 --- a/google/cloud/bigtable/internal/metrics_test.cc +++ b/google/cloud/bigtable/internal/metrics_test.cc @@ -287,20 +287,68 @@ TEST(GetResponseParamsFromMetadata, EmptyHeader) { EXPECT_FALSE(result); } -TEST(GetPeerInfoFromMetadata, NonEmptyHeader) { - google::bigtable::v2::PeerInfo expected_peer_info; - expected_peer_info.set_transport_type( - google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); - expected_peer_info.set_application_frontend_region("my-region"); - expected_peer_info.set_application_frontend_subzone("my-subzone"); +struct CreateServerMetadataOptions { + bool include_peer_info = true; + bool include_response_params = true; +}; + +// Helper function to create server metadata by populating the initial and +// trailers based on the options. +void CreateServerMetadata(grpc::ClientContext& client_context, + CreateServerMetadataOptions const& options = + CreateServerMetadataOptions{}) { + RpcMetadata server_metadata; + if (options.include_peer_info) { + google::bigtable::v2::PeerInfo peer_info; + peer_info.set_transport_type( + google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); + peer_info.set_application_frontend_region("my-region"); + peer_info.set_application_frontend_subzone("my-subzone"); + server_metadata.headers.emplace( + "bigtable-peer-info", + absl::WebSafeBase64Escape(peer_info.SerializeAsString())); + } + if (options.include_response_params) { + google::bigtable::v2::ResponseParams expected_response_params; + expected_response_params.set_cluster_id("my-cluster"); + expected_response_params.set_zone_id("my-zone"); + server_metadata.trailers.emplace( + "x-goog-ext-425905942-bin", + expected_response_params.SerializeAsString()); + } + + SetServerMetadata(client_context, server_metadata); +} + +TEST(GetPeerInfoFromMetadata, NonEmptyHeaderInitial) { grpc::ClientContext client_context; + CreateServerMetadata(client_context, + CreateServerMetadataOptions{true, false}); + + auto result = GetPeerInfoFromServerMetadata(client_context); + ASSERT_TRUE(result); + EXPECT_THAT(result->transport_type(), + Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS)); + EXPECT_THAT(result->application_frontend_region(), Eq("my-region")); + EXPECT_THAT(result->application_frontend_subzone(), Eq("my-subzone")); +} + +TEST(GetPeerInfoFromMetadata, NonEmptyHeaderTrailers) { + grpc::ClientContext client_context; + google::bigtable::v2::PeerInfo peer_info; + peer_info.set_transport_type( + google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); + peer_info.set_application_frontend_region("my-region"); + peer_info.set_application_frontend_subzone("my-subzone"); RpcMetadata server_metadata; + // Set the trailers instead of headers to test that the function handles + // both. server_metadata.trailers.emplace( "bigtable-peer-info", - absl::Base64Escape(expected_peer_info.SerializeAsString())); + absl::WebSafeBase64Escape(peer_info.SerializeAsString())); SetServerMetadata(client_context, server_metadata); - auto result = GetPeerInfoFromTrailingMetadata(client_context); + auto result = GetPeerInfoFromServerMetadata(client_context); ASSERT_TRUE(result); EXPECT_THAT(result->transport_type(), Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS)); @@ -310,20 +358,20 @@ TEST(GetPeerInfoFromMetadata, NonEmptyHeader) { TEST(GetPeerInfoFromMetadata, EmptyHeader) { grpc::ClientContext client_context; - RpcMetadata server_metadata; - SetServerMetadata(client_context, server_metadata); + // The server metadata is empty when both options are false. + CreateServerMetadata(client_context, + CreateServerMetadataOptions{false, false}); - EXPECT_THAT(GetPeerInfoFromTrailingMetadata(client_context), - Eq(absl::nullopt)); + EXPECT_THAT(GetPeerInfoFromServerMetadata(client_context), Eq(std::nullopt)); } -TEST(GetPeerInfoFromMetadata, EmptyString) { +TEST(GetPeerInfoFromMetadata, EmptyStringInitial) { grpc::ClientContext client_context; RpcMetadata server_metadata; - server_metadata.trailers.emplace("bigtable-peer-info", ""); + server_metadata.headers.emplace("bigtable-peer-info", ""); SetServerMetadata(client_context, server_metadata); - auto result = GetPeerInfoFromTrailingMetadata(client_context); + auto result = GetPeerInfoFromServerMetadata(client_context); ASSERT_TRUE(result); EXPECT_THAT(result->transport_type(), Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN)); @@ -331,14 +379,62 @@ TEST(GetPeerInfoFromMetadata, EmptyString) { EXPECT_THAT(result->application_frontend_subzone(), Eq("")); } -TEST(GetPeerInfoFromMetadata, InvalidBase64) { +TEST(GetPeerInfoFromMetadata, EmptyStringTrailers) { grpc::ClientContext client_context; RpcMetadata server_metadata; - server_metadata.trailers.emplace("bigtable-peer-info", "invalid-base64!"); + server_metadata.trailers.emplace("bigtable-peer-info", ""); SetServerMetadata(client_context, server_metadata); - EXPECT_THAT(GetPeerInfoFromTrailingMetadata(client_context), - Eq(absl::nullopt)); + auto result = GetPeerInfoFromServerMetadata(client_context); + ASSERT_TRUE(result); + EXPECT_THAT(result->transport_type(), + Eq(google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_UNKNOWN)); + EXPECT_THAT(result->application_frontend_region(), Eq("")); + EXPECT_THAT(result->application_frontend_subzone(), Eq("")); +} + +TEST(GetPeerInfoFromMetadata, InvalidBase64Initial) { + { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.headers.emplace("bigtable-peer-info", "invalid-base64!"); + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromServerMetadata(client_context), + Eq(std::nullopt)); + } + { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.headers.emplace("bigtable-peer-info", + "invalid+websafe/base64"); + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromServerMetadata(client_context), + Eq(std::nullopt)); + } +} + +TEST(GetPeerInfoFromMetadata, InvalidBase64Trailers) { + { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.trailers.emplace("bigtable-peer-info", "invalid-base64!"); + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromServerMetadata(client_context), + Eq(std::nullopt)); + } + { + grpc::ClientContext client_context; + RpcMetadata server_metadata; + server_metadata.trailers.emplace("bigtable-peer-info", + "invalid+websafe/base64"); + SetServerMetadata(client_context, server_metadata); + + EXPECT_THAT(GetPeerInfoFromServerMetadata(client_context), + Eq(std::nullopt)); + } } std::unordered_map MakeAttributesMap( @@ -357,41 +453,6 @@ std::unordered_map MakeAttributesMap( return m; } -struct CreateServerMetadataOptions { - bool include_peer_info = true; - bool include_response_params = true; -}; - -// Helper function to create server metadata by populating the trailers based on -// the options. -void CreateServerMetadata(grpc::ClientContext& client_context, - CreateServerMetadataOptions const& options = - CreateServerMetadataOptions{}) { - RpcMetadata server_metadata; - if (options.include_peer_info) { - google::bigtable::v2::PeerInfo peer_info; - peer_info.set_transport_type( - google::bigtable::v2::PeerInfo::TRANSPORT_TYPE_DIRECT_ACCESS); - peer_info.set_application_frontend_region("my-region"); - peer_info.set_application_frontend_subzone("my-subzone"); - server_metadata.trailers.emplace( - "bigtable-peer-info", - absl::Base64Escape(peer_info.SerializeAsString())); - } - if (options.include_response_params) { - google::bigtable::v2::ResponseParams expected_response_params; - expected_response_params.set_cluster_id("my-cluster"); - expected_response_params.set_zone_id("my-zone"); - server_metadata.trailers.emplace( - "x-goog-ext-425905942-bin", - expected_response_params.SerializeAsString()); - } - - if (!server_metadata.trailers.empty()) { - SetServerMetadata(client_context, server_metadata); - } -} - void SetClusterZone(grpc::ClientContext& client_context) { CreateServerMetadata(client_context, CreateServerMetadataOptions{false, true});