1 package context
2 3 import "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
16 type valueOnlyContext struct {
17 context.Context
18 19 preserveExpiredValues bool
20 valuesCtx context.Context
21 }
22 23 var _ 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.
28 func (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.
52 func WithSuppressCancel(ctx context.Context) context.Context {
53 return &valueOnlyContext{
54 Context: context.Background(),
55 valuesCtx: ctx,
56 57 preserveExpiredValues: GetPreserveExpiredValues(ctx),
58 }
59 }
60 61 type 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
69 func 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.
75 func GetPreserveExpiredValues(ctx context.Context) bool {
76 v := ctx.Value(preserveExpiredValuesKey{})
77 if v != nil {
78 return v.(bool)
79 }
80 return false
81 }
82