Skip to content

Commit 5fbf875

Browse files
mbissambissa
committed
review comments
1 parent 0d427dc commit 5fbf875

File tree

6 files changed

+36
-29
lines changed

6 files changed

+36
-29
lines changed

‎balancer/rls/metrics_test.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,10 @@ import (
2626
"go.opentelemetry.io/otel/sdk/metric"
2727
"go.opentelemetry.io/otel/sdk/metric/metricdata"
2828
"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
29+
2930
"google.golang.org/grpc"
3031
"google.golang.org/grpc/credentials/insecure"
32+
estats "google.golang.org/grpc/experimental/stats"
3133
rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
3234
"google.golang.org/grpc/internal/stubserver"
3335
rlstest "google.golang.org/grpc/internal/testutils/rls"
@@ -135,7 +137,7 @@ func (s) TestRLSTargetPickMetric(t *testing.T) {
135137
client := testgrpc.NewTestServiceClient(cc)
136138
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
137139
defer cancel()
138-
ctx = grpc.NewContextWithCustomLabel(ctx, "target-pick-custom")
140+
ctx = estats.NewContextWithCustomLabel(ctx, "target-pick-custom")
139141
_, err = client.EmptyCall(ctx, &testpb.Empty{})
140142
if err != nil {
141143
t.Fatalf("client.EmptyCall failed with error: %v", err)
@@ -248,7 +250,7 @@ func (s) TestRLSDefaultTargetPickMetric(t *testing.T) {
248250
client := testgrpc.NewTestServiceClient(cc)
249251
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
250252
defer cancel()
251-
ctx = grpc.NewContextWithCustomLabel(ctx, "default-target-pick-custom")
253+
ctx = estats.NewContextWithCustomLabel(ctx, "default-target-pick-custom")
252254
if _, err = client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
253255
t.Fatalf("client.EmptyCall failed with error: %v", err)
254256
}
@@ -347,7 +349,7 @@ func (s) TestRLSFailedRPCMetric(t *testing.T) {
347349

348350
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
349351
defer cancel()
350-
ctx = grpc.NewContextWithCustomLabel(ctx, "failed-pick-custom")
352+
ctx = estats.NewContextWithCustomLabel(ctx, "failed-pick-custom")
351353
client := testgrpc.NewTestServiceClient(cc)
352354
if _, err = client.EmptyCall(ctx, &testpb.Empty{}); err == nil {
353355
t.Fatalf("client.EmptyCall error = %v, expected a non nil error", err)

‎balancer/rls/picker.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ import (
2525
"sync/atomic"
2626
"time"
2727

28-
"google.golang.org/grpc"
2928
"google.golang.org/grpc/balancer"
3029
"google.golang.org/grpc/balancer/rls/internal/keys"
3130
"google.golang.org/grpc/codes"
@@ -199,7 +198,7 @@ func (p *rlsPicker) delegateToChildPoliciesLocked(dcEntry *cacheEntry, info bala
199198
res, err := state.Picker.Pick(info)
200199
if err != nil {
201200
pr := errToPickResult(err)
202-
customLabel, _ := grpc.CustomLabelFromContext(info.Ctx)
201+
customLabel := estats.CustomLabelFromContext(info.Ctx)
203202
return res, func() {
204203
if pr == "queue" {
205204
// Don't record metrics for queued Picks.
@@ -215,7 +214,7 @@ func (p *rlsPicker) delegateToChildPoliciesLocked(dcEntry *cacheEntry, info bala
215214
res.Metadata.Append(rlsDataHeaderName, dcEntry.headerData)
216215
}
217216
return res, func() {
218-
customLabel, _ := grpc.CustomLabelFromContext(info.Ctx)
217+
customLabel := estats.CustomLabelFromContext(info.Ctx)
219218
targetPicksMetric.Record(p.metricsRecorder, 1, p.grpcTarget, p.rlsServerTarget, cpw.target, "complete", customLabel)
220219
}, nil
221220
}
@@ -230,7 +229,7 @@ func (p *rlsPicker) delegateToChildPoliciesLocked(dcEntry *cacheEntry, info bala
230229
// target if one is configured, or fails the pick with the given error. Returns
231230
// a function to be invoked to record metrics.
232231
func (p *rlsPicker) useDefaultPickIfPossible(info balancer.PickInfo, errOnNoDefault error) (balancer.PickResult, func(), error) {
233-
customLabel, _ := grpc.CustomLabelFromContext(info.Ctx)
232+
customLabel := estats.CustomLabelFromContext(info.Ctx)
234233

235234
if p.defaultPolicy != nil {
236235
state := (*balancer.State)(atomic.LoadPointer(&p.defaultPolicy.state))

‎experimental/stats/metrics.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,27 @@
2020
package stats
2121

2222
import (
23+
"context"
24+
2325
"google.golang.org/grpc/internal"
2426
"google.golang.org/grpc/stats"
2527
)
2628

29+
type customLabelKey struct{}
30+
31+
// NewContextWithCustomLabel returns a new context with the provided custom label
32+
// attached. The label will be propagated to all metric instruments specified in gRFC A108.
33+
func NewContextWithCustomLabel(ctx context.Context, label string) context.Context {
34+
return context.WithValue(ctx, customLabelKey{}, label)
35+
}
36+
37+
// CustomLabelFromContext returns the custom label from the context if it exists.
38+
// If the custom label is not present, it returns an empty string.
39+
func CustomLabelFromContext(ctx context.Context) string {
40+
label, _ := ctx.Value(customLabelKey{}).(string)
41+
return label
42+
}
43+
2744
// MetricsRecorder records on metrics derived from metric registry.
2845
// Implementors must embed UnimplementedMetricsRecorder.
2946
type MetricsRecorder interface {

‎rpc_util.go‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1175,19 +1175,3 @@ const (
11751175
)
11761176

11771177
const grpcUA = "grpc-go/" + Version
1178-
1179-
type customLabelKey struct{}
1180-
1181-
// NewContextWithCustomLabel returns a new context with the provided custom label
1182-
// attached. The label will be propagated to per-call metrics (e.g., OpenTelemetry
1183-
// and RLS load balancer metrics).
1184-
func NewContextWithCustomLabel(ctx context.Context, label string) context.Context {
1185-
return context.WithValue(ctx, customLabelKey{}, label)
1186-
}
1187-
1188-
// CustomLabelFromContext returns the custom label from the context if it exists.
1189-
// If the custom label is not present, it returns an empty string and false.
1190-
func CustomLabelFromContext(ctx context.Context) (string, bool) {
1191-
label, ok := ctx.Value(customLabelKey{}).(string)
1192-
return label, ok
1193-
}

‎stats/opentelemetry/client_metrics.go‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ func (h *clientMetricsHandler) perCallMetrics(ctx context.Context, err error, st
141141
}
142142
for _, o := range h.options.MetricsOptions.OptionalLabels {
143143
if o == "grpc.client.call.custom" {
144-
label, _ := grpc.CustomLabelFromContext(ctx)
144+
label := estats.CustomLabelFromContext(ctx)
145145
attributes = append(attributes, otelattribute.String(o, label))
146146
}
147147
}
@@ -220,8 +220,10 @@ func (h *clientMetricsHandler) processRPCEvent(ctx context.Context, s stats.RPCS
220220
}
221221
for _, o := range h.options.MetricsOptions.OptionalLabels {
222222
if o == "grpc.client.call.custom" {
223-
label, _ := grpc.CustomLabelFromContext(ctx)
224-
attributes = append(attributes, otelattribute.String(o, label))
223+
label := estats.CustomLabelFromContext(ctx)
224+
if label != "" {
225+
attributes = append(attributes, otelattribute.String(o, label))
226+
}
225227
}
226228
}
227229

@@ -280,8 +282,10 @@ func (h *clientMetricsHandler) processRPCEnd(ctx context.Context, ai *attemptInf
280282
if val, ok := ai.xdsLabels[o]; ok {
281283
attributes = append(attributes, otelattribute.String(o, val))
282284
} else if o == "grpc.client.call.custom" {
283-
label, _ := grpc.CustomLabelFromContext(ctx)
284-
attributes = append(attributes, otelattribute.String(o, label))
285+
label := estats.CustomLabelFromContext(ctx)
286+
if label != "" {
287+
attributes = append(attributes, otelattribute.String(o, label))
288+
}
285289
}
286290
}
287291

‎stats/opentelemetry/csm/observability_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"go.opentelemetry.io/otel/sdk/metric/metricdata"
3131
"google.golang.org/grpc"
3232
"google.golang.org/grpc/encoding/gzip"
33+
estats "google.golang.org/grpc/experimental/stats"
3334
istats "google.golang.org/grpc/internal/stats"
3435
"google.golang.org/grpc/internal/stubserver"
3536
testgrpc "google.golang.org/grpc/interop/grpc_testing"
@@ -478,7 +479,7 @@ func (s) TestXDSLabels(t *testing.T) {
478479
}
479480

480481
defer ss.Stop()
481-
ss.Client.UnaryCall(grpc.NewContextWithCustomLabel(ctx, "my-custom-label"), &testpb.SimpleRequest{Payload: &testpb.Payload{
482+
ss.Client.UnaryCall(estats.NewContextWithCustomLabel(ctx, "my-custom-label"), &testpb.SimpleRequest{Payload: &testpb.Payload{
482483
Body: make([]byte, 10000),
483484
}}, grpc.UseCompressor(gzip.Name))
484485

0 commit comments

Comments
 (0)