diff --git a/lib/results.go b/lib/results.go index e9822e7e..d871b043 100644 --- a/lib/results.go +++ b/lib/results.go @@ -6,6 +6,8 @@ import ( "encoding/base64" "encoding/csv" "encoding/gob" + "encoding/json" + "fmt" "io" "net/http" "net/textproto" @@ -157,14 +159,20 @@ func (dec Decoder) Decode(r *Result) error { return dec(r) } type Encoder func(*Result) error // NewEncoder returns a new Result encoder closure for the given io.Writer -func NewEncoder(r io.Writer) Encoder { - enc := gob.NewEncoder(r) +func NewEncoder(w io.Writer) Encoder { + enc := gob.NewEncoder(w) return func(r *Result) error { return enc.Encode(r) } } // Encode is an an adapter method calling the Encoder function itself with the // given parameters. -func (enc Encoder) Encode(r *Result) error { return enc(r) } +func (enc Encoder) Encode(r *Result) error { + if err := encodeProtocol(r); err != nil { + return err + } + + return enc(r) +} // NewCSVEncoder returns an Encoder that dumps the given *Result as a CSV // record. The columns are: UNIX timestamp in ns since epoch, @@ -186,6 +194,7 @@ func NewCSVEncoder(w io.Writer) Encoder { r.Method, r.URL, base64.StdEncoding.EncodeToString(headerBytes(r.Headers)), + r.Protocol, }) if err != nil { return err @@ -207,9 +216,9 @@ func headerBytes(h http.Header) []byte { } // NewCSVDecoder returns a Decoder that decodes CSV encoded Results. -func NewCSVDecoder(r io.Reader) Decoder { - dec := csv.NewReader(r) - dec.FieldsPerRecord = 12 +func NewCSVDecoder(rd io.Reader) Decoder { + dec := csv.NewReader(rd) + dec.FieldsPerRecord = 13 dec.TrimLeadingSpace = true return func(r *Result) error { @@ -267,6 +276,10 @@ func NewCSVDecoder(r io.Reader) Decoder { r.Headers = http.Header(hdr) } + if rec[12] != "" { + r.Protocol = rec[12] + } + return err } } @@ -278,27 +291,70 @@ type jsonResult Result // NewJSONEncoder returns an Encoder that dumps the given *Results as a JSON // object. func NewJSONEncoder(w io.Writer) Encoder { - var jw jwriter.Writer + var enc jwriter.Writer return func(r *Result) error { - (*jsonResult)(r).MarshalEasyJSON(&jw) - if jw.Error != nil { - return jw.Error + (*jsonResult)(r).MarshalEasyJSON(&enc) + if enc.Error != nil { + return enc.Error } - jw.RawByte('\n') - _, err := jw.DumpTo(w) + enc.RawByte('\n') + _, err := enc.DumpTo(w) return err } } // NewJSONDecoder returns a Decoder that decodes JSON encoded Results. -func NewJSONDecoder(r io.Reader) Decoder { - rd := bufio.NewReader(r) +func NewJSONDecoder(rd io.Reader) Decoder { + dec := bufio.NewReader(rd) return func(r *Result) (err error) { var jl jlexer.Lexer - if jl.Data, err = rd.ReadBytes('\n'); err != nil { + if jl.Data, err = dec.ReadBytes('\n'); err != nil { return err } (*jsonResult)(r).UnmarshalEasyJSON(&jl) return jl.Error() } } + +func encodeProtocol(r *Result) error { + if r.Protocol == "gql" { + err := parseGraphQLErrors(r) + return err + } + + return nil +} + +type GQLResponse struct { + Data interface{} + Errors []GQLError +} + +type GQLError struct { + Extensions interface{} + Message string +} + +func parseGraphQLErrors(r *Result) error { + if r.Code < 200 || r.Code >= 400 { + return nil + } + + var res GQLResponse + err := json.Unmarshal(r.Body, &res) + if err != nil { + return err + } + + if res.Errors != nil && len(res.Errors) > 0 { + for i, e := range res.Errors { + if i == 0 { + r.Error = e.Message + } else { + r.Error = fmt.Sprintf("%v, %v", r.Error, e.Message) + } + } + } + + return nil +} diff --git a/lib/results_easyjson.go b/lib/results_easyjson.go index 844e3b95..7d3b0dc1 100644 --- a/lib/results_easyjson.go +++ b/lib/results_easyjson.go @@ -4,12 +4,11 @@ package vegeta import ( json "encoding/json" - http "net/http" - time "time" - easyjson "github.com/mailru/easyjson" jlexer "github.com/mailru/easyjson/jlexer" jwriter "github.com/mailru/easyjson/jwriter" + http "net/http" + time "time" ) // suppress unused package warning @@ -105,6 +104,8 @@ func easyjsonBd1621b8DecodeGithubComTsenartVegetaV12Lib(in *jlexer.Lexer, out *j } in.Delim('}') } + case "protocol": + out.Protocol = string(in.String()) default: in.SkipRecursive() } @@ -206,14 +207,33 @@ func easyjsonBd1621b8EncodeGithubComTsenartVegetaV12Lib(out *jwriter.Writer, in out.RawByte('}') } } + { + const prefix string = ",\"protocol\":" + out.RawString(prefix) + out.String(string(in.Protocol)) + } out.RawByte('}') } +// MarshalJSON supports json.Marshaler interface +func (v jsonResult) MarshalJSON() ([]byte, error) { + w := jwriter.Writer{} + easyjsonBd1621b8EncodeGithubComTsenartVegetaV12Lib(&w, v) + return w.Buffer.BuildBytes(), w.Error +} + // MarshalEasyJSON supports easyjson.Marshaler interface func (v jsonResult) MarshalEasyJSON(w *jwriter.Writer) { easyjsonBd1621b8EncodeGithubComTsenartVegetaV12Lib(w, v) } +// UnmarshalJSON supports json.Unmarshaler interface +func (v *jsonResult) UnmarshalJSON(data []byte) error { + r := jlexer.Lexer{Data: data} + easyjsonBd1621b8DecodeGithubComTsenartVegetaV12Lib(&r, v) + return r.Error() +} + // UnmarshalEasyJSON supports easyjson.Unmarshaler interface func (v *jsonResult) UnmarshalEasyJSON(l *jlexer.Lexer) { easyjsonBd1621b8DecodeGithubComTsenartVegetaV12Lib(l, v) diff --git a/lib/results_test.go b/lib/results_test.go index a6da0175..66b8fa99 100644 --- a/lib/results_test.go +++ b/lib/results_test.go @@ -96,10 +96,10 @@ func TestResultEncoding(t *testing.T) { BytesIn: rapid.Uint64().Draw(t, "bytes_in"), BytesOut: rapid.Uint64().Draw(t, "bytes_out"), Error: rapid.StringMatching(`^\w+$`).Draw(t, "error"), - Body: rapid.SliceOf(rapid.Byte()).Draw(t, "body"), - Method: rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$"). - Draw(t, "method"), - URL: rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&%\$#\=~]*)$`).Draw(t, "url"), + Body: []byte("{\"data\":{\"vegeta\":\"punch\"}}"), + Method: rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$").Draw(t, "method"), + URL: rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&%\$#\=~]*)$`).Draw(t, "url"), + Protocol: rapid.StringMatching("^(http|gql)$").Draw(t, "protocol"), } if len(hdrs) > 0 { @@ -115,7 +115,7 @@ func TestResultEncoding(t *testing.T) { var buf bytes.Buffer enc := tc.enc(&buf) for j := 0; j < 2; j++ { - if err := enc(&want); err != nil { + if err := enc.Encode(&want); err != nil { t.Fatal(err) } } @@ -128,7 +128,7 @@ func TestResultEncoding(t *testing.T) { } for j := 0; j < 2; j++ { var got Result - if err := dec(&got); err != nil { + if err := dec.Decode(&got); err != nil { t.Fatalf("err: %q buffer: %s", err, encoded) } @@ -142,6 +142,66 @@ func TestResultEncoding(t *testing.T) { } } +func TestGQLEncoding(t *testing.T) { + rapid.Check(t, func(t *rapid.T) { + hdrs := rapid.MapOf( + rapid.StringMatching("^[!#$%&'*+\\-.^_`|~0-9a-zA-Z]+$"), + rapid.SliceOfN(rapid.StringMatching(`^[0-9a-zA-Z]+$`), 1, -1), + ).Draw(t, "headers") + + in := Result{ + Attack: "gql-test", + Seq: rapid.Uint64().Draw(t, "seq"), + Code: 200, + Timestamp: time.Unix(rapid.Int64Range(0, 1e8).Draw(t, "timestamp"), 0), + Latency: time.Duration(rapid.Int64Min(0).Draw(t, "latency")), + BytesIn: rapid.Uint64().Draw(t, "bytes_in"), + BytesOut: rapid.Uint64().Draw(t, "bytes_out"), + Error: "", + Body: []byte("{\"data\":{},\"errors\":[{\"message\":\"no punch\"}]}"), + Method: rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$").Draw(t, "method"), + URL: rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&%\$#\=~]*)$`).Draw(t, "url"), + Protocol: "gql", + } + + if len(hdrs) > 0 { + in.Headers = make(http.Header, len(hdrs)) + } + + for k, vs := range hdrs { + for _, v := range vs { + in.Headers.Add(k, v) + } + } + + want := in + want.Error = "no punch" + + var buf bytes.Buffer + enc := NewJSONEncoder(&buf) + for j := 0; j < 2; j++ { + if err := enc.Encode(&in); err != nil { + t.Fatal(err) + } + } + + encoded := buf.String() + + dec := NewJSONDecoder(&buf) + for j := 0; j < 2; j++ { + var got Result + if err := dec.Decode(&got); err != nil { + t.Fatalf("err: %q buffer: %s", err, encoded) + } + + if !got.Equal(want) { + t.Logf("encoded: %s", encoded) + t.Fatalf("mismatch: %s", cmp.Diff(got, want)) + } + } + }) +} + func BenchmarkResultEncodings(b *testing.B) { b.StopTimer() b.ResetTimer()