context.go 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. package metrics
  2. import (
  3. "context"
  4. "encoding/json"
  5. )
  6. // AddToPropagateCtx adds a key and value that will be propagated through RPC calls
  7. func AddToPropagateCtx(ctx context.Context, key string, val interface{}) context.Context {
  8. propagate := ToMap(ctx)
  9. propagate[key] = val
  10. return context.WithValue(ctx, PropagateCtxKey, propagate)
  11. }
  12. // GetFromPropagateCtx get a value from the propagate
  13. func GetFromPropagateCtx(ctx context.Context, key string) interface{} {
  14. propagate := ToMap(ctx)
  15. if val, ok := propagate[key]; ok {
  16. return val
  17. }
  18. return nil
  19. }
  20. // ToMap returns the values that will be propagated through RPC calls in map[string]interface{} format
  21. func ToMap(ctx context.Context) map[string]interface{} {
  22. if ctx == nil {
  23. return map[string]interface{}{}
  24. }
  25. p := ctx.Value(PropagateCtxKey)
  26. if p != nil {
  27. return p.(map[string]interface{})
  28. }
  29. return map[string]interface{}{}
  30. }
  31. // FromMap creates a new context from a map with propagated values
  32. func FromMap(val map[string]interface{}) context.Context {
  33. return context.WithValue(context.Background(), PropagateCtxKey, val)
  34. }
  35. // Encode returns the given propagatable context encoded in binary format
  36. func Encode(ctx context.Context) ([]byte, error) {
  37. m := ToMap(ctx)
  38. if len(m) > 0 {
  39. return json.Marshal(m)
  40. }
  41. return nil, nil
  42. }
  43. // Decode returns a context given a binary encoded message
  44. func Decode(m []byte) (context.Context, error) {
  45. if len(m) == 0 {
  46. // TODO maybe return an error
  47. return nil, nil
  48. }
  49. mp := make(map[string]interface{})
  50. err := json.Unmarshal(m, &mp)
  51. if err != nil {
  52. return nil, err
  53. }
  54. return FromMap(mp), nil
  55. }