[Yarpc-Go] Implement Metrics Tag Decorator

api/metrics/metricstagdecorator/metricstagdecorator.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 metricstagdecorator
+
+import (
+	"context"
+)
+
+// MetricsTagDecorator is used for adding custom tags to YARPC metrics.
+type MetricsTagDecorator interface {
+	ProvideTags(ctx context.Context) 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,15 +27,19 @@ 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/yarpc/internal/sampledlogger"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
 
 var (
 	_timeNow          = time.Now // for tests
 	_defaultGraphSize = 128
+	_logInterval      = time.Duration(5 * time.Minute)
+
 	// Latency buckets for histograms. At some point, we may want to make these
 	// configurable.
 	_bucketsMs = bucket.NewRPCLatency()
@@ -72,9 +76,10 @@ type graph struct {
 	logger           *zap.Logger
 	extract          ContextExtractor
 	ignoreMetricsTag *metricsTagIgnore
-
-	edgesMu sync.RWMutex
-	edges   map[string]*edge
+	decoratorTags    []metricstagdecorator.MetricsTagDecorator
+	sampledLogger    *sampledlogger.SampledLogger
+	edgesMu          sync.RWMutex
+	edges            map[string]*edge
 
 	inboundLevels, outboundLevels levels
 }
@@ -164,13 +169,15 @@ 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 {
 	return graph{
 		edges:            make(map[string]*edge, _defaultGraphSize),
 		meter:            meter,
 		logger:           logger,
 		extract:          extract,
 		ignoreMetricsTag: newMetricsTagIgnore(metricTagsIgnore),
+		decoratorTags:    metricsTagsDecorators,
+		sampledLogger:    sampledlogger.NewSampledLogger(_logInterval, logger),
 		inboundLevels: levels{
 			success:          zapcore.DebugLevel,
 			failure:          zapcore.ErrorLevel,
@@ -220,7 +227,7 @@ func (g *graph) begin(ctx context.Context, rpcType transport.Type, direction dir
 	if !g.ignoreMetricsTag.rpcType {
 		d.Add(rpcType.String())
 	}
-	e := g.getOrCreateEdge(d.Digest(), req, string(direction), rpcType)
+	e := g.getOrCreateEdge(ctx, d.Digest(), req, string(direction), rpcType)
 	d.Free()
 
 	levels := &g.inboundLevels
@@ -240,11 +247,11 @@ func (g *graph) begin(ctx context.Context, rpcType transport.Type, direction dir
 	}
 }
 
-func (g *graph) getOrCreateEdge(key []byte, req *transport.Request, direction string, rpcType transport.Type) *edge {
+func (g *graph) getOrCreateEdge(ctx context.Context, key []byte, req *transport.Request, direction string, rpcType transport.Type) *edge {
 	if e := g.getEdge(key); e != nil {
 		return e
 	}
-	return g.createEdge(key, req, direction, rpcType)
+	return g.createEdge(ctx, key, req, direction, rpcType)
 }
 
 func (g *graph) getEdge(key []byte) *edge {
@@ -254,7 +261,7 @@ func (g *graph) getEdge(key []byte) *edge {
 	return e
 }
 
-func (g *graph) createEdge(key []byte, req *transport.Request, direction string, rpcType transport.Type) *edge {
+func (g *graph) createEdge(ctx context.Context, key []byte, req *transport.Request, direction string, rpcType transport.Type) *edge {
 	g.edgesMu.Lock()
 	// Since we'll rarely hit this code path, the overhead of defer is acceptable.
 	defer g.edgesMu.Unlock()
@@ -264,7 +271,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(ctx, g.logger, g.sampledLogger, g.meter, g.ignoreMetricsTag, g.decoratorTags, req, direction, rpcType)
 	g.edges[string(key)] = e
 	return e
 }
@@ -309,9 +316,17 @@ 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(ctx context.Context, logger *zap.Logger, sampledLogger *sampledlogger.SampledLogger, meter *metrics.Scope, tagToIgnore *metricsTagIgnore, decoratorTags []metricstagdecorator.MetricsTagDecorator, req *transport.Request, direction string, rpcType transport.Type) *edge {
 	tags := tagToIgnore.tags(req, direction, rpcType)
 
+	// Merge custom decorator tags into the YARPC metrics base tag set.
+	customDecoratorTags := getDecoratorTags(ctx, sampledLogger, decoratorTags)
+	if customDecoratorTags != nil {
+		for key, value := range *customDecoratorTags {
+			tags[key] = value
+		}
+	}
+
 	// metrics for all RPCs
 	calls, err := meter.Counter(metrics.Spec{
 		Name:      "calls",
@@ -603,3 +618,23 @@ func unknownIfEmpty(t string) string {
 	}
 	return t
 }
+
+// getDecoratorTags collects tags from all provided MetricsTagDecorators.
+func getDecoratorTags(ctx context.Context, sampledLogger *sampledlogger.SampledLogger, metricsTagsDecorators []metricstagdecorator.MetricsTagDecorator) *metrics.Tags {
+	if len(metricsTagsDecorators) == 0 {
+		return nil
+	}
+
+	decoratorTags := metrics.Tags{}
+	for _, decorator := range metricsTagsDecorators {
+		tags := decorator.ProvideTags(ctx)
+		for key, value := range tags {
+			if oldValue, exists := decoratorTags[key]; exists {
+				sampledLogger.Warn("MetricsTagsDecorators is overwriting metric tag", zap.String("key", key), zap.String("old", oldValue), zap.String("new", value))
+			}
+			decoratorTags[key] = value
+		}
+	}
+
+	return &decoratorTags
+}

internal/observability/graph_test.go

@@ -26,12 +26,27 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/net/metrics"
+	"go.uber.org/yarpc/api/metrics/metricstagdecorator"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/api/transport/transporttest"
+	"go.uber.org/yarpc/internal/sampledlogger"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest/observer"
 )
 
+const (
+	testMetricsTagsDecoratorKey   = "test_tag"
+	testMetricsTagsDecoratorValue = "test_value"
+)
+
+type TestMetricsTagsDecorator struct{}
+
+func (d *TestMetricsTagsDecorator) ProvideTags(ctx context.Context) map[string]string {
+	return map[string]string{
+		testMetricsTagsDecoratorKey: testMetricsTagsDecoratorValue,
+	}
+}
+
 func TestHandleWithReservedField(t *testing.T) {
 	root := metrics.New()
 	meter := root.Scope()
@@ -148,6 +163,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",
+	}
+
+	// Creating new edge with MetricsTagsDecorators
+	decorator := []metricstagdecorator.MetricsTagDecorator{&TestMetricsTagsDecorator{}}
+	graph := newGraph(meter, zap.NewNop(), nil, nil, decorator)
+
+	_ = newEdge(context.Background(), graph.logger, graph.sampledLogger, graph.meter, graph.ignoreMetricsTag, graph.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 +216,11 @@ func TestEdgeNopFallbacks(t *testing.T) {
 	}
 
 	// Should succeed, covered by middleware tests.
-	_ = newEdge(zap.NewNop(), meter, &metricsTagIgnore{}, req, string(_directionOutbound), transport.Unary)
+	_ = newEdge(context.Background(), zap.NewNop(), sampledlogger.NewDefaultSampledLogger(), 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(context.Background(), zap.NewNop(), sampledlogger.NewDefaultSampledLogger(), 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)

internal/sampledlogger/sampledlogger.go

@@ -0,0 +1,89 @@
+package sampledlogger
+
+import (
+	"sync/atomic"
+	"time"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+type SampledLogger struct {
+	logger      *zap.Logger
+	logInterval time.Duration
+	lastLogTime atomic.Value
+}
+
+// NewSampledLogger creates a new SampledLogger with custom interval and zap.Logger.
+func NewSampledLogger(interval time.Duration, logger *zap.Logger) *SampledLogger {
+	return &SampledLogger{
+		logger:      logger,
+		logInterval: interval,
+	}
+}
+
+// NewDefaultSampledLogger creates a SampledLogger with zap.NewNop() and 5-minute interval.
+func NewDefaultSampledLogger() *SampledLogger {
+	return NewSampledLogger(5*time.Minute, zap.NewNop())
+}
+
+// log performs rate-limited logging for the given level and message.
+func (sl *SampledLogger) log(level zapcore.Level, msg string, fields ...zap.Field) {
+	now := time.Now()
+	last := sl.lastLogTime.Load().(time.Time)
+
+	if last.IsZero() || now.Sub(last) > sl.logInterval {
+		switch level {
+		case zapcore.DebugLevel:
+			sl.logger.Debug(msg, fields...)
+		case zapcore.InfoLevel:
+			sl.logger.Info(msg, fields...)
+		case zapcore.WarnLevel:
+			sl.logger.Warn(msg, fields...)
+		case zapcore.ErrorLevel:
+			sl.logger.Error(msg, fields...)
+		case zapcore.DPanicLevel:
+			sl.logger.DPanic(msg, fields...)
+		case zapcore.PanicLevel:
+			sl.logger.Panic(msg, fields...)
+		case zapcore.FatalLevel:
+			sl.logger.Fatal(msg, fields...)
+		}
+		sl.lastLogTime.Store(now)
+	}
+}
+
+// Debug logs a debug-level message with rate limiting.
+func (sl *SampledLogger) Debug(msg string, fields ...zap.Field) {
+	sl.log(zapcore.DebugLevel, msg, fields...)
+}
+
+// Info logs an info-level message with rate limiting.
+func (sl *SampledLogger) Info(msg string, fields ...zap.Field) {
+	sl.log(zapcore.InfoLevel, msg, fields...)
+}
+
+// Warn logs a warn-level message with rate limiting.
+func (sl *SampledLogger) Warn(msg string, fields ...zap.Field) {
+	sl.log(zapcore.WarnLevel, msg, fields...)
+}
+
+// Error logs an error-level message with rate limiting.
+func (sl *SampledLogger) Error(msg string, fields ...zap.Field) {
+	sl.log(zapcore.ErrorLevel, msg, fields...)
+}
+
+// DPanic logs a DPanic-level message with rate limiting.
+func (sl *SampledLogger) DPanic(msg string, fields ...zap.Field) {
+	sl.log(zapcore.DPanicLevel, msg, fields...)
+}
+
+// Panic logs a panic-level message with rate limiting.
+func (sl *SampledLogger) Panic(msg string, fields ...zap.Field) {
+	sl.log(zapcore.PanicLevel, msg, fields...)
+}
+
+// Fatal logs a fatal-level message with rate limiting.
+func (sl *SampledLogger) Fatal(msg string, fields ...zap.Field) {
+	sl.log(zapcore.FatalLevel, msg, fields...)
+}