blob: 6b5c0c2d933558957b2e58c636b74669c5363b0c [file] [log] [blame]
Joey Armstrong5f51f2e2023-01-17 17:06:26 -05001// Copyright The OpenTelemetry Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Package baggage provides types and functions to manage W3C Baggage.
16package baggage
17
18import (
19 "context"
20
21 "go.opentelemetry.io/otel/label"
22)
23
24type rawMap map[label.Key]label.Value
25type keySet map[label.Key]struct{}
26
27// Map is an immutable storage for correlations.
28type Map struct {
29 m rawMap
30}
31
32// MapUpdate contains information about correlation changes to be
33// made.
34type MapUpdate struct {
35 // DropSingleK contains a single key to be dropped from
36 // correlations. Use this to avoid an overhead of a slice
37 // allocation if there is only one key to drop.
38 DropSingleK label.Key
39 // DropMultiK contains all the keys to be dropped from
40 // correlations.
41 DropMultiK []label.Key
42
43 // SingleKV contains a single key-value pair to be added to
44 // correlations. Use this to avoid an overhead of a slice
45 // allocation if there is only one key-value pair to add.
46 SingleKV label.KeyValue
47 // MultiKV contains all the key-value pairs to be added to
48 // correlations.
49 MultiKV []label.KeyValue
50}
51
52func newMap(raw rawMap) Map {
53 return Map{
54 m: raw,
55 }
56}
57
58// NewEmptyMap creates an empty correlations map.
59func NewEmptyMap() Map {
60 return newMap(nil)
61}
62
63// NewMap creates a map with the contents of the update applied. In
64// this function, having an update with DropSingleK or DropMultiK
65// makes no sense - those fields are effectively ignored.
66func NewMap(update MapUpdate) Map {
67 return NewEmptyMap().Apply(update)
68}
69
70// Apply creates a copy of the map with the contents of the update
71// applied. Apply will first drop the keys from DropSingleK and
72// DropMultiK, then add key-value pairs from SingleKV and MultiKV.
73func (m Map) Apply(update MapUpdate) Map {
74 delSet, addSet := getModificationSets(update)
75 mapSize := getNewMapSize(m.m, delSet, addSet)
76
77 r := make(rawMap, mapSize)
78 for k, v := range m.m {
79 // do not copy items we want to drop
80 if _, ok := delSet[k]; ok {
81 continue
82 }
83 // do not copy items we would overwrite
84 if _, ok := addSet[k]; ok {
85 continue
86 }
87 r[k] = v
88 }
89 if update.SingleKV.Key.Defined() {
90 r[update.SingleKV.Key] = update.SingleKV.Value
91 }
92 for _, kv := range update.MultiKV {
93 r[kv.Key] = kv.Value
94 }
95 if len(r) == 0 {
96 r = nil
97 }
98 return newMap(r)
99}
100
101func getModificationSets(update MapUpdate) (delSet, addSet keySet) {
102 deletionsCount := len(update.DropMultiK)
103 if update.DropSingleK.Defined() {
104 deletionsCount++
105 }
106 if deletionsCount > 0 {
107 delSet = make(map[label.Key]struct{}, deletionsCount)
108 for _, k := range update.DropMultiK {
109 delSet[k] = struct{}{}
110 }
111 if update.DropSingleK.Defined() {
112 delSet[update.DropSingleK] = struct{}{}
113 }
114 }
115
116 additionsCount := len(update.MultiKV)
117 if update.SingleKV.Key.Defined() {
118 additionsCount++
119 }
120 if additionsCount > 0 {
121 addSet = make(map[label.Key]struct{}, additionsCount)
122 for _, k := range update.MultiKV {
123 addSet[k.Key] = struct{}{}
124 }
125 if update.SingleKV.Key.Defined() {
126 addSet[update.SingleKV.Key] = struct{}{}
127 }
128 }
129
130 return
131}
132
133func getNewMapSize(m rawMap, delSet, addSet keySet) int {
134 mapSizeDiff := 0
135 for k := range addSet {
136 if _, ok := m[k]; !ok {
137 mapSizeDiff++
138 }
139 }
140 for k := range delSet {
141 if _, ok := m[k]; ok {
142 if _, inAddSet := addSet[k]; !inAddSet {
143 mapSizeDiff--
144 }
145 }
146 }
147 return len(m) + mapSizeDiff
148}
149
150// Value gets a value from correlations map and returns a boolean
151// value indicating whether the key exist in the map.
152func (m Map) Value(k label.Key) (label.Value, bool) {
153 value, ok := m.m[k]
154 return value, ok
155}
156
157// HasValue returns a boolean value indicating whether the key exist
158// in the map.
159func (m Map) HasValue(k label.Key) bool {
160 _, has := m.Value(k)
161 return has
162}
163
164// Len returns a length of the map.
165func (m Map) Len() int {
166 return len(m.m)
167}
168
169// Foreach calls a passed callback once on each key-value pair until
170// all the key-value pairs of the map were iterated or the callback
171// returns false, whichever happens first.
172func (m Map) Foreach(f func(label.KeyValue) bool) {
173 for k, v := range m.m {
174 if !f(label.KeyValue{
175 Key: k,
176 Value: v,
177 }) {
178 return
179 }
180 }
181}
182
183type correlationsType struct{}
184
185// SetHookFunc describes a type of a callback that is called when
186// storing baggage in the context.
187type SetHookFunc func(context.Context) context.Context
188
189// GetHookFunc describes a type of a callback that is called when
190// getting baggage from the context.
191type GetHookFunc func(context.Context, Map) Map
192
193// value under this key is either of type Map or correlationsData
194var correlationsKey = &correlationsType{}
195
196type correlationsData struct {
197 m Map
198 setHook SetHookFunc
199 getHook GetHookFunc
200}
201
202func (d correlationsData) isHookless() bool {
203 return d.setHook == nil && d.getHook == nil
204}
205
206type hookKind int
207
208const (
209 hookKindSet hookKind = iota
210 hookKindGet
211)
212
213func (d *correlationsData) overrideHook(kind hookKind, setHook SetHookFunc, getHook GetHookFunc) {
214 switch kind {
215 case hookKindSet:
216 d.setHook = setHook
217 case hookKindGet:
218 d.getHook = getHook
219 }
220}
221
222// ContextWithSetHook installs a hook function that will be invoked
223// every time ContextWithMap is called. To avoid unnecessary callback
224// invocations (recursive or not), the callback can temporarily clear
225// the hooks from the context with the ContextWithNoHooks function.
226//
227// Note that NewContext also calls ContextWithMap, so the hook will be
228// invoked.
229//
230// Passing nil SetHookFunc creates a context with no set hook to call.
231//
232// This function should not be used by applications or libraries. It
233// is mostly for interoperation with other observability APIs.
234func ContextWithSetHook(ctx context.Context, hook SetHookFunc) context.Context {
235 return contextWithHook(ctx, hookKindSet, hook, nil)
236}
237
238// ContextWithGetHook installs a hook function that will be invoked
239// every time MapFromContext is called. To avoid unnecessary callback
240// invocations (recursive or not), the callback can temporarily clear
241// the hooks from the context with the ContextWithNoHooks function.
242//
243// Note that NewContext also calls MapFromContext, so the hook will be
244// invoked.
245//
246// Passing nil GetHookFunc creates a context with no get hook to call.
247//
248// This function should not be used by applications or libraries. It
249// is mostly for interoperation with other observability APIs.
250func ContextWithGetHook(ctx context.Context, hook GetHookFunc) context.Context {
251 return contextWithHook(ctx, hookKindGet, nil, hook)
252}
253
254func contextWithHook(ctx context.Context, kind hookKind, setHook SetHookFunc, getHook GetHookFunc) context.Context {
255 switch v := ctx.Value(correlationsKey).(type) {
256 case correlationsData:
257 v.overrideHook(kind, setHook, getHook)
258 if v.isHookless() {
259 return context.WithValue(ctx, correlationsKey, v.m)
260 }
261 return context.WithValue(ctx, correlationsKey, v)
262 case Map:
263 return contextWithOneHookAndMap(ctx, kind, setHook, getHook, v)
264 default:
265 m := NewEmptyMap()
266 return contextWithOneHookAndMap(ctx, kind, setHook, getHook, m)
267 }
268}
269
270func contextWithOneHookAndMap(ctx context.Context, kind hookKind, setHook SetHookFunc, getHook GetHookFunc, m Map) context.Context {
271 d := correlationsData{m: m}
272 d.overrideHook(kind, setHook, getHook)
273 if d.isHookless() {
274 return ctx
275 }
276 return context.WithValue(ctx, correlationsKey, d)
277}
278
279// ContextWithNoHooks creates a context with all the hooks
280// disabled. Also returns old set and get hooks. This function can be
281// used to temporarily clear the context from hooks and then reinstate
282// them by calling ContextWithSetHook and ContextWithGetHook functions
283// passing the hooks returned by this function.
284//
285// This function should not be used by applications or libraries. It
286// is mostly for interoperation with other observability APIs.
287func ContextWithNoHooks(ctx context.Context) (context.Context, SetHookFunc, GetHookFunc) {
288 switch v := ctx.Value(correlationsKey).(type) {
289 case correlationsData:
290 return context.WithValue(ctx, correlationsKey, v.m), v.setHook, v.getHook
291 default:
292 return ctx, nil, nil
293 }
294}
295
296// ContextWithMap returns a context with the Map entered into it.
297func ContextWithMap(ctx context.Context, m Map) context.Context {
298 switch v := ctx.Value(correlationsKey).(type) {
299 case correlationsData:
300 v.m = m
301 ctx = context.WithValue(ctx, correlationsKey, v)
302 if v.setHook != nil {
303 ctx = v.setHook(ctx)
304 }
305 return ctx
306 default:
307 return context.WithValue(ctx, correlationsKey, m)
308 }
309}
310
311// ContextWithNoCorrelationData returns a context stripped of correlation
312// data.
313func ContextWithNoCorrelationData(ctx context.Context) context.Context {
314 return context.WithValue(ctx, correlationsKey, nil)
315}
316
317// NewContext returns a context with the map from passed context
318// updated with the passed key-value pairs.
319func NewContext(ctx context.Context, keyvalues ...label.KeyValue) context.Context {
320 return ContextWithMap(ctx, MapFromContext(ctx).Apply(MapUpdate{
321 MultiKV: keyvalues,
322 }))
323}
324
325// MapFromContext gets the current Map from a Context.
326func MapFromContext(ctx context.Context) Map {
327 switch v := ctx.Value(correlationsKey).(type) {
328 case correlationsData:
329 if v.getHook != nil {
330 return v.getHook(ctx, v.m)
331 }
332 return v.m
333 case Map:
334 return v
335 default:
336 return NewEmptyMap()
337 }
338}