src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

suppress_expired.go (2866B)

      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 }