stats: Add optional locality label in cluster_impl picker (#7434)

stats/opentelemetry/client_metrics.go

@@ -150,7 +150,12 @@ func (h *clientStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo)
  var labels *istats.Labels
  if labels = istats.GetLabels(ctx); labels == nil {
  labels = &istats.Labels{
- TelemetryLabels: make(map[string]string),
+ // The defaults for all the per call labels from a plugin that
+ // executes on the callpath that this OpenTelemetry component
+ // currently supports.
+ TelemetryLabels: map[string]string{
+ "grpc.lb.locality": "",
+ },
  }
  ctx = istats.SetLabels(ctx, labels)
  }
@@ -232,6 +237,8 @@ func (h *clientStatsHandler) processRPCEnd(ctx context.Context, ai *attemptInfo,
  }
 
  for _, o := range h.options.MetricsOptions.OptionalLabels {
+ // TODO: Add a filter for converting to unknown if not present in the
+ // CSM Plugin Option layer by adding an optional labels API.
  if val, ok := ai.xdsLabels[o]; ok {
  attributes = append(attributes, otelattribute.String(o, val))
  }

stats/opentelemetry/csm/observability.go

@@ -70,6 +70,10 @@ func (o *perTargetDialOption) DialOptionForTarget(parsedTarget url.URL) grpc.Dia
 
 func dialOptionWithCSMPluginOption(options opentelemetry.Options, po otelinternal.PluginOption) grpc.DialOption {
  options.MetricsOptions.OptionalLabels = []string{"csm.service_name", "csm.service_namespace_name"} // Attach the two xDS Optional Labels for this component to not filter out.
+ return dialOptionSetCSM(options, po)
+}
+
+func dialOptionSetCSM(options opentelemetry.Options, po otelinternal.PluginOption) grpc.DialOption {
  otelinternal.SetPluginOption.(func(options *opentelemetry.Options, po otelinternal.PluginOption))(&options, po)
  return opentelemetry.DialOption(options)
 }

stats/opentelemetry/csm/observability_test.go

@@ -425,9 +425,12 @@ func (s) TestCSMPluginOptionStreaming(t *testing.T) {
 func unaryInterceptorAttachXDSLabels(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
  ctx = istats.SetLabels(ctx, &istats.Labels{
  TelemetryLabels: map[string]string{
- // mock what the cluster impl would write here ("csm." xDS Labels)
+ // mock what the cluster impl would write here ("csm." xDS Labels
+ // and locality label)
  "csm.service_name": "service_name_val",
  "csm.service_namespace_name": "service_namespace_val",
+
+ "grpc.lb.locality": "grpc.lb.locality_val",
  },
  })
 
@@ -441,8 +444,9 @@ func unaryInterceptorAttachXDSLabels(ctx context.Context, method string, req, re
 // Optional Labels turned on. It then configures an interceptor to attach
 // labels, representing the cluster_impl picker. It then makes a unary RPC, and
 // expects xDS Labels labels to be attached to emitted relevant metrics. Full
-// xDS System alongside OpenTelemetry will be tested with interop. (there is
-// a test for xDS -> Stats handler and this tests -> OTel -> emission).
+// xDS System alongside OpenTelemetry will be tested with interop. (there is a
+// test for xDS -> Stats handler and this tests -> OTel -> emission). It also
+// tests the optional per call locality label in the same manner.
 func (s) TestXDSLabels(t *testing.T) {
  ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
  defer cancel()
@@ -457,11 +461,11 @@ func (s) TestXDSLabels(t *testing.T) {
  }
 
  po := newPluginOption(ctx)
- dopts := []grpc.DialOption{dialOptionWithCSMPluginOption(opentelemetry.Options{
+ dopts := []grpc.DialOption{dialOptionSetCSM(opentelemetry.Options{
  MetricsOptions: opentelemetry.MetricsOptions{
  MeterProvider: provider,
  Metrics: opentelemetry.DefaultMetrics(),
- OptionalLabels: []string{"csm.service_name", "csm.service_namespace_name"},
+ OptionalLabels: []string{"csm.service_name", "csm.service_namespace_name", "grpc.lb.locality"},
  },
  }, po), grpc.WithUnaryInterceptor(unaryInterceptorAttachXDSLabels)}
  if err := ss.Start(nil, dopts...); err != nil {
@@ -489,6 +493,7 @@ func (s) TestXDSLabels(t *testing.T) {
 
  serviceNameAttr := attribute.String("csm.service_name", "service_name_val")
  serviceNamespaceAttr := attribute.String("csm.service_namespace_name", "service_namespace_val")
+ localityAttr := attribute.String("grpc.lb.locality", "grpc.lb.locality_val")
  meshIDAttr := attribute.String("csm.mesh_id", "unknown")
  workloadCanonicalServiceAttr := attribute.String("csm.workload_canonical_service", "unknown")
  remoteWorkloadTypeAttr := attribute.String("csm.remote_workload_type", "unknown")
@@ -500,6 +505,7 @@ func (s) TestXDSLabels(t *testing.T) {
  unaryStatusAttr,
  serviceNameAttr,
  serviceNamespaceAttr,
+ localityAttr,
  meshIDAttr,
  workloadCanonicalServiceAttr,
  remoteWorkloadTypeAttr,

test/xds/xds_telemetry_labels_test.go

@@ -28,11 +28,12 @@ import (
  "google.golang.org/grpc/internal/stubserver"
  "google.golang.org/grpc/internal/testutils"
  "google.golang.org/grpc/internal/testutils/xds/e2e"
+ testgrpc "google.golang.org/grpc/interop/grpc_testing"
+ testpb "google.golang.org/grpc/interop/grpc_testing"
  "google.golang.org/grpc/stats"
 
  v3corepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
- testgrpc "google.golang.org/grpc/interop/grpc_testing"
- testpb "google.golang.org/grpc/interop/grpc_testing"
+ "github.com/google/go-cmp/cmp"
  "google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -43,6 +44,9 @@ const serviceNamespaceKeyCSM = "csm.service_namespace_name"
 const serviceNameValue = "grpc-service"
 const serviceNamespaceValue = "grpc-service-namespace"
 
+const localityKey = "grpc.lb.locality"
+const localityValue = `{"region":"region-1","zone":"zone-1","subZone":"subzone-1"}`
+
 // TestTelemetryLabels tests that telemetry labels from CDS make their way to
 // the stats handler. The stats handler sets the mutable context value that the
 // cluster impl picker will write telemetry labels to, and then the stats
@@ -126,13 +130,14 @@ func (fsh *fakeStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
  // aren't started. All of these should have access to the desired telemetry
  // labels.
  case *stats.OutPayload, *stats.InPayload, *stats.End:
- if label, ok := fsh.labels.TelemetryLabels[serviceNameKeyCSM]; !ok || label != serviceNameValue {
- fsh.t.Fatalf("for telemetry label %v, want: %v, got: %v", serviceNameKeyCSM, serviceNameValue, label)
+ want := map[string]string{
+ serviceNameKeyCSM: serviceNameValue,
+ serviceNamespaceKeyCSM: serviceNamespaceValue,
+ localityKey: localityValue,
  }
- if label, ok := fsh.labels.TelemetryLabels[serviceNamespaceKeyCSM]; !ok || label != serviceNamespaceValue {
- fsh.t.Fatalf("for telemetry label %v, want: %v, got: %v", serviceNamespaceKeyCSM, serviceNamespaceValue, label)
+ if diff := cmp.Diff(fsh.labels.TelemetryLabels, want); diff != "" {
+ fsh.t.Fatalf("fsh.labels.TelemetryLabels (-got +want): %v", diff)
  }
-
  default:
  // Nothing to assert for the other stats.Handler callouts.
  }

xds/internal/balancer/clusterimpl/picker.go

@@ -19,6 +19,8 @@
 package clusterimpl
 
 import (
+ "context"
+
  v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
  "google.golang.org/grpc/balancer"
  "google.golang.org/grpc/codes"
@@ -96,14 +98,23 @@ func (b *clusterImplBalancer) newPicker(config *dropConfigs) *picker {
  }
 }
 
+func telemetryLabels(ctx context.Context) map[string]string {
+ if ctx == nil {
+ return nil
+ }
+ labels := stats.GetLabels(ctx)
+ if labels == nil {
+ return nil
+ }
+ return labels.TelemetryLabels
+}
+
 func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
  // Unconditionally set labels if present, even dropped or queued RPC's can
  // use these labels.
- if info.Ctx != nil {
- if labels := stats.GetLabels(info.Ctx); labels != nil && labels.TelemetryLabels != nil {
- for key, value := range d.telemetryLabels {
- labels.TelemetryLabels[key] = value
- }
+ if labels := telemetryLabels(info.Ctx); labels != nil {
+ for key, value := range d.telemetryLabels {
+ labels[key] = value
  }
  }
 
@@ -156,6 +167,10 @@ func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
  return pr, err
  }
 
+ if labels := telemetryLabels(info.Ctx); labels != nil {
+ labels["grpc.lb.locality"] = lIDStr
+ }
+
  if d.loadStore != nil {
  d.loadStore.CallStarted(lIDStr)
  oldDone := pr.Done