[Yarpc-Go] Implement Metrics Tag Decorator

Update

Comment address

address comment

resolve lint error

update copyright

changing interface location

gRPC reflection server (#2371)

* gRPC reflection server

* license header

* add README.md for gRPC reflection

* updated .codecov.yml

[Yarpc-Go] Implement Metrics Tag Decorator

Update

Comment address

address comment

resolve lint error

update copyright

changing interface location

gRPC reflection server (#2371)

* gRPC reflection server

* license header

* add README.md for gRPC reflection

* updated .codecov.yml

nit

Revert "gRPC reflection server (#2371)"

This reverts commit 056f218.

nit

Author: sakshi4016

api/metrics/metricstagdecorators/metricstagdecorators.go

@@ -0,0 +1,30 @@
+// Copyright (c) 2025 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package metricstagdecorators
+
+// DecoratorProperties is used for passing request-specific properties
+// to metrics decorators.
+type DecoratorProperties struct{}
+
+// MetricsTagsDecorators is used for adding custom tags to YARPC metrics.
+type MetricsTagsDecorators interface {
+	ProvideTags(properties DecoratorProperties) map[string]string
+}

config.go

@@ -28,6 +28,7 @@ import (
 	"github.com/uber-go/tally"
 	"go.uber.org/net/metrics"
 	"go.uber.org/net/metrics/tallypush"
+	"go.uber.org/yarpc/api/metrics/metricstagdecorators"
 	"go.uber.org/yarpc/api/middleware"
 	"go.uber.org/yarpc/internal/observability"
 	"go.uber.org/zap"
@@ -157,6 +158,8 @@ type MetricsConfig struct {
 	// TagsBlocklist enlists tags' keys that should be suppressed from all the metrics
 	// emitted from w/in YARPC middleware.
 	TagsBlocklist []string
+	// MetricTagDecorators populates yarpc metrics scope with custom tags.
+	MetricTagsDecorators []metricstagdecorators.MetricsTagsDecorators
 }
 
 func (c MetricsConfig) scope(name string, logger *zap.Logger) (*metrics.Scope, context.CancelFunc) {

dispatcher.go

@@ -103,10 +103,11 @@ func addObservingMiddleware(cfg Config, meter *metrics.Scope, logger *zap.Logger
 	}
 
 	observer := observability.NewMiddleware(observability.Config{
-		Logger:              logger,
-		Scope:               meter,
-		ContextExtractor:    extractor,
-		MetricTagsBlocklist: cfg.Metrics.TagsBlocklist,
+		Logger:               logger,
+		Scope:                meter,
+		ContextExtractor:     extractor,
+		MetricTagsBlocklist:  cfg.Metrics.TagsBlocklist,
+		MetricTagsDecorators: cfg.Metrics.MetricTagsDecorators,
 		Levels: observability.LevelsConfig{
 			Default: observability.DirectionalLevelsConfig{
 				Success:          cfg.Logging.Levels.Success,

internal/observability/graph.go

@@ -27,6 +27,7 @@ import (
 
 	"go.uber.org/net/metrics"
 	"go.uber.org/net/metrics/bucket"
+	"go.uber.org/yarpc/api/metrics/metricstagdecorators"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/internal/digester"
 	"go.uber.org/zap"
@@ -68,10 +69,11 @@ const (
 // A graph represents a collection of services: each service is a node, and we
 // collect stats for each caller-callee-transport-encoding-procedure-rk-sk-rd edge.
 type graph struct {
-	meter            *metrics.Scope
-	logger           *zap.Logger
-	extract          ContextExtractor
-	ignoreMetricsTag *metricsTagIgnore
+	meter                *metrics.Scope
+	logger               *zap.Logger
+	extract              ContextExtractor
+	ignoreMetricsTag     *metricsTagIgnore
+	metricTagsDecorators []metricstagdecorators.MetricsTagsDecorators
 
 	edgesMu sync.RWMutex
 	edges   map[string]*edge
@@ -164,13 +166,14 @@ func (m *metricsTagIgnore) tags(req *transport.Request, direction string, rpcTyp
 	return tags
 }
 
-func newGraph(meter *metrics.Scope, logger *zap.Logger, extract ContextExtractor, metricTagsIgnore []string) graph {
+func newGraph(meter *metrics.Scope, logger *zap.Logger, extract ContextExtractor, metricTagsIgnore []string, metricTagsDecorators []metricstagdecorators.MetricsTagsDecorators) graph {
 	return graph{
-		edges:            make(map[string]*edge, _defaultGraphSize),
-		meter:            meter,
-		logger:           logger,
-		extract:          extract,
-		ignoreMetricsTag: newMetricsTagIgnore(metricTagsIgnore),
+		edges:                make(map[string]*edge, _defaultGraphSize),
+		meter:                meter,
+		logger:               logger,
+		extract:              extract,
+		ignoreMetricsTag:     newMetricsTagIgnore(metricTagsIgnore),
+		metricTagsDecorators: metricTagsDecorators,
 		inboundLevels: levels{
 			success:          zapcore.DebugLevel,
 			failure:          zapcore.ErrorLevel,
@@ -264,7 +267,7 @@ func (g *graph) createEdge(key []byte, req *transport.Request, direction string,
 		return e
 	}
 
-	e := newEdge(g.logger, g.meter, g.ignoreMetricsTag, req, direction, rpcType)
+	e := newEdge(g.logger, g.meter, g.ignoreMetricsTag, g.metricTagsDecorators, req, direction, rpcType)
 	g.edges[string(key)] = e
 	return e
 }
@@ -309,9 +312,23 @@ type streamEdge struct {
 
 // newEdge constructs a new edge. Since Registries enforce metric uniqueness,
 // edges should be cached and re-used for each RPC.
-func newEdge(logger *zap.Logger, meter *metrics.Scope, tagToIgnore *metricsTagIgnore, req *transport.Request, direction string, rpcType transport.Type) *edge {
+func newEdge(logger *zap.Logger, meter *metrics.Scope, tagToIgnore *metricsTagIgnore, metricTagsDecorators []metricstagdecorators.MetricsTagsDecorators, req *transport.Request, direction string, rpcType transport.Type) *edge {
 	tags := tagToIgnore.tags(req, direction, rpcType)
 
+	if len(metricTagsDecorators) > 0 {
+		decoratorProperties := metricstagdecorators.DecoratorProperties{}
+
+		for _, decorator := range metricTagsDecorators {
+			decoratorTags := decorator.ProvideTags(decoratorProperties)
+			for key, value := range decoratorTags {
+				if _, exists := tags[key]; exists {
+					logger.Warn("MetricsTagsDecorators is overwriting metric tag", zap.String("key", key), zap.String("old", tags[key]), zap.String("new", value))
+				}
+				tags[key] = value
+			}
+		}
+	}
+
 	// metrics for all RPCs
 	calls, err := meter.Counter(metrics.Spec{
 		Name:      "calls",

internal/observability/graph_test.go

@@ -26,12 +26,26 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/net/metrics"
+	"go.uber.org/yarpc/api/metrics/metricstagdecorators"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/api/transport/transporttest"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest/observer"
 )
 
+const (
+	testMetricsTagsDecoratorKey   = "test_tag"
+	testMetricsTagsDecoratorValue = "test_value"
+)
+
+type TestMetricsTagsDecorator struct{}
+
+func (d *TestMetricsTagsDecorator) ProvideTags(request metricstagdecorators.DecoratorProperties) map[string]string {
+	return map[string]string{
+		testMetricsTagsDecoratorKey: testMetricsTagsDecoratorValue,
+	}
+}
+
 func TestHandleWithReservedField(t *testing.T) {
 	root := metrics.New()
 	meter := root.Scope()
@@ -148,6 +162,39 @@ func TestMetricsTagIgnore(t *testing.T) {
 	}
 }
 
+func TestPopulatingMetricsTagsDecorators(t *testing.T) {
+	root := metrics.New()
+	meter := root.Scope()
+	req := &transport.Request{
+		Caller:          "caller",
+		Service:         "service",
+		Transport:       "",
+		Encoding:        "proto",
+		Procedure:       "procedure",
+		RoutingKey:      "rk",
+		RoutingDelegate: "rd",
+	}
+
+	decorators := []metricstagdecorators.MetricsTagsDecorators{&TestMetricsTagsDecorator{}}
+
+	_ = newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, decorators, req, string(_directionOutbound), transport.Unary)
+
+	for _, counter := range root.Snapshot().Counters {
+		assert.Equal(t, counter.Tags[testMetricsTagsDecoratorKey], testMetricsTagsDecoratorValue,
+			"Expected testMetricsTagsDecorator tag to populate in %s metrics tags", counter.Name)
+	}
+
+	for _, gauge := range root.Snapshot().Gauges {
+		assert.Equal(t, gauge.Tags[testMetricsTagsDecoratorKey], testMetricsTagsDecoratorValue,
+			"Expected testMetricsTagsDecorator tag to populate in %s metrics tags", gauge.Name)
+	}
+
+	for _, histogram := range root.Snapshot().Histograms {
+		assert.Equal(t, histogram.Tags[testMetricsTagsDecoratorKey], testMetricsTagsDecoratorValue,
+			"Expected testMetricsTagsDecorator tag to populate in %s metrics tags", histogram.Name)
+	}
+}
+
 func TestEdgeNopFallbacks(t *testing.T) {
 	// If we fail to create any of the metrics required for the edge, we should
 	// fall back to no-op implementations. The easiest way to trigger failures
@@ -166,11 +213,11 @@ func TestEdgeNopFallbacks(t *testing.T) {
 	}
 
 	// Should succeed, covered by middleware tests.
-	_ = newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, req, string(_directionOutbound), transport.Unary)
+	_ = newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, nil, req, string(_directionOutbound), transport.Unary)
 
 	// Should fall back to no-op metrics.
 	// Usage of nil metrics should not panic, should not observe changes.
-	e := newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, req, string(_directionOutbound), transport.Unary)
+	e := newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, nil, req, string(_directionOutbound), transport.Unary)
 
 	e.calls.Inc()
 	assert.Equal(t, int64(0), e.calls.Load(), "Expected to fall back to no-op metrics.")

internal/observability/middleware.go

@@ -25,6 +25,7 @@ import (
 	"sync"
 
 	"go.uber.org/net/metrics"
+	"go.uber.org/yarpc/api/metrics/metricstagdecorators"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/yarpcerrors"
 	"go.uber.org/zap"
@@ -99,6 +100,9 @@ type Config struct {
 
 	// Levels specify log levels for various classes of requests.
 	Levels LevelsConfig
+
+	// MetricsTagDecorators populates yarpc metrics scope with custom tags.
+	MetricTagsDecorators []metricstagdecorators.MetricsTagsDecorators
 }
 
 // LevelsConfig specifies log level overrides for inbound traffic, outbound
@@ -147,7 +151,7 @@ type DirectionalLevelsConfig struct {
 // NewMiddleware constructs an observability middleware with the provided
 // configuration.
 func NewMiddleware(cfg Config) *Middleware {
-	m := &Middleware{newGraph(cfg.Scope, cfg.Logger, cfg.ContextExtractor, cfg.MetricTagsBlocklist)}
+	m := &Middleware{newGraph(cfg.Scope, cfg.Logger, cfg.ContextExtractor, cfg.MetricTagsBlocklist, cfg.MetricTagsDecorators)}
 
 	// Apply the default levels
 	applyLogLevelsConfig(&m.graph.inboundLevels, &cfg.Levels.Default)