func main() {
var (
rate = flag.Uint64("rate", 50, "Requests per second")
targetsf = flag.String("targets", "targets.txt", "Targets file")
ordering = flag.String("ordering", "random", "Attack ordering [sequential, random]")
duration = flag.Duration("duration", 10*time.Second, "Duration of the test")
reporter = flag.String("reporter", "text", "Reporter to use [text, plot:timings]")
output = flag.String("output", "stdout", "Reporter output file")
)
flag.Parse()
if flag.NFlag() == 0 {
flag.Usage()
return
}
if *rate == 0 {
log.Fatal("rate can't be zero")
}
targets, err := vegeta.NewTargetsFromFile(*targetsf)
if err != nil {
log.Fatal(err)
}
switch *ordering {
case "random":
targets.Shuffle(time.Now().UnixNano())
case "sequential":
break
default:
log.Fatalf("Unknown ordering %s", *ordering)
}
if *duration == 0 {
log.Fatal("Duration provided is invalid")
}
var rep vegeta.Reporter
switch *reporter {
case "text":
rep = vegeta.NewTextReporter()
case "plot:timings":
rep = vegeta.NewTimingsPlotReporter()
default:
log.Println("Reporter provided is not supported. using text")
rep = vegeta.NewTextReporter()
}
var out io.Writer
switch *output {
case "stdout":
out = os.Stdout
default:
file, err := os.Create(*output)
if err != nil {
log.Fatalf("Couldn't open `%s` for writing report: %s", *output, err)
}
defer file.Close()
out = file
}
log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), *ordering, *duration)
vegeta.Attack(targets, *rate, *duration, rep)
log.Println("Done!")
log.Printf("Writing report to '%s'...", *output)
if rep.Report(out) != nil {
log.Println("Failed to report!")
}
}
// report validates the report arguments, sets up the required resources
// and writes the report
func report(reporter, inputs, output string) error {
if len(reporter) < 4 {
return fmt.Errorf("bad reporter: %s", reporter)
}
var rep vegeta.Reporter
switch reporter[:4] {
case "text":
rep = vegeta.ReportText
case "json":
rep = vegeta.ReportJSON
case "plot":
rep = vegeta.ReportPlot
case "hist":
if len(reporter) < 6 {
return fmt.Errorf("bad buckets: '%s'", reporter[4:])
}
var hist vegeta.HistogramReporter
if err := hist.Set(reporter[4:len(reporter)]); err != nil {
return err
}
rep = hist
}
files := strings.Split(inputs, ",")
srcs := make([]io.Reader, len(files))
for i, f := range files {
in, err := file(f, false)
if err != nil {
return err
}
defer in.Close()
srcs[i] = in
}
out, err := file(output, true)
if err != nil {
return err
}
defer out.Close()
var results vegeta.Results
res, errs := vegeta.Collect(srcs...)
sig := make(chan os.Signal, 1)
signal.Notify(sig, os.Interrupt)
outer:
for {
select {
case _ = <-sig:
break outer
case r, ok := <-res:
if !ok {
break outer
}
results = append(results, r)
case err, ok := <-errs:
if !ok {
break outer
}
return err
}
}
sort.Sort(results)
data, err := rep.Report(results)
if err != nil {
return err
}
_, err = out.Write(data)
return err
}
// report validates the report arguments, sets up the required resources
// and writes the report
func report(reporter, inputs, output string) error {
if len(reporter) < 4 {
return fmt.Errorf("bad reporter: %s", reporter)
}
files := strings.Split(inputs, ",")
srcs := make([]io.Reader, len(files))
for i, f := range files {
in, err := file(f, false)
if err != nil {
return err
}
defer in.Close()
srcs[i] = in
}
dec := vegeta.NewDecoder(srcs...)
out, err := file(output, true)
if err != nil {
return err
}
defer out.Close()
var (
rep vegeta.Reporter
report vegeta.Report
)
switch reporter[:4] {
case "text":
var m vegeta.Metrics
rep, report = vegeta.NewTextReporter(&m), &m
case "json":
var m vegeta.Metrics
rep, report = vegeta.NewJSONReporter(&m), &m
case "plot":
var rs vegeta.Results
rep, report = vegeta.NewPlotReporter(&rs), &rs
case "hist":
if len(reporter) < 6 {
return fmt.Errorf("bad buckets: '%s'", reporter[4:])
}
var hist vegeta.Histogram
if err := hist.Buckets.UnmarshalText([]byte(reporter[4:])); err != nil {
return err
}
rep, report = vegeta.NewHistogramReporter(&hist), &hist
default:
return fmt.Errorf("unknown reporter: %q", reporter)
}
sigch := make(chan os.Signal, 1)
signal.Notify(sigch, os.Interrupt)
decode:
for {
select {
case <-sigch:
break decode
default:
var r vegeta.Result
if err = dec.Decode(&r); err != nil {
if err == io.EOF {
break decode
}
return err
}
report.Add(&r)
}
}
if c, ok := report.(vegeta.Closer); ok {
c.Close()
}
return rep.Report(out)
}