context.go

 1// Copyright The OpenTelemetry Authors
 2// SPDX-License-Identifier: Apache-2.0
 3
 4package baggage // import "go.opentelemetry.io/otel/internal/baggage"
 5
 6import "context"
 7
 8type baggageContextKeyType int
 9
10const baggageKey baggageContextKeyType = iota
11
12// SetHookFunc is a callback called when storing baggage in the context.
13type SetHookFunc func(context.Context, List) context.Context
14
15// GetHookFunc is a callback called when getting baggage from the context.
16type GetHookFunc func(context.Context, List) List
17
18type baggageState struct {
19	list List
20
21	setHook SetHookFunc
22	getHook GetHookFunc
23}
24
25// ContextWithSetHook returns a copy of parent with hook configured to be
26// invoked every time ContextWithBaggage is called.
27//
28// Passing nil SetHookFunc creates a context with no set hook to call.
29func ContextWithSetHook(parent context.Context, hook SetHookFunc) context.Context {
30	var s baggageState
31	if v, ok := parent.Value(baggageKey).(baggageState); ok {
32		s = v
33	}
34
35	s.setHook = hook
36	return context.WithValue(parent, baggageKey, s)
37}
38
39// ContextWithGetHook returns a copy of parent with hook configured to be
40// invoked every time FromContext is called.
41//
42// Passing nil GetHookFunc creates a context with no get hook to call.
43func ContextWithGetHook(parent context.Context, hook GetHookFunc) context.Context {
44	var s baggageState
45	if v, ok := parent.Value(baggageKey).(baggageState); ok {
46		s = v
47	}
48
49	s.getHook = hook
50	return context.WithValue(parent, baggageKey, s)
51}
52
53// ContextWithList returns a copy of parent with baggage. Passing nil list
54// returns a context without any baggage.
55func ContextWithList(parent context.Context, list List) context.Context {
56	var s baggageState
57	if v, ok := parent.Value(baggageKey).(baggageState); ok {
58		s = v
59	}
60
61	s.list = list
62	ctx := context.WithValue(parent, baggageKey, s)
63	if s.setHook != nil {
64		ctx = s.setHook(ctx, list)
65	}
66
67	return ctx
68}
69
70// ListFromContext returns the baggage contained in ctx.
71func ListFromContext(ctx context.Context) List {
72	switch v := ctx.Value(baggageKey).(type) {
73	case baggageState:
74		if v.getHook != nil {
75			return v.getHook(ctx, v.list)
76		}
77		return v.list
78	default:
79		return nil
80	}
81}