diff --git a/metrics/context.go b/metrics/context.go index 5bb349aa..b5529e97 100644 --- a/metrics/context.go +++ b/metrics/context.go @@ -22,10 +22,22 @@ func WithRegistry(ctx context.Context, registry Registry) context.Context { } func FromContext(ctx context.Context) Registry { - if registry, ok := ctx.Value(registryKey).(Registry); ok { + registry, ok := ctx.Value(registryKey).(Registry) + if !ok { + registry = DefaultMetricsRegistry + } + rootRegistry, ok := registry.(*rootRegistry) + if !ok { + return registry + } + tagsContainer, ok := ctx.Value(tagsKey).(*tagsContainer) + if !ok { return registry } - return DefaultMetricsRegistry + return &childRegistry{ + root: rootRegistry, + tags: tagsContainer.Tags, + } } // AddTags adds the provided tags to the provided context. If the provided context already has tag storage, the new tags diff --git a/metrics/context_test.go b/metrics/context_test.go index 79a236ed..ea1d0cb6 100644 --- a/metrics/context_test.go +++ b/metrics/context_test.go @@ -23,6 +23,21 @@ func TestFromContext(t *testing.T) { assert.Equal(t, FromContext(ctx), reg) } +func TestFromContextWithTags(t *testing.T) { + ctx := context.Background() + reg := &rootRegistry{ + registry: metrics.NewPrefixedRegistry("foo"), + } + + ctx = WithRegistry(ctx, reg) + ctx = AddTags(ctx, MustNewTag("bar", "baz")) + + assert.Equal(t, FromContext(ctx), &childRegistry{ + root: reg, + tags: Tags{MustNewTag("bar", "baz")}, + }) +} + func TestDefaultFromContext(t *testing.T) { ctx := context.Background() assert.Equal(t, FromContext(ctx), DefaultMetricsRegistry)