blob: 6b5c0c2d933558957b2e58c636b74669c5363b0c [file] [log] [blame]
// 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()
}
}