/
rpcx.go
59 lines (52 loc) · 1.35 KB
/
rpcx.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
package logtrace
import (
"context"
"github.com/smallnest/rpcx/share"
"github.com/tal-tech/xtools/jsutil"
)
var rpcxMetaKey = share.ReqMetaDataKey
func InjectTraceNodeToRpcx(ctx context.Context) context.Context {
//IncrementRpcId(ctx)
meta := ExtractTraceNodeFromXesContext(ctx)
traceRpcId := meta.Get("x_rpcid")
if len(traceRpcId) == 0 {
return ctx
}
metaStr, err := jsutil.Json.MarshalToString(meta.ForkMap())
if err != nil {
metaStr = ""
}
rpcxMeta := ctx.Value(rpcxMetaKey)
if rpcxMeta == nil {
//todo: why enter here
logMeta := map[string]string{
GetMetadataKey(): metaStr,
}
ctx = context.WithValue(ctx, rpcxMetaKey, logMeta)
} else {
if logMeta, ok := rpcxMeta.(map[string]string); ok {
logMeta[GetMetadataKey()] = metaStr
ctx = context.WithValue(ctx, rpcxMetaKey, logMeta)
}
}
return ctx
}
func ExtractTraceNodeToXexContext(ctx context.Context) context.Context {
reqMeta, ok := ctx.Value(rpcxMetaKey).(map[string]string)
if ok {
logMetaKey := GetMetadataKey()
metaStr, ok1 := reqMeta[logMetaKey]
if ok1 {
var mapVal map[string]string
if jsutil.Json.UnmarshalFromString(metaStr, &mapVal) == nil {
tNode := ExtractTraceNodeFromXesContext(ctx)
for k, v := range mapVal {
tNode.Set(k, v)
}
ctx = context.WithValue(ctx, GetMetadataKey(), tNode)
}
AppendNewRpcId(ctx)
}
}
return ctx
}