diff --git a/.gitignore b/.gitignore index d0c0e875..295f2d43 100644 --- a/.gitignore +++ b/.gitignore @@ -32,3 +32,4 @@ vendor *.lz .DS_Store +.idea diff --git a/go.mod b/go.mod index cf7036e2..86de5473 100644 --- a/go.mod +++ b/go.mod @@ -23,5 +23,6 @@ require ( github.com/miekg/dns v1.1.17 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e + github.com/valyala/fastjson v1.5.0 golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297 ) diff --git a/go.sum b/go.sum index 83081219..db9a8d1e 100644 --- a/go.sum +++ b/go.sum @@ -38,6 +38,8 @@ github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 h1:7z3LSn867ex6 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU= github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e h1:bB5SXzQmSUsJCmjPDN9fKYx3SSDER5diSjlN6TefTCc= github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e/go.mod h1:SWZznP1z5Ki7hDT2ioqiFKEse8K9tU2OUvaRI0NeGQo= +github.com/valyala/fastjson v1.5.0 h1:DGrb4wEYso2HdGLyLmNoyNCQnCWfjd8yhghPv5/5YQg= +github.com/valyala/fastjson v1.5.0/go.mod h1:CLCAqky6SMuOcxStkYQvblddUtoRxhYMGLrsQns1aXY= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472 h1:Gv7RPwsi3eZ2Fgewe3CBsuOebPwO27PoXzRpJPsvSSM= golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= diff --git a/lib/results.go b/lib/results.go index 3d62cb75..15aa3ef5 100644 --- a/lib/results.go +++ b/lib/results.go @@ -13,9 +13,6 @@ import ( "strconv" "strings" "time" - - "github.com/mailru/easyjson/jlexer" - jwriter "github.com/mailru/easyjson/jwriter" ) func init() { @@ -267,34 +264,3 @@ func NewCSVDecoder(r io.Reader) Decoder { } } -//go:generate easyjson -no_std_marshalers -output_filename results_easyjson.go results.go -//easyjson:json -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 - return func(r *Result) error { - (*jsonResult)(r).MarshalEasyJSON(&jw) - if jw.Error != nil { - return jw.Error - } - jw.RawByte('\n') - _, err := jw.DumpTo(w) - return err - } -} - -// NewJSONDecoder returns a Decoder that decodes JSON encoded Results. -func NewJSONDecoder(r io.Reader) Decoder { - rd := bufio.NewReader(r) - return func(r *Result) (err error) { - var jl jlexer.Lexer - if jl.Data, err = rd.ReadBytes('\n'); err != nil { - return err - } - (*jsonResult)(r).UnmarshalEasyJSON(&jl) - return jl.Error() - } -} diff --git a/lib/results_easyjson.go b/lib/results_easyjson.go deleted file mode 100644 index 47c1edd3..00000000 --- a/lib/results_easyjson.go +++ /dev/null @@ -1,205 +0,0 @@ -// Code generated by easyjson for marshaling/unmarshaling. DO NOT EDIT. - -package vegeta - -import ( - json "encoding/json" - "net/http" - time "time" - - easyjson "github.com/mailru/easyjson" - jlexer "github.com/mailru/easyjson/jlexer" - jwriter "github.com/mailru/easyjson/jwriter" -) - -// suppress unused package warning -var ( - _ *json.RawMessage - _ *jlexer.Lexer - _ *jwriter.Writer - _ easyjson.Marshaler -) - -func easyjsonBd1621b8DecodeGithubComTsenartVegetaLib(in *jlexer.Lexer, out *jsonResult) { - isTopLevel := in.IsStart() - if in.IsNull() { - if isTopLevel { - in.Consumed() - } - in.Skip() - return - } - in.Delim('{') - for !in.IsDelim('}') { - key := in.UnsafeString() - in.WantColon() - if in.IsNull() { - in.Skip() - in.WantComma() - continue - } - switch key { - case "attack": - out.Attack = string(in.String()) - case "seq": - out.Seq = uint64(in.Uint64()) - case "code": - out.Code = uint16(in.Uint16()) - case "timestamp": - if data := in.Raw(); in.Ok() { - in.AddError((out.Timestamp).UnmarshalJSON(data)) - } - case "latency": - out.Latency = time.Duration(in.Int64()) - case "bytes_out": - out.BytesOut = uint64(in.Uint64()) - case "bytes_in": - out.BytesIn = uint64(in.Uint64()) - case "error": - out.Error = string(in.String()) - case "body": - if in.IsNull() { - in.Skip() - out.Body = nil - } else { - out.Body = in.Bytes() - } - case "method": - out.Method = string(in.String()) - case "url": - out.URL = string(in.String()) - case "headers": - out.Headers = easyjsonUnmarshalHeaders(in) - default: - in.SkipRecursive() - } - in.WantComma() - } - in.Delim('}') - if isTopLevel { - in.Consumed() - } -} -func easyjsonBd1621b8EncodeGithubComTsenartVegetaLib(out *jwriter.Writer, in jsonResult) { - out.RawByte('{') - first := true - _ = first - { - const prefix string = ",\"attack\":" - out.RawString(prefix[1:]) - out.String(string(in.Attack)) - } - { - const prefix string = ",\"seq\":" - out.RawString(prefix) - out.Uint64(uint64(in.Seq)) - } - { - const prefix string = ",\"code\":" - out.RawString(prefix) - out.Uint16(uint16(in.Code)) - } - { - const prefix string = ",\"timestamp\":" - out.RawString(prefix) - out.Raw((in.Timestamp).MarshalJSON()) - } - { - const prefix string = ",\"latency\":" - out.RawString(prefix) - out.Int64(int64(in.Latency)) - } - { - const prefix string = ",\"bytes_out\":" - out.RawString(prefix) - out.Uint64(uint64(in.BytesOut)) - } - { - const prefix string = ",\"bytes_in\":" - out.RawString(prefix) - out.Uint64(uint64(in.BytesIn)) - } - { - const prefix string = ",\"error\":" - out.RawString(prefix) - out.String(string(in.Error)) - } - { - const prefix string = ",\"body\":" - out.RawString(prefix) - out.Base64Bytes(in.Body) - } - { - const prefix string = ",\"method\":" - out.RawString(prefix) - out.String(string(in.Method)) - } - { - const prefix string = ",\"url\":" - out.RawString(prefix) - out.String(string(in.URL)) - } - { - const prefix string = ",\"headers\":" - out.RawString(prefix) - easyjsonMarshalHeaders(out, in.Headers) - } - out.RawByte('}') -} - -func easyjsonUnmarshalHeaders(in *jlexer.Lexer) http.Header { - h := http.Header{} - in.Delim('[') - for !in.IsDelim(']') { - for in.IsDelim('{') { - in.Delim('{') - var key string - var values []string - for !in.IsDelim('}') { - k := in.UnsafeString() - in.WantColon() - if in.IsNull() { - in.Skip() - in.WantComma() - continue - } - switch k { - case "key": - key = in.String() - case "value": - values = append(values, in.String()) - } - in.WantComma() - } - h[key] = values - in.Delim('}') - in.WantComma() - } - } - in.Delim(']') - return h -} - -func easyjsonMarshalHeaders(w *jwriter.Writer, h http.Header) { - w.RawByte('[') - for key, values := range h { - for _, value := range values { - w.RawString(`{"key":`) - w.String(key) - w.RawString(`,"value":`) - w.String(value) - w.RawByte('}') - } - } - w.RawByte(']') -} - -// MarshalEasyJSON supports easyjson.Marshaler interface -func (v jsonResult) MarshalEasyJSON(w *jwriter.Writer) { - easyjsonBd1621b8EncodeGithubComTsenartVegetaLib(w, v) -} - -// UnmarshalEasyJSON supports easyjson.Unmarshaler interface -func (v *jsonResult) UnmarshalEasyJSON(l *jlexer.Lexer) { - easyjsonBd1621b8DecodeGithubComTsenartVegetaLib(l, v) -} diff --git a/lib/results_json.go b/lib/results_json.go new file mode 100644 index 00000000..95c294ea --- /dev/null +++ b/lib/results_json.go @@ -0,0 +1,268 @@ +package vegeta + +import ( + "bufio" + "encoding/base64" + "io" + "net/http" + "reflect" + "strconv" + "strings" + "time" + "unsafe" + + "github.com/valyala/fastjson" +) + +// NewJSONEncoder returns an Encoder that dumps the given *Results as a JSON +// object. +func NewJSONEncoder(w io.Writer) Encoder { + buf := make([]byte, 0, 4096) + return func(r *Result) error { + buf = buf[:0] + buf = append(buf, `{"attack":"`...) + if len(r.Attack) > 0 { + buf = appendJSONString(buf, r.Attack) + } + + buf = append(buf, `","seq":`...) + buf = strconv.AppendUint(buf, r.Seq, 10) + buf = append(buf, `,"code":`...) + buf = strconv.AppendUint(buf, uint64(r.Code), 10) + buf = append(buf, `,"timestamp":"`...) + buf = r.Timestamp.AppendFormat(buf, time.RFC3339Nano) + buf = append(buf, `","latency":`...) + buf = strconv.AppendInt(buf, int64(r.Latency), 10) + buf = append(buf, `,"bytes_out":`...) + buf = strconv.AppendUint(buf, r.BytesOut, 10) + buf = append(buf, `,"bytes_in":`...) + buf = strconv.AppendUint(buf, r.BytesIn, 10) + + buf = append(buf, `,"error":"`...) + if len(r.Error) > 0 { + buf = appendJSONString(buf, r.Error) + } + + buf = append(buf, `","body":"`...) + if len(r.Body) > 0 { + buf = appendBase64(buf, r.Body) + } + + buf = append(buf, `","method":"`...) + buf = append(buf, r.Method...) + buf = append(buf, `","url":"`...) + buf = appendJSONString(buf, r.URL) + + buf = append(buf, `","headers":{`...) + for k, vs := range r.Headers { + buf = append(buf, '"') + buf = appendJSONString(buf, k) + buf = append(buf, `":[`...) + for _, v := range vs { + buf = append(buf, '"') + buf = appendJSONString(buf, v) + buf = append(buf, `",`...) + } + if len(vs) > 0 { + buf = buf[:len(buf)-1] + } + buf = append(buf, `],`...) + } + if len(r.Headers) > 0 { + buf = buf[:len(buf)-1] + } + buf = append(buf, "}}\n"...) + + _, err := w.Write(buf) + return err + } +} + +// NewJSONDecoder returns a Decoder that decodes JSON encoded Results. +func NewJSONDecoder(r io.Reader) Decoder { + var p fastjson.Parser + rd := bufio.NewReader(r) + return func(r *Result) (err error) { + line, err := rd.ReadBytes('\n') + if err != nil { + return err + } + + v, err := p.ParseBytes(line) + if err != nil { + return err + } + + r.Attack = string(v.GetStringBytes("attack")) + r.Seq = v.GetUint64("seq") + r.Code = uint16(v.GetUint("code")) + + r.Timestamp, err = time.Parse(time.RFC3339Nano, string(v.GetStringBytes("timestamp"))) + if err != nil { + return err + } + + r.Latency = time.Duration(v.GetInt64("latency")) + r.BytesIn = v.GetUint64("bytes_in") + r.BytesOut = v.GetUint64("bytes_out") + r.Error = string(v.GetStringBytes("error")) + + body := v.GetStringBytes("body") + r.Body = make([]byte, base64.StdEncoding.DecodedLen(len(body))) + n, err := base64.StdEncoding.Decode(r.Body, body) + if err != nil { + return err + } + r.Body = r.Body[:n] + + r.Method = string(v.GetStringBytes("method")) + r.URL = string(v.GetStringBytes("url")) + + headers, err := v.Get("headers").Object() + if err != nil { + return err + } + + r.Headers = make(http.Header, headers.Len()) + headers.Visit(func(key []byte, v *fastjson.Value) { + if err != nil { // Previous visit errored + return + } + + var vs []*fastjson.Value + if vs, err = v.Array(); err != nil { + return + } + + k := string(key) + for _, v := range vs { + r.Headers[k] = append(r.Headers[k], string(v.GetStringBytes())) + } + }) + + return err + } +} + +func appendBase64(buf []byte, bs []byte) []byte { + n := base64.StdEncoding.EncodedLen(len(bs)) + buf = expand(buf, n) + base64.StdEncoding.Encode(buf[len(buf):len(buf)+n], bs) + return buf[:len(buf)+n] +} + +// expand grows the given buf to have enough capacity to hold n +// extra bytes beyond the current len +func expand(buf []byte, n int) []byte { + l := len(buf) + free := cap(buf) - l + grow := n - free + if grow > 0 { + buf = append(buf[:cap(buf)], make([]byte, grow)...)[:l] + } + return buf +} + +// The following code was copied and adapted from https://github.com/valyala/quicktemplate + +func s2b(s string) []byte { + sh := (*reflect.StringHeader)(unsafe.Pointer(&s)) + bh := reflect.SliceHeader{ + Data: sh.Data, + Len: sh.Len, + Cap: sh.Len, + } + return *(*[]byte)(unsafe.Pointer(&bh)) +} + +func b2s(z []byte) string { + return *(*string)(unsafe.Pointer(&z)) +} + +func appendJSONString(buf []byte, s string) []byte { + if len(s) > 24 && + strings.IndexByte(s, '"') < 0 && + strings.IndexByte(s, '\\') < 0 && + strings.IndexByte(s, '\n') < 0 && + strings.IndexByte(s, '\r') < 0 && + strings.IndexByte(s, '\t') < 0 && + strings.IndexByte(s, '\f') < 0 && + strings.IndexByte(s, '\b') < 0 && + strings.IndexByte(s, '<') < 0 && + strings.IndexByte(s, '\'') < 0 && + strings.IndexByte(s, 0) < 0 { + + // fast path - nothing to escape + return append(buf, s2b(s)...) + } + + // slow path + write := func(bs []byte) { buf = append(buf, bs...) } + b := s2b(s) + j := 0 + n := len(b) + if n > 0 { + // Hint the compiler to remove bounds checks in the loop below. + _ = b[n-1] + } + for i := 0; i < n; i++ { + switch b[i] { + case '"': + write(b[j:i]) + write(strBackslashQuote) + j = i + 1 + case '\\': + write(b[j:i]) + write(strBackslashBackslash) + j = i + 1 + case '\n': + write(b[j:i]) + write(strBackslashN) + j = i + 1 + case '\r': + write(b[j:i]) + write(strBackslashR) + j = i + 1 + case '\t': + write(b[j:i]) + write(strBackslashT) + j = i + 1 + case '\f': + write(b[j:i]) + write(strBackslashF) + j = i + 1 + case '\b': + write(b[j:i]) + write(strBackslashB) + j = i + 1 + case '<': + write(b[j:i]) + write(strBackslashLT) + j = i + 1 + case '\'': + write(b[j:i]) + write(strBackslashQ) + j = i + 1 + case 0: + write(b[j:i]) + write(strBackslashZero) + j = i + 1 + } + } + write(b[j:]) + + return buf +} + +var ( + strBackslashQuote = []byte(`\"`) + strBackslashBackslash = []byte(`\\`) + strBackslashN = []byte(`\n`) + strBackslashR = []byte(`\r`) + strBackslashT = []byte(`\t`) + strBackslashF = []byte(`\u000c`) + strBackslashB = []byte(`\u0008`) + strBackslashLT = []byte(`\u003c`) + strBackslashQ = []byte(`\u0027`) + strBackslashZero = []byte(`\u0000`) +) diff --git a/lib/results_test.go b/lib/results_test.go index d035d3f8..8df292be 100644 --- a/lib/results_test.go +++ b/lib/results_test.go @@ -2,7 +2,9 @@ package vegeta import ( "bytes" + "encoding/json" "io" + "io/ioutil" "math/rand" "net/http" "reflect" @@ -50,6 +52,16 @@ func TestResultDecoding(t *testing.T) { } func TestResultEncoding(t *testing.T) { + newStdJSONEncoder := func(w io.Writer) Encoder { + enc := json.NewEncoder(w) + return func(r *Result) error { return enc.Encode(r) } + } + + newStdJSONDecoder := func(r io.Reader) Decoder { + dec := json.NewDecoder(r) + return func(r *Result) error { return dec.Decode(r) } + } + for _, tc := range []struct { encoding string enc func(io.Writer) Encoder @@ -61,6 +73,8 @@ func TestResultEncoding(t *testing.T) { {"gob", NewEncoder, NewDecoder}, {"csv", NewCSVEncoder, NewCSVDecoder}, {"json", NewJSONEncoder, NewJSONDecoder}, + {"json-dec-compat", NewJSONEncoder, newStdJSONDecoder}, + {"json-enc-compat", newStdJSONEncoder, NewJSONDecoder}, } { tc := tc t.Run(tc.encoding, func(t *testing.T) { @@ -77,6 +91,8 @@ func TestResultEncoding(t *testing.T) { BytesOut: bsOut, Error: e, Body: body, + Method: "GET", + URL: "http://vegeta.test", Headers: http.Header{"Foo": []string{"bar"}}, } @@ -141,8 +157,7 @@ func BenchmarkResultEncodings(b *testing.B) { {"csv", NewCSVEncoder, NewCSVDecoder}, {"json", NewJSONEncoder, NewJSONDecoder}, } { - var buf bytes.Buffer - enc := tc.enc(&buf) + enc := tc.enc(ioutil.Discard) b.Run(tc.encoding+"-encode", func(b *testing.B) { for i := 0; i < b.N; i++ { @@ -150,10 +165,17 @@ func BenchmarkResultEncodings(b *testing.B) { } }) + var buf bytes.Buffer + enc = tc.enc(&buf) + for _, r := range results { + enc.Encode(&r) + } + dec := tc.dec(&buf) b.Run(tc.encoding+"-decode", func(b *testing.B) { + var r Result for i := 0; i < b.N; i++ { - dec.Decode(&results[i%len(results)]) + dec.Decode(&r) } }) }