// attack validates the attack arguments, sets up the // required resources, launches the attack and writes the results func attack(rate uint64, duration time.Duration, targetsf, ordering, output string, header http.Header) error { if rate == 0 { return fmt.Errorf(errRatePrefix + "can't be zero") } if duration == 0 { return fmt.Errorf(errDurationPrefix + "can't be zero") } in, err := file(targetsf, false) if err != nil { return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) } defer in.Close() targets, err := vegeta.NewTargetsFrom(in) if err != nil { return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) } targets.SetHeader(header) switch ordering { case "random": targets.Shuffle(time.Now().UnixNano()) case "sequential": break default: return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering) } out, err := file(output, true) if err != nil { return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err) } defer out.Close() log.Printf("Vegeta is attacking %d targets in %s order for %s with %d requests/sec...\n", len(targets), ordering, duration, rate) results := vegeta.Attack(targets, rate, duration) log.Println("Done!") log.Printf("Writing results to '%s'...", output) if err := results.Encode(out); err != nil { return err } return nil }
func rapidCmd(args []string) command { fs := flag.NewFlagSet("rapid", flag.ExitOnError) targetsf := fs.String("targets", "stdin", "Targets file") ordering := fs.String("ordering", "random", "Attack ordering [sequential, random]") duration := fs.Duration("duration", 10*time.Second, "Duration of the attack on every rate") output := fs.String("output", "stdout", "Output file") hdrs := headers{Header: make(http.Header)} fs.Var(hdrs, "header", "Targets request header") var rateFlag rateList = []uint64{} fs.Var(&rateFlag, "rates", "One or more rates, comma separated in requests per second") fs.Parse(args) in, err := file(*targetsf, false) if err != nil { return nil } defer in.Close() targets, err := vegeta.NewTargetsFrom(in) if err != nil { return nil } return func() error { results := make(vegeta.Results, 0) var err error = nil for _, rate := range rateFlag { if results, err = rapidAttack(rate, *duration, &targets, *ordering, *output, hdrs.Header, results); err != nil { return err } } if err = writeResults(results, *output); err != nil { return err } return nil } }