-
Notifications
You must be signed in to change notification settings - Fork 34
/
context.go
73 lines (63 loc) · 2.01 KB
/
context.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
// Copyright (c) 2018 Palantir Technologies. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package metrics
import (
"context"
)
type mContextKey string
const (
registryKey = mContextKey("metrics-registry")
tagsKey = mContextKey("metrics-tags")
)
var DefaultMetricsRegistry = NewRootMetricsRegistry()
func WithRegistry(ctx context.Context, registry Registry) context.Context {
return context.WithValue(ctx, registryKey, registry)
}
func FromContext(ctx context.Context) Registry {
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 &childRegistry{
root: rootRegistry,
tags: tagsContainer.Tags,
}
}
// AddTags adds the provided tags to the provided context. If no tags are provided, the context is returned unchanged.
// Otherwise, a new context is returned with the new tags appended to any tags stored on the parent context.
// This function does not perform any de-duplication (that is, if a tag in the provided tags has the
// same key as an existing one, it will still be appended).
func AddTags(ctx context.Context, tags ...Tag) context.Context {
if len(tags) == 0 {
return ctx
}
container, ok := ctx.Value(tagsKey).(*tagsContainer)
if !ok || container == nil {
container = &tagsContainer{}
}
newTags := make(Tags, len(container.Tags)+len(tags))
copy(newTags, container.Tags)
copy(newTags[len(container.Tags):], tags)
return context.WithValue(ctx, tagsKey, &tagsContainer{
Tags: newTags,
})
}
// TagsFromContext returns the tags stored on the provided context. May be nil if no tags have been set on the context.
func TagsFromContext(ctx context.Context) Tags {
if tagsContainer, ok := ctx.Value(tagsKey).(*tagsContainer); ok {
return tagsContainer.Tags
}
return nil
}
type tagsContainer struct {
Tags Tags
}