Zipkin thrift cli tool refactor

plugins/inputs/zipkin/cmd/thrift_serialize/thrift_serialize.go

@@ -1,7 +1,32 @@
+/*
+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 unmarshall 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"
@@ -11,101 +36,118 @@ import (
 )
 
 var (
-	filename    *string
-	outFileName *string
+	filename    string
+	outFileName string
+	inputType   string
 )
 
-/*func init() {
-	filename = flag.String("input", "", usage)
-	outFileName = flag.String("output", "", usage)
-}
+const usage = `./json_serialize -input <input> -output output -input-type<json|thrift>`
 
-const usage = `./json_serialize -input <input> -output output`*/
+func init() {
+	flag.StringVar(&filename, "input", "", usage)
+	flag.StringVar(&outFileName, "output", "", usage)
+	flag.StringVar(&inputType, "input-type", "thrift", usage)
+}
 
 func main() {
-	b, err := ioutil.ReadFile("../testdata/json/distributed_trace_sample.json")
+	flag.Parse()
+	contents, err := ioutil.ReadFile(filename)
 	if err != nil {
-		log.Fatalf("error: %v\n", err)
+		log.Fatalf("Error reading file: %v\n", err)
+	}
+
+	switch inputType {
+	case "json":
+		raw, err := jsonToZipkinThrift(contents)
+		if err != nil {
+			log.Fatalf("%v\n", err)
+		}
+		ioutil.WriteFile(outFileName, raw, 0644)
+	case "thrift":
+		raw, err := thriftToJSONSpans(contents)
+		if err != nil {
+			log.Fatalf("%v\n", err)
+		}
+		ioutil.WriteFile(outFileName, raw, 0644)
+	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
-	span, err := serializeJSON(b)
-	spans = append(spans, span)
+	err := json.Unmarshal(jsonRaw, &spans)
+	if err != nil {
+		return nil, fmt.Errorf("error unmarshalling: %v", err)
+	}
 
+	var zspans []*zipkincore.Span
 	if err != nil {
-		log.Fatalf("error: %v\n", err)
+		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 {
-		panic(err)
+		return nil, fmt.Errorf("error in beginning thrift write: %v", err)
 	}
 
-	for _, s := range spans {
-		err = s.Write(transport)
+	for _, span := range zspans {
+		err = span.Write(transport)
 		if err != nil {
-			panic(err)
+			return nil, fmt.Errorf("error converting zipkin struct to thrift: %v", err)
 		}
 	}
 
 	if err = transport.WriteListEnd(); err != nil {
-		panic(err)
+		return nil, fmt.Errorf("error finishing thrift write: %v", err)
 	}
 
-	log.Println(buf.Buffer.String())
-
-	ioutil.WriteFile("../testdata/out.dat", buf.Buffer.Bytes(), 0644)
-
-	b, err = ioutil.ReadFile("../testdata/out.dat")
-
-	//log.Println("read bytes: ", b)
-	if err != nil {
-		log.Printf("%v\n", err)
-	}
-	deserializeThrift(b)
-}
-
-func serializeJSON(b []byte) (*zipkincore.Span, error) {
-	var span *zipkincore.Span
-	err := json.Unmarshal(b, &span)
-	return span, err
+	return buf.Bytes(), nil
 }
 
-func deserializeThrift(b []byte) {
+func thriftToJSONSpans(thriftData []byte) ([]byte, error) {
 	buffer := thrift.NewTMemoryBuffer()
-	if _, err := buffer.Write(b); err != nil {
-		log.Println("Error in buffer write: ", err)
-		return
+	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 {
-		log.Printf("Error in ReadListBegin: %s", err)
-		return
+		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 {
-			log.Printf("Error reading into zipkin struct: %s", err)
-			return
+			err = fmt.Errorf("Error reading into zipkin struct: %v", err)
+			return nil, err
 		}
 		spans = append(spans, zs)
 	}
 
 	err = transport.ReadListEnd()
 	if err != nil {
-		log.Printf("%s", err)
-		return
+		err = fmt.Errorf("error ending thrift read: %v", err)
+		return nil, err
 	}
 
-	//marshal json for debugging purposes
 	out, _ := json.MarshalIndent(spans, "", "    ")
-	log.Println(string(out))
+	return out, nil
 }

plugins/inputs/zipkin/testdata/json/distributed_trace_sample.json

@@ -1,4 +1,4 @@
-{
+[{
   "trace_id": 6802735349851856000,
   "name": "main.dud",
   "id": 6802735349851856000,
@@ -27,4 +27,4 @@
   ],
   "binary_annotations": [],
   "debug": true
-}
+}]