-
Notifications
You must be signed in to change notification settings - Fork 5.5k
/
thrift_serialize.go
157 lines (128 loc) · 3.89 KB
/
thrift_serialize.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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
/*
A small cli utility meant to convert json to zipkin thrift binary format, and
vice versa.
To convert from json to thrift,
the json is unmarshalled, converted to zipkincore.Span structures, and
marshalled into thrift binary protocol. The json must be in an array format (even if it only has one object),
because the tool automatically tries to unmarshal the json into an array of structs.
To convert from thrift to json,
the opposite process must happen. The thrift binary data must be read into an array of
zipkin span structures, and those spans must be marshalled into json.
Usage:
./thrift_serialize -input <input-file> -output <output-file> -deserialize<true|false>
If `deserialize` is set to true (false by default), the tool will interpret the input file as
thrift, and write it as json to the output file.
Otherwise, the input file will be interpreted as json, and the output will be encoded as thrift.
*/
package main
import (
"encoding/json"
"errors"
"flag"
"fmt"
"io/ioutil"
"log"
"github.com/apache/thrift/lib/go/thrift"
"github.com/openzipkin/zipkin-go-opentracing/thrift/gen-go/zipkincore"
)
var (
filename string
outFileName string
inputType string
)
const usage = `./json_serialize -input <input> -output output -input-type<json|thrift>`
func init() {
flag.StringVar(&filename, "input", "", usage)
flag.StringVar(&outFileName, "output", "", usage)
flag.StringVar(&inputType, "input-type", "thrift", usage)
}
func main() {
flag.Parse()
contents, err := ioutil.ReadFile(filename)
if err != nil {
log.Fatalf("Error reading file: %v\n", err)
}
switch inputType {
case "json":
raw, err := jsonToZipkinThrift(contents)
if err != nil {
log.Fatalf("%v\n", err)
}
if err := ioutil.WriteFile(outFileName, raw, 0644); err != nil {
log.Fatalf("%v", err)
}
case "thrift":
raw, err := thriftToJSONSpans(contents)
if err != nil {
log.Fatalf("%v\n", err)
}
if err := ioutil.WriteFile(outFileName, raw, 0644); err != nil {
log.Fatalf("%v", err)
}
default:
log.Fatalf("Unsupported input type")
}
}
func jsonToZipkinThrift(jsonRaw []byte) ([]byte, error) {
if len(jsonRaw) == 0 {
return nil, errors.New("no data")
}
if string(jsonRaw)[0] != '[' {
return nil, errors.New("cannot unmarshal non array type")
}
var spans []*zipkincore.Span
err := json.Unmarshal(jsonRaw, &spans)
if err != nil {
return nil, fmt.Errorf("error unmarshalling: %v", err)
}
var zspans []*zipkincore.Span
if err != nil {
return nil, err
}
zspans = append(zspans, spans...)
fmt.Println(spans)
buf := thrift.NewTMemoryBuffer()
transport := thrift.NewTBinaryProtocolTransport(buf)
if err = transport.WriteListBegin(thrift.STRUCT, len(spans)); err != nil {
return nil, fmt.Errorf("error in beginning thrift write: %v", err)
}
for _, span := range zspans {
err = span.Write(transport)
if err != nil {
return nil, fmt.Errorf("error converting zipkin struct to thrift: %v", err)
}
}
if err = transport.WriteListEnd(); err != nil {
return nil, fmt.Errorf("error finishing thrift write: %v", err)
}
return buf.Bytes(), nil
}
func thriftToJSONSpans(thriftData []byte) ([]byte, error) {
buffer := thrift.NewTMemoryBuffer()
if _, err := buffer.Write(thriftData); err != nil {
err = fmt.Errorf("error in buffer write: %v", err)
return nil, err
}
transport := thrift.NewTBinaryProtocolTransport(buffer)
_, size, err := transport.ReadListBegin()
if err != nil {
err = fmt.Errorf("error in ReadListBegin: %v", err)
return nil, err
}
var spans []*zipkincore.Span
for i := 0; i < size; i++ {
zs := &zipkincore.Span{}
if err = zs.Read(transport); err != nil {
err = fmt.Errorf("Error reading into zipkin struct: %v", err)
return nil, err
}
spans = append(spans, zs)
}
err = transport.ReadListEnd()
if err != nil {
err = fmt.Errorf("error ending thrift read: %v", err)
return nil, err
}
out, _ := json.MarshalIndent(spans, "", " ")
return out, nil
}