/
util.go
110 lines (89 loc) · 2.18 KB
/
util.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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package awsxray
import (
"crypto/rand"
"fmt"
"strings"
)
// getRandom generates a random hex encoded string
func getRandom(i int) string {
b := make([]byte, i)
for {
// keep trying till we get it
if _, err := rand.Read(b); err != nil {
continue
}
return fmt.Sprintf("%x", b)
}
}
// GetTraceId returns trace id from header or blank
func GetTraceId(header string) string {
for _, h := range strings.Split(header, ";") {
th := strings.TrimSpace(h)
if strings.HasPrefix(th, "Root=") {
return strings.TrimPrefix(th, "Root=")
}
}
if len(header) > 0 {
// return as is
return header
}
return ""
}
// GetParentId returns parent id from header or blank
func GetParentId(header string) string {
for _, h := range strings.Split(header, ";") {
th := strings.TrimSpace(h)
if strings.HasPrefix(th, "Parent=") {
return strings.TrimPrefix(th, "Parent=")
}
}
return ""
}
// SetParentId will set the parent id of a trace in the header
func SetParentId(header, parentId string) string {
parentHeader := fmt.Sprintf("Parent=%s", parentId)
// no existing header?
if len(header) == 0 {
return parentHeader
}
headers := strings.Split(header, ";")
for i, h := range headers {
th := strings.TrimSpace(h)
if len(th) == 0 {
continue
}
// get Parent=Id match
if strings.HasPrefix(th, "Parent=") {
// set parent header
headers[i] = parentHeader
// return entire header
return strings.Join(headers, "; ")
}
}
// no match; set new parent header
return strings.Join(append(headers, parentHeader), "; ")
}
// SetTraceId will set the trace id in the header
func SetTraceId(header, traceId string) string {
traceHeader := fmt.Sprintf("Root=%s", traceId)
// no existing header?
if len(header) == 0 {
return traceHeader
}
headers := strings.Split(header, ";")
for i, h := range headers {
th := strings.TrimSpace(h)
if len(th) == 0 {
continue
}
// get Root=Id match
if strings.HasPrefix(th, "Root=") {
// set trace header
headers[i] = traceHeader
// return entire header
return strings.Join(headers, "; ")
}
}
// no match; set new trace header as first entry
return strings.Join(append([]string{traceHeader}, headers...), "; ")
}