[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

limiting logging per edge

Author: sakshi4016

api/metrics/metricstagdecorator/metricstagdecorator.go

@@ -0,0 +1,26 @@
+// 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 metricstagdecorator
+
+// MetricsTagDecorator is used for adding custom tags to YARPC metrics.
+type MetricsTagDecorator interface {
+	ProvideTags() 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/metricstagdecorator"
 	"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
+	// MetricsTagsDecorators populates yarpc metrics scope with custom tags.
+	MetricsTagsDecorators []metricstagdecorator.MetricsTagDecorator
 }
 
 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,
+		MetricsTagsDecorators: cfg.Metrics.MetricsTagsDecorators,
 		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/metricstagdecorator"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/internal/digester"
 	"go.uber.org/zap"
@@ -72,6 +73,7 @@ type graph struct {
 	logger           *zap.Logger
 	extract          ContextExtractor
 	ignoreMetricsTag *metricsTagIgnore
+	decoratorTags    metrics.Tags
 
 	edgesMu sync.RWMutex
 	edges   map[string]*edge
@@ -164,13 +166,27 @@ 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, metricsTagsDecorators []metricstagdecorator.MetricsTagDecorator) graph {
+	decoratorTags := metrics.Tags{}
+	if len(metricsTagsDecorators) > 0 {
+		for _, decorator := range metricsTagsDecorators {
+			tags := decorator.ProvideTags()
+			for key, value := range tags {
+				if oldValue, exists := decoratorTags[key]; exists {
+					logger.Warn("MetricsTagsDecorators is overwriting metric tag", zap.String("key", key), zap.String("old", oldValue), zap.String("new", value))
+				}
+				decoratorTags[key] = value
+			}
+		}
+	}
+
 	return graph{
 		edges:            make(map[string]*edge, _defaultGraphSize),
 		meter:            meter,
 		logger:           logger,
 		extract:          extract,
 		ignoreMetricsTag: newMetricsTagIgnore(metricTagsIgnore),
+		decoratorTags:    decoratorTags,
 		inboundLevels: levels{
 			success:          zapcore.DebugLevel,
 			failure:          zapcore.ErrorLevel,
@@ -264,7 +280,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.decoratorTags, req, direction, rpcType)
 	g.edges[string(key)] = e
 	return e
 }
@@ -309,9 +325,15 @@ 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, decoratorTags metrics.Tags, req *transport.Request, direction string, rpcType transport.Type) *edge {
 	tags := tagToIgnore.tags(req, direction, rpcType)
 
+	if len(decoratorTags) > 0 {
+		for key, value := range decoratorTags {
+			tags[key] = value
+		}
+	}
+
 	// metrics for all RPCs
 	calls, err := meter.Counter(metrics.Spec{
 		Name:      "calls",

internal/observability/graph_test.go

@@ -32,6 +32,19 @@ import (
 	"go.uber.org/zap/zaptest/observer"
 )
 
+const (
+	testMetricsTagsDecoratorKey   = "test_tag"
+	testMetricsTagsDecoratorValue = "test_value"
+)
+
+type TestMetricsTagsDecorator struct{}
+
+func (d *TestMetricsTagsDecorator) ProvideTags() map[string]string {
+	return map[string]string{
+		testMetricsTagsDecoratorKey: testMetricsTagsDecoratorValue,
+	}
+}
+
 func TestHandleWithReservedField(t *testing.T) {
 	root := metrics.New()
 	meter := root.Scope()
@@ -148,6 +161,41 @@ 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",
+	}
+
+	decoratorTags := metrics.Tags{
+		testMetricsTagsDecoratorKey: testMetricsTagsDecoratorValue,
+	}
+
+	_ = newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, decoratorTags, 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 +214,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/metricstagdecorator"
 	"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.
+	MetricsTagsDecorators []metricstagdecorator.MetricsTagDecorator
 }
 
 // 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.MetricsTagsDecorators)}
 
 	// Apply the default levels
 	applyLogLevelsConfig(&m.graph.inboundLevels, &cfg.Levels.Default)