context.go 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. package utils
  2. import (
  3. "context"
  4. "adm-ods/common.in/span"
  5. beegoctx "github.com/astaxie/beego/context"
  6. "github.com/smallnest/rpcx/share"
  7. )
  8. func NewContextFromBeegoCtx(ctx *beegoctx.Context) context.Context {
  9. // request metadata
  10. metadata := make(map[string]string)
  11. if spanID, err := span.ParseSpanID(ctx.Input.Header("span")); err == nil {
  12. metadata["span"] = span.NewSpanID(*spanID).String()
  13. }
  14. return context.WithValue(context.Background(), share.ReqMetaDataKey, metadata)
  15. }
  16. func NewContextFromRpcCtx(ctx context.Context) context.Context {
  17. metadata := make(map[string]string)
  18. metadata["span"] = span.NewRootSpanID().String() // default new a root spanid
  19. if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
  20. if mapMD := reqMetadata.(map[string]string); mapMD != nil {
  21. if v, ok := mapMD["span"]; ok {
  22. if spanID, err := span.ParseSpanID(v); err == nil {
  23. metadata["span"] = span.NewSpanID(*spanID).String()
  24. }
  25. }
  26. }
  27. }
  28. return context.WithValue(context.Background(), share.ReqMetaDataKey, metadata)
  29. }
  30. func GetSpanFromRpcCtx(ctx context.Context) (spanID span.SpanID) {
  31. if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
  32. if mapMD := reqMetadata.(map[string]string); mapMD != nil {
  33. if v, ok := mapMD["span"]; ok {
  34. if s, err := span.ParseSpanID(v); err == nil {
  35. spanID = *s
  36. }
  37. }
  38. }
  39. }
  40. return
  41. }
  42. func GetMethodFromRpcCtx(ctx context.Context) (method string) {
  43. if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
  44. if mapMD := reqMetadata.(map[string]string); mapMD != nil {
  45. if v, ok := mapMD["method"]; ok {
  46. method = v
  47. }
  48. }
  49. }
  50. return
  51. }
  52. func SetMethodToRpcCtx(ctx context.Context, method string) {
  53. if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
  54. if mapMD := reqMetadata.(map[string]string); mapMD != nil {
  55. mapMD["method"] = method
  56. }
  57. }
  58. }