experimental/stats: Expose Telemetry Label Callback

experimental/stats/telemetry.go

@@ -0,0 +1,59 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package stats
+
+import (
+	"context"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+// LabelCallback is a function that is executed when telemetry
+// label keys are updated
+type LabelCallback func(string, string)
+
+type telemetryLabelCallbackKey struct{}
+
+// WithTelemetryLabelCallback registers a callback function that is executed whenever
+// telemetry labels will be updated. This does _not_ require opentelemetry instrumentation
+// to be configured on the client or server.
+func WithTelemetryLabelCallback(ctx context.Context, callback LabelCallback) context.Context {
+	if callback == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, telemetryLabelCallbackKey{}, callback)
+}
+
+// ExecuteTelemetryLabelCallback runs the registered callback on the context with the provided
+// key and value. If no callback is registered it does nothing.
+//
+// If the registered callback panics it will be swallowed and logged
+func ExecuteTelemetryLabelCallback(ctx context.Context, key, value string) {
+	if ctx == nil {
+		return
+	}
+	if callback, ok := ctx.Value(telemetryLabelCallbackKey{}).(LabelCallback); ok {
+		defer func() {
+			if r := recover(); r != nil {
+				grpclog.Component("experimental-stats").Warningf("LabelCallback panicked: %v", r)
+			}
+		}()
+		callback(key, value)
+	}
+}

experimental/stats/telemetry_test.go

@@ -0,0 +1,87 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package stats
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+// TestTelemetryLabels tests registering a callback function with the context and
+// the effects of executing the callback on a local label state tracker. Each test
+// case constructs a new context with the provided callback registered.
+func (s) TestTelemetryLabels(t *testing.T) {
+	commonLabelValues := map[string]string{"grpc.lb.backend_service": "grpc.lb.backend_service_val", "grpc.lb.locality_val": "grpc.lb.locality_val"}
+	tracker := map[string]string{}
+
+	tests := map[string]struct {
+		callback         func(string, string)
+		additionalLabels map[string]string
+		wantLabels       map[string]string
+	}{
+		"NilCallback": {
+			callback:   nil,
+			wantLabels: map[string]string{},
+		},
+		"NoOPCallback": {
+			callback:   func(string, string) {},
+			wantLabels: map[string]string{},
+		},
+		"PanicCallback": {
+			callback:   func(string, string) { panic("intentional panic") },
+			wantLabels: map[string]string{},
+		},
+		"MutatingCallback": {
+			callback: func(key, value string) {
+				tracker[key] = value
+			},
+			wantLabels: map[string]string{"grpc.lb.backend_service": "grpc.lb.backend_service_val", "grpc.lb.locality_val": "grpc.lb.locality_val"},
+		},
+		"OverrideLabelsWithCallback": {
+			callback: func(key, value string) {
+				tracker[key] = value
+			},
+			additionalLabels: map[string]string{"grpc.lb.backend_service": "grpc.lb.backend_service_other_val"},
+			wantLabels:       map[string]string{"grpc.lb.backend_service": "grpc.lb.backend_service_other_val", "grpc.lb.locality_val": "grpc.lb.locality_val"},
+		},
+	}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			// resest the tracker at the end of every test
+			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+			t.Cleanup(func() {
+				tracker = map[string]string{}
+				cancel()
+			})
+			ctx = WithTelemetryLabelCallback(ctx, test.callback)
+			for k, v := range commonLabelValues {
+				ExecuteTelemetryLabelCallback(ctx, k, v)
+			}
+			for k, v := range test.additionalLabels {
+				ExecuteTelemetryLabelCallback(ctx, k, v)
+			}
+			if diff := cmp.Diff(tracker, test.wantLabels); diff != "" {
+				t.Fatalf("tracked labels did not match expcted values (-got, +want): %v", diff)
+			}
+		})
+	}
+}

internal/xds/balancer/clusterimpl/picker.go

@@ -26,6 +26,7 @@ import (
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
+	estats "google.golang.org/grpc/experimental/stats"
 	"google.golang.org/grpc/internal/stats"
 	"google.golang.org/grpc/internal/wrr"
 	xdsinternal "google.golang.org/grpc/internal/xds"
@@ -164,6 +165,8 @@ func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
 		return pr, err
 	}
 
+	estats.ExecuteTelemetryLabelCallback(info.Ctx, "grpc.lb.locality", xdsinternal.LocalityString(lID))
+	estats.ExecuteTelemetryLabelCallback(info.Ctx, "grpc.lb.backend_service", d.clusterName)
 	if labels != nil {
 		labels["grpc.lb.locality"] = xdsinternal.LocalityString(lID)
 		labels["grpc.lb.backend_service"] = d.clusterName