From 4bbc7a0296bd12655521f0479aee19cb54c34207 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Senart?= Date: Tue, 10 Sep 2013 21:47:37 +0100 Subject: [PATCH] Split Attack and Report into sub commands --- attack.go | 67 +++++++++++++++++++ main_test.go => attack_test.go | 41 +++++------- main.go | 114 +++++++++------------------------ report.go | 59 +++++++++++++++++ 4 files changed, 170 insertions(+), 111 deletions(-) create mode 100644 attack.go rename main_test.go => attack_test.go (59%) create mode 100644 report.go diff --git a/attack.go b/attack.go new file mode 100644 index 0000000..b10e5ab --- /dev/null +++ b/attack.go @@ -0,0 +1,67 @@ +package main + +import ( + "fmt" + vegeta "github.com/tsenart/vegeta/lib" + "io" + "log" + "os" + "time" +) + +const ( + errRatePrefix = "Rate: " + errDurationPrefix = "Duration: " + errOutputFilePrefix = "Output file: " + errTargetsFilePrefix = "Targets file: " + errOrderingPrefix = "Ordering: " + errReportingPrefix = "Reporting: " +) + +// attack is the command function that validates the attack arguments, sets up the +// required resources, launches the attack and reports the results +func attack(rate uint64, duration time.Duration, targetsf, ordering, output string) error { + if rate == 0 { + return fmt.Errorf(errRatePrefix + "can't be zero") + } + + if duration == 0 { + return fmt.Errorf(errDurationPrefix + "can't be zero") + } + + targets, err := vegeta.NewTargetsFromFile(targetsf) + if err != nil { + return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) + } + + switch ordering { + case "random": + targets.Shuffle(time.Now().UnixNano()) + case "sequential": + break + default: + return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering) + } + + var out io.Writer + switch output { + case "stdout": + out = os.Stdout + default: + file, err := os.Create(output) + if err != nil { + return fmt.Errorf(errOutputFilePrefix+"(%s): %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) + results := vegeta.Attack(targets, rate, duration) + log.Println("Done!") + log.Printf("Writing results to '%s'...", output) + if err := results.WriteTo(out); err != nil { + return err + } + return nil +} diff --git a/main_test.go b/attack_test.go similarity index 59% rename from main_test.go rename to attack_test.go index 9cd75a0..824386d 100644 --- a/main_test.go +++ b/attack_test.go @@ -14,31 +14,31 @@ func init() { } func TestRateValidation(t *testing.T) { - rate, duration, targetsf, ordering, reporter, output := defaultArguments() + rate, duration, targetsf, ordering, output := defaultArguments() rate = 0 - err := run(rate, duration, targetsf, ordering, reporter, output) + err := attack(rate, duration, targetsf, ordering, output) if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errRatePrefix)) { t.Errorf("Rate 0 shouldn't be valid: %s", err) } } func TestDurationValidation(t *testing.T) { - rate, duration, targetsf, ordering, reporter, output := defaultArguments() + rate, duration, targetsf, ordering, output := defaultArguments() duration = 0 - err := run(rate, duration, targetsf, ordering, reporter, output) + err := attack(rate, duration, targetsf, ordering, output) if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errDurationPrefix)) { t.Errorf("Duration 0 shouldn't be valid: %s", err) } } func TestOutputValidation(t *testing.T) { - rate, duration, targetsf, ordering, reporter, _ := defaultArguments() + rate, duration, targetsf, ordering, _ := defaultArguments() // Good cases for _, output := range []string{"stdout", "/dev/null"} { - err := run(rate, duration, targetsf, ordering, reporter, output) + err := attack(rate, duration, targetsf, ordering, output) if err != nil { t.Errorf("Output file `%s` should be valid: %s", output, err) } @@ -46,46 +46,35 @@ func TestOutputValidation(t *testing.T) { // Bad case badOutput := "" - err := run(rate, duration, targetsf, ordering, reporter, badOutput) + err := attack(rate, duration, targetsf, ordering, badOutput) if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errOutputFilePrefix)) { t.Errorf("Output file `%s` shouldn't be valid: %s", badOutput, err) } } -func TestReporter(t *testing.T) { - rate, duration, targetsf, ordering, _, output := defaultArguments() - - for _, reporter := range []string{"text", "plot:timings"} { - err := run(rate, duration, targetsf, ordering, reporter, output) - if err != nil { - t.Errorf("Reporter `%s` shouldn't return an error: %s", reporter, err) - } - } -} - func TestTargetsValidation(t *testing.T) { - rate, duration, goodFile, ordering, reporter, output := defaultArguments() + rate, duration, goodFile, ordering, output := defaultArguments() // Good case - err := run(rate, duration, goodFile, ordering, reporter, output) + err := attack(rate, duration, goodFile, ordering, output) if err != nil { t.Errorf("Targets file `%s` should be valid: %s", goodFile, err) } // Bad case badFile := "randomInexistingFile12345.txt" - err = run(rate, duration, badFile, ordering, reporter, output) + err = attack(rate, duration, badFile, ordering, output) if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errTargetsFilePrefix)) { t.Errorf("Targets file `%s` shouldn't be valid: %s", badFile, err) } } func TestOrderingValidation(t *testing.T) { - rate, duration, targetsf, _, reporter, output := defaultArguments() + rate, duration, targetsf, _, output := defaultArguments() // Good cases for _, ordering := range []string{"random", "sequential"} { - err := run(rate, duration, targetsf, ordering, reporter, output) + err := attack(rate, duration, targetsf, ordering, output) if err != nil { t.Errorf("Ordering `%s` should be valid: %s", ordering, err) } @@ -93,12 +82,12 @@ func TestOrderingValidation(t *testing.T) { // Bad case badOrdering := "lolcat" - err := run(rate, duration, targetsf, badOrdering, reporter, output) + err := attack(rate, duration, targetsf, badOrdering, output) if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errOrderingPrefix)) { t.Errorf("Ordering `%s` shouldn't be valid: %s", badOrdering, err) } } -func defaultArguments() (uint64, time.Duration, string, string, string, string) { - return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "text", "/dev/null" +func defaultArguments() (uint64, time.Duration, string, string, string) { + return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "/dev/null" } diff --git a/main.go b/main.go index 63363f6..31b62c9 100644 --- a/main.go +++ b/main.go @@ -2,105 +2,49 @@ package main import ( "flag" - "fmt" - vegeta "github.com/tsenart/vegeta/lib" - "io" "log" - "os" "runtime" "time" ) 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, json, plot:timings]") - output = flag.String("output", "stdout", "Reporter output file") - cpus = flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use") - ) + // Global flags + cpus := flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use") flag.Parse() - - if flag.NFlag() == 0 { - flag.Usage() - return - } - runtime.GOMAXPROCS(*cpus) - if err := run(*rate, *duration, *targetsf, *ordering, *reporter, *output); err != nil { - log.Fatal(err) + args := flag.Args() + if len(args) < 1 { + log.Fatal("Unspecified command") } -} + cmd, cmdf := args[0], flag.NewFlagSet(args[0], flag.ExitOnError) -const ( - errRatePrefix = "Rate: " - errDurationPrefix = "Duration: " - errOutputFilePrefix = "Output file: " - errTargetsFilePrefix = "Targets file: " - errOrderingPrefix = "Ordering: " - errReportingPrefix = "Reporting: " -) + switch cmd { + case "attack": + rate := cmdf.Uint64("rate", 50, "Requests per second") + targetsf := cmdf.String("targets", "targets.txt", "Targets file") + ordering := cmdf.String("ordering", "random", "Attack ordering [sequential, random]") + duration := cmdf.Duration("duration", 10*time.Second, "Duration of the test") + output := cmdf.String("output", "stdout", "Vegeta Results file") -// run is an utility function that validates the attack arguments, sets up the -// required resources, launches the attack and reports the results -func run(rate uint64, duration time.Duration, targetsf, ordering, reporter, output string) error { - if rate == 0 { - return fmt.Errorf(errRatePrefix + "can't be zero") - } - - if duration == 0 { - return fmt.Errorf(errDurationPrefix + "can't be zero") - } - - var out io.Writer - switch output { - case "stdout": - out = os.Stdout - default: - file, err := os.Create(output) - if err != nil { - return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err) + if err := cmdf.Parse(args[1:]); err != nil { + log.Fatal(err) } - defer file.Close() - out = file - } + if err := attack(*rate, *duration, *targetsf, *ordering, *output); err != nil { + log.Fatal(err) + } + case "report": + reporter := cmdf.String("reporter", "text", "Reporter [text, json, plot:timings]") + input := cmdf.String("input", "stdin", "Vegeta Results file") + output := cmdf.String("output", "stdout", "Output file") - var rep vegeta.Reporter - switch reporter { - case "text": - rep = vegeta.ReportText - case "json": - rep = vegeta.ReportJSON - case "plot:timings": - rep = vegeta.ReportTimingsPlot + if err := cmdf.Parse(args[1:]); err != nil { + log.Fatal(err) + } + if err := report(*reporter, *input, *output); err != nil { + log.Fatal(err) + } default: - log.Println("Reporter provided is not supported. Using text") - rep = vegeta.ReportText + log.Fatalf("Unknown command: %s", cmd) } - - targets, err := vegeta.NewTargetsFromFile(targetsf) - if err != nil { - return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) - } - - switch ordering { - case "random": - targets.Shuffle(time.Now().UnixNano()) - case "sequential": - break - default: - return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering) - } - - log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), ordering, duration) - results := vegeta.Attack(targets, rate, duration) - log.Println("Done!") - log.Printf("Writing report to '%s'...", output) - if err = rep(results, out); err != nil { - return fmt.Errorf(errReportingPrefix+"%s", err) - } - return nil } diff --git a/report.go b/report.go new file mode 100644 index 0000000..f99eb6d --- /dev/null +++ b/report.go @@ -0,0 +1,59 @@ +package main + +import ( + vegeta "github.com/tsenart/vegeta/lib" + "io" + "log" + "os" +) + +func report(reporter, input, output string) error { + var rep vegeta.Reporter + switch reporter { + case "text": + rep = vegeta.ReportText + case "json": + rep = vegeta.ReportJSON + case "plot:timings": + rep = vegeta.ReportTimingsPlot + default: + log.Println("Reporter provided is not supported. Using text") + rep = vegeta.ReportText + } + + var in io.Reader + switch input { + case "stdin": + in = os.Stdin + default: + file, err := os.Open(input) + if err != nil { + return err + } + defer file.Close() + in = file + } + + var out io.Writer + switch output { + case "stdout": + out = os.Stdout + default: + file, err := os.Create(output) + if err != nil { + return err + } + defer file.Close() + out = file + } + + results := vegeta.Results{} + if err := results.ReadFrom(in); err != nil { + return err + } + if err := rep(results, out); err != nil { + return err + } + + return nil +}