[xDS] add optional locality label to per-call metrics (#36049)

Closes #36049 COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36049 from markdroth:xds_locality_label_on_per_call_metrics 9c5aeb7fb1 PiperOrigin-RevId: 612648213
2024-03-04 17:41:31 -08:00 · 2024-03-04 17:41:31 -08:00 · db9ca10b35
parent 692b7fcb7f
commit db9ca10b35
8 changed files with 146 additions and 62 deletions
--- a/src/core/BUILD
+++ b/src/core/BUILD
@ -5277,6 +5277,7 @@ grpc_cc_library(
        "lb_policy",
        "lb_policy_factory",
        "lb_policy_registry",
+        "match",
        "pollset_set",
        "ref_counted",
        "resolved_address",
--- a/src/core/ext/xds/xds_client_stats.h
+++ b/src/core/ext/xds/xds_client_stats.h
@ -63,7 +63,11 @@ class XdsLocalityName : public RefCounted<XdsLocalityName> {
  XdsLocalityName(std::string region, std::string zone, std::string sub_zone)
      : region_(std::move(region)),
        zone_(std::move(zone)),
-        sub_zone_(std::move(sub_zone)) {}
+        sub_zone_(std::move(sub_zone)),
+        locality_labels_(
+            std::make_shared<std::map<std::string, std::string>>()) {
+    locality_labels_->emplace("grpc.lb.locality", AsHumanReadableString());
+  }

  bool operator==(const XdsLocalityName& other) const {
    return region_ == other.region_ && zone_ == other.zone_ &&
@ -85,6 +89,9 @@ class XdsLocalityName : public RefCounted<XdsLocalityName> {
  const std::string& region() const { return region_; }
  const std::string& zone() const { return zone_; }
  const std::string& sub_zone() const { return sub_zone_; }
+  std::shared_ptr<std::map<std::string, std::string>> locality_labels() const {
+    return locality_labels_;
+  }

  std::string AsHumanReadableString() const {
    return absl::StrFormat("{region=\"%s\", zone=\"%s\", sub_zone=\"%s\"}",
@ -104,6 +111,7 @@ class XdsLocalityName : public RefCounted<XdsLocalityName> {
  std::string region_;
  std::string zone_;
  std::string sub_zone_;
+  std::shared_ptr<std::map<std::string, std::string>> locality_labels_;
 };

 // Drop stats for an xds cluster.
@ -223,6 +231,8 @@ class XdsClusterLocalityStats : public RefCounted<XdsClusterLocalityStats> {
  void AddCallFinished(const std::map<absl::string_view, double>* named_metrics,
                       bool fail = false);

+  XdsLocalityName* locality_name() const { return name_.get(); }
+
 private:
  struct Stats {
    std::atomic<uint64_t> total_successful_requests{0};
--- a/src/core/lib/channel/call_tracer.h
+++ b/src/core/lib/channel/call_tracer.h
@ -128,8 +128,9 @@ class ClientCallTracer : public CallTracerAnnotationInterface {
  class CallAttemptTracer : public CallTracerInterface {
   public:
    enum class OptionalLabelComponent : std::uint8_t {
-      kXdsServiceLabels = 0,
-      kSize = 1,  // keep last
+      kXdsServiceLabels,
+      kXdsLocalityLabels,
+      kSize,  // keep last
    };

    ~CallAttemptTracer() override {}
--- a/src/core/load_balancing/xds/xds_cluster_impl.cc
+++ b/src/core/load_balancing/xds/xds_cluster_impl.cc
@ -51,6 +51,7 @@
 #include "src/core/lib/config/core_configuration.h"
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/gprpp/debug_location.h"
+#include "src/core/lib/gprpp/match.h"
 #include "src/core/lib/gprpp/orphanable.h"
 #include "src/core/lib/gprpp/ref_counted.h"
 #include "src/core/lib/gprpp/ref_counted_ptr.h"
@ -191,18 +192,46 @@ class XdsClusterImplLb : public LoadBalancingPolicy {
 private:
  class StatsSubchannelWrapper : public DelegatingSubchannel {
   public:
+    // If load reporting is enabled and we have an XdsClusterLocalityStats
+    // object, that object already contains the locality labels.  We
+    // need to store the locality labels directly only in the case where
+    // load reporting is disabled.
+    using LocalityData = absl::variant<
+        std::shared_ptr<std::map<std::string, std::string>> /*locality_labels*/,
+        RefCountedPtr<XdsClusterLocalityStats> /*locality_stats*/>;
+
    StatsSubchannelWrapper(
        RefCountedPtr<SubchannelInterface> wrapped_subchannel,
-        RefCountedPtr<XdsClusterLocalityStats> locality_stats)
+        LocalityData locality_data)
        : DelegatingSubchannel(std::move(wrapped_subchannel)),
-          locality_stats_(std::move(locality_stats)) {}
+          locality_data_(std::move(locality_data)) {}
+
+    std::shared_ptr<std::map<std::string, std::string>> locality_labels()
+        const {
+      return Match(
+          locality_data_,
+          [](const std::shared_ptr<std::map<std::string, std::string>>&
+                 locality_labels) {
+            return locality_labels;
+          },
+          [](const RefCountedPtr<XdsClusterLocalityStats>& locality_stats) {
+            return locality_stats->locality_name()->locality_labels();
+          });
+    }

    XdsClusterLocalityStats* locality_stats() const {
-      return locality_stats_.get();
+      return Match(
+          locality_data_,
+          [](const std::shared_ptr<std::map<std::string, std::string>>&) {
+            return static_cast<XdsClusterLocalityStats*>(nullptr);
+          },
+          [](const RefCountedPtr<XdsClusterLocalityStats>& locality_stats) {
+            return locality_stats.get();
+          });
    }

   private:
-    RefCountedPtr<XdsClusterLocalityStats> locality_stats_;
+    LocalityData locality_data_;
  };

  // A picker that wraps the picker from the child to perform drops.
@ -375,8 +404,9 @@ XdsClusterImplLb::Picker::Picker(XdsClusterImplLb* xds_cluster_impl_lb,
 LoadBalancingPolicy::PickResult XdsClusterImplLb::Picker::Pick(
    LoadBalancingPolicy::PickArgs args) {
  auto* call_state = static_cast<ClientChannelLbCallState*>(args.call_state);
-  if (call_state->GetCallAttemptTracer() != nullptr) {
-    call_state->GetCallAttemptTracer()->AddOptionalLabels(
+  auto* call_attempt_tracer = call_state->GetCallAttemptTracer();
+  if (call_attempt_tracer != nullptr) {
+    call_attempt_tracer->AddOptionalLabels(
        OptionalLabelComponent::kXdsServiceLabels, service_labels_);
  }
  // Handle EDS drops.
@ -404,16 +434,22 @@ LoadBalancingPolicy::PickResult XdsClusterImplLb::Picker::Pick(
  PickResult result = picker_->Pick(args);
  auto* complete_pick = absl::get_if<PickResult::Complete>(&result.result);
  if (complete_pick != nullptr) {
+    auto* subchannel_wrapper =
+        static_cast<StatsSubchannelWrapper*>(complete_pick->subchannel.get());
+    // Add locality labels to per-call metrics if needed.
+    if (call_attempt_tracer != nullptr) {
+      call_attempt_tracer->AddOptionalLabels(
+          OptionalLabelComponent::kXdsLocalityLabels,
+          subchannel_wrapper->locality_labels());
+    }
+    // Handle load reporting.
    RefCountedPtr<XdsClusterLocalityStats> locality_stats;
-    if (drop_stats_ != nullptr) {  // If load reporting is enabled.
-      auto* subchannel_wrapper =
-          static_cast<StatsSubchannelWrapper*>(complete_pick->subchannel.get());
-      // Handle load reporting.
+    if (subchannel_wrapper->locality_stats() != nullptr) {
      locality_stats = subchannel_wrapper->locality_stats()->Ref(
          DEBUG_LOCATION, "SubchannelCallTracker");
-      // Unwrap subchannel to pass back up the stack.
-      complete_pick->subchannel = subchannel_wrapper->wrapped_subchannel();
    }
+    // Unwrap subchannel to pass back up the stack.
+    complete_pick->subchannel = subchannel_wrapper->wrapped_subchannel();
    // Inject subchannel call tracker to record call completion.
    complete_pick->subchannel_call_tracker =
        std::make_unique<SubchannelCallTracker>(
@ -743,36 +779,39 @@ RefCountedPtr<SubchannelInterface> XdsClusterImplLb::Helper::CreateSubchannel(
    const grpc_resolved_address& address, const ChannelArgs& per_address_args,
    const ChannelArgs& args) {
  if (parent()->shutting_down_) return nullptr;
-  // If load reporting is enabled, wrap the subchannel such that it
-  // includes the locality stats object, which will be used by the Picker.
+  // Wrap the subchannel so that we pass along the locality labels and
+  // (if load reporting is enabled) the locality stats object, which
+  // will be used by the picker.
+  auto locality_name = per_address_args.GetObjectRef<XdsLocalityName>();
+  RefCountedPtr<XdsClusterLocalityStats> locality_stats;
  if (parent()->cluster_resource_->lrs_load_reporting_server.has_value()) {
-    auto locality_name = per_address_args.GetObjectRef<XdsLocalityName>();
-    RefCountedPtr<XdsClusterLocalityStats> locality_stats =
-        parent()->xds_client_->AddClusterLocalityStats(
-            parent()->cluster_resource_->lrs_load_reporting_server.value(),
-            parent()->config_->cluster_name(),
-            GetEdsResourceName(*parent()->cluster_resource_),
-            std::move(locality_name));
-    if (locality_stats != nullptr) {
-      return MakeRefCounted<StatsSubchannelWrapper>(
-          parent()->channel_control_helper()->CreateSubchannel(
-              address, per_address_args, args),
-          std::move(locality_stats));
+    locality_stats = parent()->xds_client_->AddClusterLocalityStats(
+        parent()->cluster_resource_->lrs_load_reporting_server.value(),
+        parent()->config_->cluster_name(),
+        GetEdsResourceName(*parent()->cluster_resource_), locality_name);
+    if (locality_stats == nullptr) {
+      gpr_log(GPR_ERROR,
+              "[xds_cluster_impl_lb %p] Failed to get locality stats object "
+              "for LRS server %s, cluster %s, EDS service name %s; load "
+              "reports will not be generated",
+              parent(),
+              parent()
+                  ->cluster_resource_->lrs_load_reporting_server->server_uri()
+                  .c_str(),
+              parent()->config_->cluster_name().c_str(),
+              GetEdsResourceName(*parent()->cluster_resource_).c_str());
    }
-    gpr_log(GPR_ERROR,
-            "[xds_cluster_impl_lb %p] Failed to get locality stats object for "
-            "LRS server %s, cluster %s, EDS service name %s; load reports will "
-            "not be generated (not wrapping subchannel)",
-            parent(),
-            parent()
-                ->cluster_resource_->lrs_load_reporting_server->server_uri()
-                .c_str(),
-            parent()->config_->cluster_name().c_str(),
-            GetEdsResourceName(*parent()->cluster_resource_).c_str());
  }
-  // Load reporting not enabled, so don't wrap the subchannel.
-  return parent()->channel_control_helper()->CreateSubchannel(
-      address, per_address_args, args);
+  StatsSubchannelWrapper::LocalityData locality_data;
+  if (locality_stats != nullptr) {
+    locality_data = std::move(locality_stats);
+  } else {
+    locality_data = locality_name->locality_labels();
+  }
+  return MakeRefCounted<StatsSubchannelWrapper>(
+      parent()->channel_control_helper()->CreateSubchannel(
+          address, per_address_args, args),
+      std::move(locality_data));
 }

 void XdsClusterImplLb::Helper::UpdateState(
--- a/test/cpp/end2end/xds/xds_cluster_end2end_test.cc
+++ b/test/cpp/end2end/xds/xds_cluster_end2end_test.cc
@ -309,12 +309,13 @@ TEST_P(CdsTest, ClusterChangeAfterAdsCallFails) {
  WaitForBackend(DEBUG_LOCATION, 1);
 }

-TEST_P(CdsTest, VerifyCsmServiceLabelsParsing) {
+TEST_P(CdsTest, MetricLabels) {
  // Injects a fake client call tracer factory. Try keep this at top.
  grpc_core::FakeClientCallTracerFactory fake_client_call_tracer_factory;
-  CreateAndStartBackends(1);
+  CreateAndStartBackends(2);
  // Populates EDS resources.
-  EdsResourceArgs args({{"locality0", CreateEndpointsForBackends()}});
+  EdsResourceArgs args({{"locality0", CreateEndpointsForBackends(0, 1)},
+                        {"locality1", CreateEndpointsForBackends(1, 2)}});
  balancer_->ads_service()->SetEdsResource(BuildEdsResource(args));
  // Populates service labels to CDS resources.
  auto cluster = default_cluster_;
@ -328,17 +329,46 @@ TEST_P(CdsTest, VerifyCsmServiceLabelsParsing) {
  channel_args.SetPointer(GRPC_ARG_INJECT_FAKE_CLIENT_CALL_TRACER_FACTORY,
                          &fake_client_call_tracer_factory);
  ResetStub(/*failover_timeout_ms=*/0, &channel_args);
-  // Sends an RPC and verifies that the service labels are recorded in the fake
-  // client call tracer.
-  CheckRpcSendOk(DEBUG_LOCATION);
-  EXPECT_THAT(fake_client_call_tracer_factory.GetLastFakeClientCallTracer()
-                  ->GetLastCallAttemptTracer()
-                  ->GetOptionalLabels(),
-              ::testing::ElementsAre(::testing::Pair(
-                  OptionalLabelComponent::kXdsServiceLabels,
-                  ::testing::Pointee(::testing::ElementsAre(
-                      ::testing::Pair("service_name", "myservice"),
-                      ::testing::Pair("service_namespace", "mynamespace"))))));
+  // Send an RPC to backend 0.
+  WaitForBackend(DEBUG_LOCATION, 0);
+  // Verify that the optional labels are recorded in the call tracer.
+  EXPECT_THAT(
+      fake_client_call_tracer_factory.GetLastFakeClientCallTracer()
+          ->GetLastCallAttemptTracer()
+          ->GetOptionalLabels(),
+      ::testing::ElementsAre(
+          ::testing::Pair(
+              OptionalLabelComponent::kXdsServiceLabels,
+              ::testing::Pointee(::testing::ElementsAre(
+                  ::testing::Pair("service_name", "myservice"),
+                  ::testing::Pair("service_namespace", "mynamespace")))),
+          ::testing::Pair(
+              OptionalLabelComponent::kXdsLocalityLabels,
+              ::testing::Pointee(::testing::ElementsAre(::testing::Pair(
+                  "grpc.lb.locality", LocalityNameString("locality0")))))));
+  // Send an RPC to backend 1.
+  WaitForBackend(DEBUG_LOCATION, 1);
+  // Verify that the optional labels are recorded in the call tracer.
+  EXPECT_THAT(
+      fake_client_call_tracer_factory.GetLastFakeClientCallTracer()
+          ->GetLastCallAttemptTracer()
+          ->GetOptionalLabels(),
+      ::testing::ElementsAre(
+          ::testing::Pair(
+              OptionalLabelComponent::kXdsServiceLabels,
+              ::testing::Pointee(::testing::ElementsAre(
+                  ::testing::Pair("service_name", "myservice"),
+                  ::testing::Pair("service_namespace", "mynamespace")))),
+          ::testing::Pair(
+              OptionalLabelComponent::kXdsLocalityLabels,
+              ::testing::Pointee(::testing::ElementsAre(::testing::Pair(
+                  "grpc.lb.locality", LocalityNameString("locality1")))))));
+  // TODO(yashkt, yijiem): This shutdown shouldn't actually be necessary.
+  // The only reason it's here is to add a delay before
+  // fake_client_call_tracer_factory goes out of scope, since there may
+  // be lingering callbacks in the call stack that are using the
+  // CallAttemptTracer even after we get here, which would then cause a
+  // crash.  Find a cleaner way to fix this.
  balancer_->Shutdown();
 }

--- a/test/cpp/end2end/xds/xds_utils.cc
+++ b/test/cpp/end2end/xds/xds_utils.cc
@ -317,6 +317,12 @@ void XdsResourceUtils::SetRouteConfiguration(
  }
 }

+std::string XdsResourceUtils::LocalityNameString(absl::string_view sub_zone) {
+  return absl::StrFormat("{region=\"%s\", zone=\"%s\", sub_zone=\"%s\"}",
+                         kDefaultLocalityRegion, kDefaultLocalityZone,
+                         sub_zone);
+}
+
 ClusterLoadAssignment XdsResourceUtils::BuildEdsResource(
    const EdsResourceArgs& args, absl::string_view eds_service_name) {
  ClusterLoadAssignment assignment;
--- a/test/cpp/end2end/xds/xds_utils.h
+++ b/test/cpp/end2end/xds/xds_utils.h
@ -195,6 +195,9 @@ class XdsResourceUtils {
                                    bool use_rds = false,
                                    const Listener* listener_to_copy = nullptr);

+  // Returns a string representing the locality with the specified sub_zone.
+  static std::string LocalityNameString(absl::string_view sub_zone);
+
  // Arguments for constructing an EDS resource.
  struct EdsResourceArgs {
    // An individual endpoint for a backend running on a specified port.
--- a/test/cpp/end2end/xds/xds_wrr_end2end_test.cc
+++ b/test/cpp/end2end/xds/xds_wrr_end2end_test.cc
@ -48,12 +48,6 @@ class WrrTest : public XdsEnd2endTest {
  void SetUp() override {
    // No-op -- tests must explicitly call InitClient().
  }
-
-  static std::string LocalityNameString(absl::string_view sub_zone) {
-    return absl::StrFormat("{region=\"%s\", zone=\"%s\", sub_zone=\"%s\"}",
-                           kDefaultLocalityRegion, kDefaultLocalityZone,
-                           sub_zone);
-  }
 };

 INSTANTIATE_TEST_SUITE_P(XdsTest, WrrTest, ::testing::Values(XdsTestType()),