diff options
Diffstat (limited to 'vendor/go.opentelemetry.io/otel/internal/baggage/baggage.go')
-rw-r--r-- | vendor/go.opentelemetry.io/otel/internal/baggage/baggage.go | 338 |
1 files changed, 338 insertions, 0 deletions
diff --git a/vendor/go.opentelemetry.io/otel/internal/baggage/baggage.go b/vendor/go.opentelemetry.io/otel/internal/baggage/baggage.go new file mode 100644 index 0000000000..6b5c0c2d93 --- /dev/null +++ b/vendor/go.opentelemetry.io/otel/internal/baggage/baggage.go @@ -0,0 +1,338 @@ +// Copyright The OpenTelemetry Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package baggage provides types and functions to manage W3C Baggage. +package baggage + +import ( + "context" + + "go.opentelemetry.io/otel/label" +) + +type rawMap map[label.Key]label.Value +type keySet map[label.Key]struct{} + +// Map is an immutable storage for correlations. +type Map struct { + m rawMap +} + +// MapUpdate contains information about correlation changes to be +// made. +type MapUpdate struct { + // DropSingleK contains a single key to be dropped from + // correlations. Use this to avoid an overhead of a slice + // allocation if there is only one key to drop. + DropSingleK label.Key + // DropMultiK contains all the keys to be dropped from + // correlations. + DropMultiK []label.Key + + // SingleKV contains a single key-value pair to be added to + // correlations. Use this to avoid an overhead of a slice + // allocation if there is only one key-value pair to add. + SingleKV label.KeyValue + // MultiKV contains all the key-value pairs to be added to + // correlations. + MultiKV []label.KeyValue +} + +func newMap(raw rawMap) Map { + return Map{ + m: raw, + } +} + +// NewEmptyMap creates an empty correlations map. +func NewEmptyMap() Map { + return newMap(nil) +} + +// NewMap creates a map with the contents of the update applied. In +// this function, having an update with DropSingleK or DropMultiK +// makes no sense - those fields are effectively ignored. +func NewMap(update MapUpdate) Map { + return NewEmptyMap().Apply(update) +} + +// Apply creates a copy of the map with the contents of the update +// applied. Apply will first drop the keys from DropSingleK and +// DropMultiK, then add key-value pairs from SingleKV and MultiKV. +func (m Map) Apply(update MapUpdate) Map { + delSet, addSet := getModificationSets(update) + mapSize := getNewMapSize(m.m, delSet, addSet) + + r := make(rawMap, mapSize) + for k, v := range m.m { + // do not copy items we want to drop + if _, ok := delSet[k]; ok { + continue + } + // do not copy items we would overwrite + if _, ok := addSet[k]; ok { + continue + } + r[k] = v + } + if update.SingleKV.Key.Defined() { + r[update.SingleKV.Key] = update.SingleKV.Value + } + for _, kv := range update.MultiKV { + r[kv.Key] = kv.Value + } + if len(r) == 0 { + r = nil + } + return newMap(r) +} + +func getModificationSets(update MapUpdate) (delSet, addSet keySet) { + deletionsCount := len(update.DropMultiK) + if update.DropSingleK.Defined() { + deletionsCount++ + } + if deletionsCount > 0 { + delSet = make(map[label.Key]struct{}, deletionsCount) + for _, k := range update.DropMultiK { + delSet[k] = struct{}{} + } + if update.DropSingleK.Defined() { + delSet[update.DropSingleK] = struct{}{} + } + } + + additionsCount := len(update.MultiKV) + if update.SingleKV.Key.Defined() { + additionsCount++ + } + if additionsCount > 0 { + addSet = make(map[label.Key]struct{}, additionsCount) + for _, k := range update.MultiKV { + addSet[k.Key] = struct{}{} + } + if update.SingleKV.Key.Defined() { + addSet[update.SingleKV.Key] = struct{}{} + } + } + + return +} + +func getNewMapSize(m rawMap, delSet, addSet keySet) int { + mapSizeDiff := 0 + for k := range addSet { + if _, ok := m[k]; !ok { + mapSizeDiff++ + } + } + for k := range delSet { + if _, ok := m[k]; ok { + if _, inAddSet := addSet[k]; !inAddSet { + mapSizeDiff-- + } + } + } + return len(m) + mapSizeDiff +} + +// Value gets a value from correlations map and returns a boolean +// value indicating whether the key exist in the map. +func (m Map) Value(k label.Key) (label.Value, bool) { + value, ok := m.m[k] + return value, ok +} + +// HasValue returns a boolean value indicating whether the key exist +// in the map. +func (m Map) HasValue(k label.Key) bool { + _, has := m.Value(k) + return has +} + +// Len returns a length of the map. +func (m Map) Len() int { + return len(m.m) +} + +// Foreach calls a passed callback once on each key-value pair until +// all the key-value pairs of the map were iterated or the callback +// returns false, whichever happens first. +func (m Map) Foreach(f func(label.KeyValue) bool) { + for k, v := range m.m { + if !f(label.KeyValue{ + Key: k, + Value: v, + }) { + return + } + } +} + +type correlationsType struct{} + +// SetHookFunc describes a type of a callback that is called when +// storing baggage in the context. +type SetHookFunc func(context.Context) context.Context + +// GetHookFunc describes a type of a callback that is called when +// getting baggage from the context. +type GetHookFunc func(context.Context, Map) Map + +// value under this key is either of type Map or correlationsData +var correlationsKey = &correlationsType{} + +type correlationsData struct { + m Map + setHook SetHookFunc + getHook GetHookFunc +} + +func (d correlationsData) isHookless() bool { + return d.setHook == nil && d.getHook == nil +} + +type hookKind int + +const ( + hookKindSet hookKind = iota + hookKindGet +) + +func (d *correlationsData) overrideHook(kind hookKind, setHook SetHookFunc, getHook GetHookFunc) { + switch kind { + case hookKindSet: + d.setHook = setHook + case hookKindGet: + d.getHook = getHook + } +} + +// ContextWithSetHook installs a hook function that will be invoked +// every time ContextWithMap is called. To avoid unnecessary callback +// invocations (recursive or not), the callback can temporarily clear +// the hooks from the context with the ContextWithNoHooks function. +// +// Note that NewContext also calls ContextWithMap, so the hook will be +// invoked. +// +// Passing nil SetHookFunc creates a context with no set hook to call. +// +// This function should not be used by applications or libraries. It +// is mostly for interoperation with other observability APIs. +func ContextWithSetHook(ctx context.Context, hook SetHookFunc) context.Context { + return contextWithHook(ctx, hookKindSet, hook, nil) +} + +// ContextWithGetHook installs a hook function that will be invoked +// every time MapFromContext is called. To avoid unnecessary callback +// invocations (recursive or not), the callback can temporarily clear +// the hooks from the context with the ContextWithNoHooks function. +// +// Note that NewContext also calls MapFromContext, so the hook will be +// invoked. +// +// Passing nil GetHookFunc creates a context with no get hook to call. +// +// This function should not be used by applications or libraries. It +// is mostly for interoperation with other observability APIs. +func ContextWithGetHook(ctx context.Context, hook GetHookFunc) context.Context { + return contextWithHook(ctx, hookKindGet, nil, hook) +} + +func contextWithHook(ctx context.Context, kind hookKind, setHook SetHookFunc, getHook GetHookFunc) context.Context { + switch v := ctx.Value(correlationsKey).(type) { + case correlationsData: + v.overrideHook(kind, setHook, getHook) + if v.isHookless() { + return context.WithValue(ctx, correlationsKey, v.m) + } + return context.WithValue(ctx, correlationsKey, v) + case Map: + return contextWithOneHookAndMap(ctx, kind, setHook, getHook, v) + default: + m := NewEmptyMap() + return contextWithOneHookAndMap(ctx, kind, setHook, getHook, m) + } +} + +func contextWithOneHookAndMap(ctx context.Context, kind hookKind, setHook SetHookFunc, getHook GetHookFunc, m Map) context.Context { + d := correlationsData{m: m} + d.overrideHook(kind, setHook, getHook) + if d.isHookless() { + return ctx + } + return context.WithValue(ctx, correlationsKey, d) +} + +// ContextWithNoHooks creates a context with all the hooks +// disabled. Also returns old set and get hooks. This function can be +// used to temporarily clear the context from hooks and then reinstate +// them by calling ContextWithSetHook and ContextWithGetHook functions +// passing the hooks returned by this function. +// +// This function should not be used by applications or libraries. It +// is mostly for interoperation with other observability APIs. +func ContextWithNoHooks(ctx context.Context) (context.Context, SetHookFunc, GetHookFunc) { + switch v := ctx.Value(correlationsKey).(type) { + case correlationsData: + return context.WithValue(ctx, correlationsKey, v.m), v.setHook, v.getHook + default: + return ctx, nil, nil + } +} + +// ContextWithMap returns a context with the Map entered into it. +func ContextWithMap(ctx context.Context, m Map) context.Context { + switch v := ctx.Value(correlationsKey).(type) { + case correlationsData: + v.m = m + ctx = context.WithValue(ctx, correlationsKey, v) + if v.setHook != nil { + ctx = v.setHook(ctx) + } + return ctx + default: + return context.WithValue(ctx, correlationsKey, m) + } +} + +// ContextWithNoCorrelationData returns a context stripped of correlation +// data. +func ContextWithNoCorrelationData(ctx context.Context) context.Context { + return context.WithValue(ctx, correlationsKey, nil) +} + +// NewContext returns a context with the map from passed context +// updated with the passed key-value pairs. +func NewContext(ctx context.Context, keyvalues ...label.KeyValue) context.Context { + return ContextWithMap(ctx, MapFromContext(ctx).Apply(MapUpdate{ + MultiKV: keyvalues, + })) +} + +// MapFromContext gets the current Map from a Context. +func MapFromContext(ctx context.Context) Map { + switch v := ctx.Value(correlationsKey).(type) { + case correlationsData: + if v.getHook != nil { + return v.getHook(ctx, v.m) + } + return v.m + case Map: + return v + default: + return NewEmptyMap() + } +} |