1package context
 2
 3import "context"
 4
 5// valueOnlyContext provides a utility to preserve only the values of a
 6// Context. Suppressing any cancellation or deadline on that context being
 7// propagated downstream of this value.
 8//
 9// If preserveExpiredValues is false (default), and the valueCtx is canceled,
10// calls to lookup values with the Values method, will always return nil. Setting
11// preserveExpiredValues to true, will allow the valueOnlyContext to lookup
12// values in valueCtx even if valueCtx is canceled.
13//
14// Based on the Go standard libraries net/lookup.go onlyValuesCtx utility.
15// https://github.com/golang/go/blob/da2773fe3e2f6106634673a38dc3a6eb875fe7d8/src/net/lookup.go
16type valueOnlyContext struct {
17	context.Context
18
19	preserveExpiredValues bool
20	valuesCtx             context.Context
21}
22
23var _ context.Context = (*valueOnlyContext)(nil)
24
25// Value looks up the key, returning its value. If configured to not preserve
26// values of expired context, and the wrapping context is canceled, nil will be
27// returned.
28func (v *valueOnlyContext) Value(key interface{}) interface{} {
29	if !v.preserveExpiredValues {
30		select {
31		case <-v.valuesCtx.Done():
32			return nil
33		default:
34		}
35	}
36
37	return v.valuesCtx.Value(key)
38}
39
40// WithSuppressCancel wraps the Context value, suppressing its deadline and
41// cancellation events being propagated downstream to consumer of the returned
42// context.
43//
44// By default the wrapped Context's Values are available downstream until the
45// wrapped Context is canceled. Once the wrapped Context is canceled, Values
46// method called on the context return will no longer lookup any key. As they
47// are now considered expired.
48//
49// To override this behavior, use WithPreserveExpiredValues on the Context
50// before it is wrapped by WithSuppressCancel. This will make the Context
51// returned by WithSuppressCancel allow lookup of expired values.
52func WithSuppressCancel(ctx context.Context) context.Context {
53	return &valueOnlyContext{
54		Context:   context.Background(),
55		valuesCtx: ctx,
56
57		preserveExpiredValues: GetPreserveExpiredValues(ctx),
58	}
59}
60
61type preserveExpiredValuesKey struct{}
62
63// WithPreserveExpiredValues adds a Value to the Context if expired values
64// should be preserved, and looked up by a Context wrapped by
65// WithSuppressCancel.
66//
67// WithPreserveExpiredValues must be added as a value to a Context, before that
68// Context is wrapped by WithSuppressCancel
69func WithPreserveExpiredValues(ctx context.Context, enable bool) context.Context {
70	return context.WithValue(ctx, preserveExpiredValuesKey{}, enable)
71}
72
73// GetPreserveExpiredValues looks up, and returns the PreserveExpressValues
74// value in the context. Returning true if enabled, false otherwise.
75func GetPreserveExpiredValues(ctx context.Context) bool {
76	v := ctx.Value(preserveExpiredValuesKey{})
77	if v != nil {
78		return v.(bool)
79	}
80	return false
81}