1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162 |
- package metrics
- import (
- "context"
- "encoding/json"
- )
- // AddToPropagateCtx adds a key and value that will be propagated through RPC calls
- func AddToPropagateCtx(ctx context.Context, key string, val interface{}) context.Context {
- propagate := ToMap(ctx)
- propagate[key] = val
- return context.WithValue(ctx, PropagateCtxKey, propagate)
- }
- // GetFromPropagateCtx get a value from the propagate
- func GetFromPropagateCtx(ctx context.Context, key string) interface{} {
- propagate := ToMap(ctx)
- if val, ok := propagate[key]; ok {
- return val
- }
- return nil
- }
- // ToMap returns the values that will be propagated through RPC calls in map[string]interface{} format
- func ToMap(ctx context.Context) map[string]interface{} {
- if ctx == nil {
- return map[string]interface{}{}
- }
- p := ctx.Value(PropagateCtxKey)
- if p != nil {
- return p.(map[string]interface{})
- }
- return map[string]interface{}{}
- }
- // FromMap creates a new context from a map with propagated values
- func FromMap(val map[string]interface{}) context.Context {
- return context.WithValue(context.Background(), PropagateCtxKey, val)
- }
- // Encode returns the given propagatable context encoded in binary format
- func Encode(ctx context.Context) ([]byte, error) {
- m := ToMap(ctx)
- if len(m) > 0 {
- return json.Marshal(m)
- }
- return nil, nil
- }
- // Decode returns a context given a binary encoded message
- func Decode(m []byte) (context.Context, error) {
- if len(m) == 0 {
- // TODO maybe return an error
- return nil, nil
- }
- mp := make(map[string]interface{})
- err := json.Unmarshal(m, &mp)
- if err != nil {
- return nil, err
- }
- return FromMap(mp), nil
- }
|