forked from tsenart/vegeta
-
Notifications
You must be signed in to change notification settings - Fork 0
/
encode.go
139 lines (115 loc) · 2.99 KB
/
encode.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
package main
import (
"bytes"
"flag"
"fmt"
"io"
"os"
"os/signal"
"strings"
vegeta "github.com/tsenart/vegeta/lib"
)
const (
encodingCSV = "csv"
encodingGob = "gob"
encodingJSON = "json"
)
const encodeUsage = `Usage: vegeta encode [options] [<file>...]
Encodes vegeta attack results from one encoding to another.
The supported encodings are Gob (binary), CSV and JSON.
Each input file may have a different encoding which is detected
automatically.
Arguments:
<file> A file with vegeta attack results encoded with one of
the supported encodings (gob | json | csv) [default: stdin]
Options:
--to Output encoding (gob | json | csv) [default: json]
--output Output file [default: stdout]
Examples:
echo "GET http://:80" | vegeta attack -rate=1/s > results.gob
cat results.gob | vegeta encode | jq -c 'del(.body)' | vegeta encode -to gob
`
func encodeCmd() command {
encs := "[" + strings.Join([]string{encodingCSV, encodingGob, encodingJSON}, ", ") + "]"
fs := flag.NewFlagSet("vegeta encode", flag.ExitOnError)
to := fs.String("to", encodingJSON, "Output encoding "+encs)
output := fs.String("output", "stdout", "Output file")
fs.Usage = func() {
fmt.Fprintln(os.Stderr, encodeUsage)
}
return command{fs, func(args []string) error {
fs.Parse(args)
files := fs.Args()
if len(files) == 0 {
files = append(files, "stdin")
}
return encode(files, *to, *output)
}}
}
func encode(files []string, to, output string) error {
srcs := make([]vegeta.Decoder, len(files))
decs := []func(io.Reader) vegeta.Decoder{
vegeta.NewDecoder,
vegeta.NewJSONDecoder,
vegeta.NewCSVDecoder,
}
for i, f := range files {
in, err := file(f, false)
if err != nil {
return err
}
defer in.Close()
// Auto-detect encoding of each file individually and buffer the read bytes
// so that they can be read in subsequent decoding attempts as well as
// in the final decoder.
var buf bytes.Buffer
var dec func(io.Reader) vegeta.Decoder
for j := range decs {
rd := io.MultiReader(bytes.NewReader(buf.Bytes()), io.TeeReader(in, &buf))
if err = decs[j](rd).Decode(&vegeta.Result{}); err == nil {
dec = decs[j]
break
}
}
if dec == nil {
return fmt.Errorf("encode: can't detect encoding of %q", f)
}
srcs[i] = dec(io.MultiReader(&buf, in))
}
dec := vegeta.NewRoundRobinDecoder(srcs...)
out, err := file(output, true)
if err != nil {
return err
}
defer out.Close()
var enc vegeta.Encoder
switch to {
case encodingCSV:
enc = vegeta.NewCSVEncoder(out)
case encodingGob:
enc = vegeta.NewEncoder(out)
case encodingJSON:
enc = vegeta.NewJSONEncoder(out)
default:
return fmt.Errorf("encode: unknown encoding %q", to)
}
sigch := make(chan os.Signal, 1)
signal.Notify(sigch, os.Interrupt)
for {
select {
case <-sigch:
return nil
default:
}
var r vegeta.Result
if err = dec.Decode(&r); err != nil {
if err == io.EOF {
break
}
return err
} else if err = enc.Encode(&r); err != nil {
return err
}
}
return nil
}