git.zd.zone/hrpc/hrpc@v0.0.12/codec/message.go (about)

     1  package codec
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"time"
     7  
     8  	"git.zd.zone/hrpc/hrpc/utils/hash"
     9  	"google.golang.org/grpc/metadata"
    10  )
    11  
    12  type MDKey string
    13  
    14  func (m MDKey) String() string {
    15  	return string(m)
    16  }
    17  
    18  const (
    19  	MDMessage MDKey = "HRPC_MD_MESSAGE"
    20  )
    21  
    22  type MSG interface {
    23  	Context() context.Context
    24  
    25  	TraceID() string
    26  	WithTraceID(s string)
    27  
    28  	RequestTimeout() time.Duration
    29  	WithRequestTimeout(t time.Duration)
    30  
    31  	Namespace() string
    32  	WithNamespace(n string)
    33  
    34  	ServerName() string
    35  	WithServerName(n string)
    36  
    37  	Client() ClientInfo
    38  	WithClient(ip, ua string)
    39  
    40  	Metadata() metadata.MD
    41  }
    42  
    43  type message struct {
    44  	context    context.Context `json:"-"`
    45  	ServerN    string          `json:"server_name"`
    46  	TID        string          `json:"trace_id"`
    47  	RTimeout   time.Duration   `json:"request_timeout"`
    48  	NameSpace  string          `json:"namespace"`
    49  	ClientInfo ClientInfo      `json:"client_info"`
    50  }
    51  
    52  type ClientInfo struct {
    53  	IP string `json:"ip"`
    54  	UA string `json:"ua"`
    55  }
    56  
    57  func (m *message) Context() context.Context {
    58  	return m.context
    59  }
    60  
    61  func (m *message) Client() ClientInfo {
    62  	return m.ClientInfo
    63  }
    64  
    65  func (m *message) WithClient(ip, ua string) {
    66  	m.ClientInfo.IP = ip
    67  	m.ClientInfo.UA = ua
    68  }
    69  
    70  func (m *message) TraceID() string {
    71  	return m.TID
    72  }
    73  
    74  func (m *message) RequestTimeout() time.Duration {
    75  	return m.RTimeout
    76  }
    77  
    78  func (m *message) Namespace() string {
    79  	return m.NameSpace
    80  }
    81  
    82  func (m *message) WithTraceID(s string) {
    83  	m.TID = s
    84  }
    85  
    86  func (m *message) WithRequestTimeout(t time.Duration) {
    87  	m.RTimeout = t
    88  }
    89  
    90  func (m *message) WithNamespace(n string) {
    91  	m.NameSpace = n
    92  }
    93  
    94  func (m *message) ServerName() string {
    95  	return m.ServerN
    96  }
    97  
    98  func (m *message) WithServerName(n string) {
    99  	m.ServerN = n
   100  }
   101  
   102  func (m *message) Metadata() metadata.MD {
   103  	d, err := json.Marshal(m)
   104  	if err != nil {
   105  		return metadata.Pairs(MDMessage.String(), "{}")
   106  	}
   107  	return metadata.Pairs(MDMessage.String(), string(d))
   108  }
   109  
   110  func Message(ctx context.Context) MSG {
   111  	defaultMsg := message{
   112  		context:  ctx,
   113  		RTimeout: time.Second * 5,
   114  		TID:      hash.SHA256(time.Now().UnixNano()),
   115  	}
   116  	md, exist := metadata.FromIncomingContext(ctx)
   117  	if !exist {
   118  		return &defaultMsg
   119  	}
   120  	vals := md.Get(MDMessage.String())
   121  	if len(vals) == 0 {
   122  		return &defaultMsg
   123  	}
   124  	var msg message
   125  	if err := json.Unmarshal([]byte(vals[0]), &msg); err != nil {
   126  		return &defaultMsg
   127  	}
   128  	msg.context = ctx
   129  	return &msg
   130  }