├── .github └── workflows │ └── build.yml ├── .gitignore ├── .golangci.yml ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── cmd ├── cmd.go ├── completion.go ├── core │ ├── command.go │ ├── errors │ │ └── errors.go │ ├── exec.go │ ├── extract.go │ ├── flags.go │ ├── input │ │ └── input.go │ ├── options │ │ └── options.go │ ├── runner │ │ ├── concurrency.go │ │ ├── runner.go │ │ └── throughput.go │ └── sync.go ├── dhcpv4 │ ├── cmd.go │ ├── dhcpv4.go │ ├── flags.go │ └── flags_test.go ├── dhcpv6 │ ├── cmd.go │ ├── dhcpv6.go │ ├── flags.go │ └── flags_test.go ├── dns │ ├── cmd.go │ ├── dns.go │ ├── flags.go │ └── flags_test.go ├── http │ ├── cmd.go │ └── http.go ├── tftp │ ├── cmd.go │ └── tftp.go └── udp │ ├── cmd.go │ └── udp.go ├── docs ├── EXTENDING.md ├── PROTOCOLS.md └── USAGE.md ├── flags ├── constraint.go ├── constraint_test.go ├── distribution.go ├── distribution_test.go ├── errors.go ├── flags_test.go ├── growth.go ├── growth_test.go ├── log.go ├── log_test.go ├── utils.go └── utils_test.go ├── go.mod ├── go.sum ├── log ├── log.go └── log_test.go ├── main.go ├── metric ├── errors.go ├── latency.go └── parser.go ├── protocols └── udp │ └── udp.go ├── recorders ├── logrus.go ├── logrus_test.go ├── progress.go ├── progress_test.go ├── statistics.go └── statistics_test.go ├── tester ├── comparator.go ├── comparator_test.go ├── constraint.go ├── constraint_test.go ├── dhcpv4 │ └── tester.go ├── dhcpv6 │ └── tester.go ├── dns │ └── tester.go ├── growth.go ├── growth_test.go ├── http │ └── tester.go ├── metric.go ├── metric_test.go ├── run │ ├── common.go │ ├── common_test.go │ ├── concurrency.go │ ├── concurrency_test.go │ ├── throughput.go │ └── throughput_test.go ├── tester.go ├── tester_test.go ├── tftp │ └── tester.go └── udp │ └── tester.go └── utils ├── completion.go ├── completion_test.go ├── net.go ├── net_test.go ├── random.go ├── random_test.go ├── regexp.go ├── regexp_test.go ├── spinner.go └── spinner_test.go /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Go 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | 7 | build: 8 | name: Test 9 | runs-on: ubuntu-latest 10 | strategy: 11 | matrix: 12 | go-version: ['1.14', '1.15'] 13 | 14 | steps: 15 | - name: Set up Go ${{ matrix.go-version }} 16 | uses: actions/setup-go@v2 17 | with: 18 | go-version: ${{ matrix.go-version }} 19 | 20 | - name: Check out code into the Go module directory 21 | uses: actions/checkout@v2 22 | 23 | - name: Get dependencies 24 | run: go mod download 25 | 26 | - name: Build 27 | run: go build -v ./... 28 | 29 | - name: Test 30 | run: | 31 | for package in $(go list ... | grep -e "github.com/facebookincubator/fbender"); do 32 | go test -race -coverprofile=profile.out -covermode=atomic "$package" 33 | if [ -f profile.out ]; then 34 | cat profile.out >> coverage.txt 35 | rm profile.out 36 | fi 37 | done 38 | 39 | - name: Upload reports to Codecov 40 | run: bash <(curl -s https://codecov.io/bash) 41 | 42 | lint: 43 | name: Lint 44 | runs-on: ubuntu-latest 45 | 46 | steps: 47 | - name: Set up Go 48 | uses: actions/setup-go@v2 49 | with: 50 | go-version: '1.15' 51 | 52 | - name: Check out code into the Go module directory 53 | uses: actions/checkout@v2 54 | 55 | - name: Get dependencies 56 | run: go mod download 57 | 58 | - name: Download Linter 59 | run: curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v1.34.1 60 | 61 | - name: Lint 62 | run: golangci-lint run 63 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Build 2 | fbender 3 | **/profile.out 4 | 5 | # IDE 6 | .idea/ 7 | *.iml 8 | 9 | ### macOS ### 10 | .DS_Store 11 | .AppleDouble 12 | .LSOverride 13 | Icon 14 | 15 | # Thumbnails 16 | ._* 17 | 18 | # Files that might appear on external disk 19 | .Spotlight-V100 20 | .Trashes 21 | 22 | ### Linux ### 23 | !.gitignore 24 | *~ 25 | -------------------------------------------------------------------------------- /.golangci.yml: -------------------------------------------------------------------------------- 1 | run: 2 | deadline: 5m 3 | 4 | linters-settings: 5 | dupl: 6 | threshold: 100 7 | errcheck: 8 | check-type-assertions: true 9 | check-blank: true 10 | govet: 11 | check-shadowing: true 12 | golint: 13 | min-confidence: 0 14 | gocyclo: 15 | min-complecity: 10 16 | 17 | linters: 18 | enable-all: true 19 | disable: 20 | - gosec 21 | - prealloc 22 | - dupl 23 | - gomnd 24 | - exhaustivestruct 25 | - paralleltest 26 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to make participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at . All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 72 | 73 | [homepage]: https://www.contributor-covenant.org 74 | 75 | For answers to common questions about this code of conduct, see 76 | https://www.contributor-covenant.org/faq 77 | 78 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to FBender 2 | We want to make contributing to this project as easy and transparent as 3 | possible. 4 | 5 | ## Pull Requests 6 | We actively welcome your pull requests. 7 | 8 | 1. Fork the repo and create your branch from `main`. 9 | 2. If you've added code that should be tested, add tests. 10 | 3. If you've changed APIs, update the documentation. 11 | 4. Ensure the test suite passes. 12 | 5. Make sure your code lints. 13 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 14 | 15 | ## Contributor License Agreement ("CLA") 16 | In order to accept your pull request, we need you to submit a CLA. You only need 17 | to do this once to work on any of Facebook's open source projects. 18 | 19 | Complete your CLA here: 20 | 21 | ## Issues 22 | We use GitHub issues to track public bugs. Please ensure your description is 23 | clear and has sufficient instructions to be able to reproduce the issue. 24 | 25 | Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe 26 | disclosure of security bugs. In those cases, please go through the process 27 | outlined on that page and do not file a public issue. 28 | 29 | ## Coding Style 30 | All files must be formatted with `gofmt -s -w`. 31 | 32 | ## License 33 | By contributing to FBender, you agree that your contributions will be licensed 34 | under the LICENSE file in the root directory of this source tree. 35 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD License 2 | 3 | For FBender software 4 | 5 | Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. 6 | 7 | Redistribution and use in source and binary forms, with or without modification, 8 | are permitted provided that the following conditions are met: 9 | 10 | * Redistributions of source code must retain the above copyright notice, this 11 | list of conditions and the following disclaimer. 12 | 13 | * Redistributions in binary form must reproduce the above copyright notice, 14 | this list of conditions and the following disclaimer in the documentation 15 | and/or other materials provided with the distribution. 16 | 17 | * Neither the name Facebook nor the names of its contributors may be used to 18 | endorse or promote products derived from this software without specific 19 | prior written permission. 20 | 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 22 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 23 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 25 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 26 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 28 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 30 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FBender 2 | [![Build Status](https://github.com/facebookincubator/fbender/workflows/Go/badge.svg)](https://github.com/facebookincubator/fbender/actions) 3 | [![codecov](https://codecov.io/gh/facebookincubator/fbender/branch/main/graph/badge.svg)](https://codecov.io/gh/facebookincubator/fbender) 4 | [![Go Report Card](https://goreportcard.com/badge/github.com/facebookincubator/fbender)](https://goreportcard.com/report/github.com/facebookincubator/fbender) 5 | 6 | FBender is a __load testing__ command line tool for generic network protocols. 7 | 8 | As a foundation for load testing lays the [Pinterest Bender](https://github.com/pinterest/bender) 9 | library. Similar to Bender, FBender provides two different approaches to load 10 | testing. The first, __Throughput__, gives the tester control over the throughput 11 | (QPS), but not over the concurrency. The second, __Concurrency__, gives the 12 | tester control over the concurrency, but not over the throughput. You can read 13 | more about that in the [Bender documentation](https://github.com/pinterest/bender#bender). 14 | 15 | FBender has been designed to be easily extendable by additional protocols. Look 16 | at the guide on how to contribute new protocols. 17 | 18 | ## Examples 19 | 20 | In the _first example_ we will be load testing a __DNS__ server __example.com__ 21 | running on the __default port__ (53). We will perform 3 consecutive tests for 22 | each __specified QPS__ (2000, 4000, 8000) each lasting for __1 minute__. The 23 | queries will be generated based on the input file __queries.txt__. We will 24 | ignore requests output. 25 | 26 | ```sh 27 | fbender dns throughput fixed \ 28 | --target example.com --duration 1m \ 29 | --input queries.txt -v error \ 30 | 2000 4000 8000 31 | ``` 32 | 33 | In the _next example_ we will be load testing a __TFTP__ server __example.com__ 34 | running on the __default port__ (69). We will perform 3 consecutive tests for 35 | each __specified number of concurrent connections__ (10, 25, 50) each lasting 36 | for __1 minute__. The queries will be generated based on the input file 37 | __files.txt__. We will ignore requests output. 38 | 39 | ```sh 40 | fbender tftp concurrency fixed \ 41 | --target example.com --duration 1m \ 42 | --input files.txt -v error \ 43 | 10 25 50 44 | ``` 45 | 46 | The _last example_ will focus on finding the SLA for a __DHCPv6__ server 47 | __example.com__. We want the timeouts not to exceed __5% of all requests__ in 48 | the measure window of __1 minute__. To get the most accurate results we will be 49 | using exponential backoff growth starting at 20 QPS with a precision of 10 QPS. 50 | The queries will be generated based on the input file __macs.txt__. We will 51 | ignore requests output. 52 | 53 | ```sh 54 | fbender dhcpv6 throughput constraints \ 55 | --target example.com --duration 1m \ 56 | --input macs.txt -v error \ 57 | --constraints "AVG(errors) < 5" \ 58 | --growth ^10 20 59 | ``` 60 | 61 | ## Building FBender 62 | 63 | ```sh 64 | go get -u github.com/facebookincubator/fbender 65 | go build github.com/facebookincubator/fbender 66 | ``` 67 | 68 | ## Installing FBender 69 | 70 | ```sh 71 | go get -u github.com/facebookincubator/fbender 72 | go install github.com/facebookincubator/fbender 73 | ``` 74 | 75 | You may want to add the following line to your .bashrc to enable autocompletion 76 | ```sh 77 | source <(fbender complete bash) 78 | ``` 79 | 80 | ## Docs 81 | 82 | * [General usage guide](https://github.com/facebookincubator/fbender/blob/main/docs/USAGE.md) 83 | * [Protocol specific usage and examples](https://github.com/facebookincubator/fbender/blob/main/docs/PROTOCOLS.md) 84 | * [Extending FBender](https://github.com/facebookincubator/fbender/blob/main/docs/EXTENDING.md) 85 | 86 | ## License 87 | 88 | FBender is BSD licensed, as found in the LICENSE file. 89 | -------------------------------------------------------------------------------- /cmd/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package cmd 10 | 11 | import ( 12 | "fmt" 13 | "os" 14 | "time" 15 | 16 | "github.com/facebookincubator/fbender/cmd/core" 17 | "github.com/facebookincubator/fbender/cmd/dhcpv4" 18 | "github.com/facebookincubator/fbender/cmd/dhcpv6" 19 | "github.com/facebookincubator/fbender/cmd/dns" 20 | "github.com/facebookincubator/fbender/cmd/http" 21 | "github.com/facebookincubator/fbender/cmd/tftp" 22 | "github.com/facebookincubator/fbender/cmd/udp" 23 | "github.com/facebookincubator/fbender/flags" 24 | "github.com/sirupsen/logrus" 25 | "github.com/spf13/cobra" 26 | ) 27 | 28 | // Subcommands are the protocol subcommands. 29 | //nolint:gochecknoglobals 30 | var Subcommands = []*cobra.Command{ 31 | dhcpv4.Command, 32 | dhcpv6.Command, 33 | dns.Command, 34 | http.Command, 35 | tftp.Command, 36 | udp.Command, 37 | } 38 | 39 | // Command is the root command for the CLI. 40 | //nolint:gochecknoglobals 41 | var Command = &cobra.Command{ 42 | Use: "fbender", 43 | Long: `FBender is a load tester tool for various protocols. It provides two different 44 | approaches to load testing: Throughput and Concurrency and each of them can have 45 | either fixed or constraints based test values. Throughput tests give the tester 46 | control over the throughput (QPS), but not over concurrency. The second gives 47 | the user control over the concurrency but not over the throughput. 48 | * fixed - runs a single test for each of the specified values. 49 | * constraint - runs tests adjusting load based on the growth and constraints. 50 | 51 | Target: 52 | Target format may vary depending on the protocol, however most of them accept 53 | ipv4, ipv6, hostname with an optional port. Use "fbender protocol --help" to get 54 | the documentation on the target format for a specific protocol. 55 | 56 | Input: 57 | Unless explicitly stated in the command documentation one request is generated 58 | per input line, skipping the lines with improper format. Use "fbender 59 | protocol help" to get the documentation on the input format for a specific 60 | protocol. The generated requests are then reused in a round-robin manner. 61 | 62 | Output: 63 | All important information is printed to stdout. Test logs can be redirected 64 | using the output flag. They can also be filtered based on the message verbosity 65 | level. Note that this filters/redirect only test logs and not the summary and 66 | other output. Available levels (both numbers and literals are accepted): 67 | * panic/0 68 | * fatal/1 69 | * error/2 70 | * warning/3 - log when an *error response* is received 71 | * info/4 - log when a *successful response* is received 72 | * debug/5 - log when a *request* is sent 73 | `, 74 | Example: ` fbender dns throughput fixed -t $TARGET 100 75 | fbender tftp concurrency fixed -t $TARGET -o /dev/null 10 76 | fbender udp throughput fixed -t $TARGET -d 5m 100 200 300 77 | fbender http concurrency constraints -t $TARGET 20 -c "MAX(errors)<5" 78 | fbender dhcpv6 throughput constraints -t $TARGET 50 -c "MIN(latency)<20" 79 | fbender dns throughput constraints -t $TARGET 40 -c -g ^10 "MAX(errors)<5"`, 80 | } 81 | 82 | func initIOFlags() { 83 | // Input 84 | Command.PersistentFlags().StringP("input", "i", "", "load test input data from a file (default )") 85 | 86 | if err := Command.MarkPersistentFlagFilename("input"); err != nil { 87 | panic(err) 88 | } 89 | 90 | // Output 91 | logOutput := flags.NewLogOutput(logrus.StandardLogger()) 92 | 93 | Command.PersistentFlags().VarP(logOutput, "output", "o", "log test output to a file") 94 | 95 | if err := Command.MarkPersistentFlagFilename("output"); err != nil { 96 | panic(err) 97 | } 98 | 99 | // Log Level 100 | logLevel := &flags.LogLevel{Logger: logrus.StandardLogger()} 101 | logLevelChoices := flags.ChoicesString(flags.LogLevelChoices()) 102 | 103 | Command.PersistentFlags().VarP(logLevel, "verbosity", "v", fmt.Sprintf("verbosity level %s", logLevelChoices)) 104 | 105 | if err := flags.BashCompletionLogLevel(Command, Command.PersistentFlags(), "verbosity"); err != nil { 106 | panic(err) 107 | } 108 | 109 | // Log format 110 | logFormat := &flags.LogFormat{Logger: logrus.StandardLogger(), Format: "json"} 111 | logFormatChoices := flags.ChoicesString(flags.LogFormatChoices()) 112 | 113 | Command.PersistentFlags().VarP(logFormat, "format", "f", fmt.Sprintf("output format %s", logFormatChoices)) 114 | 115 | if err := flags.BashCompletionLogFormat(Command, Command.PersistentFlags(), "format"); err != nil { 116 | panic(err) 117 | } 118 | } 119 | 120 | func initExecutionFlags() { 121 | // Test duration 122 | Command.PersistentFlags().DurationP("duration", "d", 1*time.Minute, "single test duration") 123 | 124 | // Requests distribution 125 | distribution := flags.NewDefaultDistribution() 126 | distributionChoices := flags.ChoicesString(flags.DistributionChoices()) 127 | 128 | Command.PersistentFlags().VarP(distribution, "dist", "D", fmt.Sprintf("requests distribution %s", distributionChoices)) 129 | 130 | if err := flags.BashCompletionDistribution(Command, Command.PersistentFlags(), "dist"); err != nil { 131 | panic(err) 132 | } 133 | 134 | // Other settings 135 | Command.PersistentFlags().IntP("buffer", "b", 2048, "buffer size of the requests generator channel") 136 | Command.PersistentFlags().DurationP("timeout", "w", 1*time.Second, "wait timeout on requests") 137 | Command.PersistentFlags().DurationP("unit", "u", 1*time.Millisecond, "histogram scaling unit") 138 | Command.PersistentFlags().Bool("nostats", false, "disable statistics") 139 | } 140 | 141 | //nolint:gochecknoinits 142 | func init() { 143 | cobra.EnablePrefixMatching = true 144 | 145 | initIOFlags() 146 | initExecutionFlags() 147 | 148 | for _, subcommand := range Subcommands { 149 | Command.AddCommand(subcommand) 150 | subcommand.PersistentFlags().StringP("target", "t", "", "endpoint to load test") 151 | 152 | if err := subcommand.MarkPersistentFlagRequired("target"); err != nil { 153 | panic(err) 154 | } 155 | } 156 | 157 | Command.AddCommand(completionCmd) 158 | core.StartPostInit() 159 | } 160 | 161 | // Execute runs the Command. 162 | func Execute() { 163 | if err := Command.Execute(); err != nil { 164 | os.Exit(1) 165 | } 166 | } 167 | -------------------------------------------------------------------------------- /cmd/completion.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package cmd 10 | 11 | import ( 12 | "os" 13 | 14 | "github.com/spf13/cobra" 15 | ) 16 | 17 | //nolint:gochecknoglobals 18 | var completionCmd = &cobra.Command{ 19 | Use: "completion", 20 | Short: "Generates completion scripts", 21 | } 22 | 23 | //nolint:gochecknoglobals 24 | var bashCompletionCmd = &cobra.Command{ 25 | Use: "bash", 26 | Short: "Generates bash completion scripts", 27 | Long: `To load completion run 28 | 29 | . <(fbender completion) 30 | 31 | To configure your bash shell to load completions for each session add to .bashrc 32 | 33 | # ~/.bashrc or ~/.profile 34 | . <(fbender completion) 35 | `, 36 | RunE: func(cmd *cobra.Command, args []string) error { 37 | return Command.GenBashCompletion(os.Stdout) 38 | }, 39 | } 40 | 41 | //nolint:gochecknoinits 42 | func init() { 43 | completionCmd.AddCommand(bashCompletionCmd) 44 | } 45 | -------------------------------------------------------------------------------- /cmd/core/command.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package core 10 | 11 | import ( 12 | "fmt" 13 | "strings" 14 | 15 | "github.com/facebookincubator/fbender/cmd/core/errors" 16 | "github.com/facebookincubator/fbender/cmd/core/options" 17 | "github.com/facebookincubator/fbender/cmd/core/runner" 18 | "github.com/spf13/cobra" 19 | ) 20 | 21 | // CommandParams is used to generate params for the runner. 22 | type CommandParams func(*cobra.Command, *options.Options) (*runner.Params, error) 23 | 24 | // CommandTemplate groups parameters used to generate the test command. 25 | type CommandTemplate struct { 26 | // Command name it will be invoked with 27 | Name string 28 | // Short and long help messages 29 | Short, Long string 30 | // Usage examples can containt {test} string which will be replaced with the 31 | // actual test name (throughput|concurrency) 32 | Fixed, Constraints string 33 | } 34 | 35 | // NewTestCommand creates a protocol command with all the test subcommands. 36 | // ${protocol} throughput fixed 37 | // ${protocol} throughput constraints 38 | // ${protocol} concurrency fixed 39 | // ${protocol} concurrency constraints 40 | //nolint:funlen 41 | func NewTestCommand(c *CommandTemplate, p CommandParams) *cobra.Command { 42 | // Help messages 43 | var ( 44 | tShort = fmt.Sprintf("%s throughput (QPS)", c.Short) 45 | tfShort = fmt.Sprintf("%s with fixed amount of QPS", tShort) 46 | tcShort = fmt.Sprintf("%s with constraints", tShort) 47 | 48 | cShort = fmt.Sprintf("%s concurrent connections", c.Short) 49 | cfShort = fmt.Sprintf("%s with fixed number of connections", cShort) 50 | ccShort = fmt.Sprintf("%s with constraints", cShort) 51 | ) 52 | 53 | // Examples 54 | var ( 55 | tfExamples = strings.ReplaceAll(c.Fixed, "{test}", "throughput") 56 | tcExamples = strings.ReplaceAll(c.Constraints, "{test}", "throughput") 57 | tExamples = fmt.Sprintf("%s\n%s", tfExamples, tcExamples) 58 | 59 | cfExamples = strings.ReplaceAll(c.Fixed, "{test}", "concurrency") 60 | ccExamples = strings.ReplaceAll(c.Constraints, "{test}", "concurrency") 61 | cExamples = fmt.Sprintf("%s\n%s", cfExamples, ccExamples) 62 | 63 | examples = fmt.Sprintf("%s\n%s", tExamples, cExamples) 64 | ) 65 | 66 | // Top level command 67 | command := &cobra.Command{ 68 | Use: c.Name, 69 | Short: c.Short, 70 | Long: fmt.Sprintf("%s.\n%s", c.Short, c.Long), 71 | Example: examples, 72 | } 73 | 74 | // Subcommands 75 | tCommand := &cobra.Command{ 76 | Use: "throughput", 77 | Short: tShort, 78 | Long: fmt.Sprintf("%s.\n%s", tShort, c.Long), 79 | Example: tExamples, 80 | } 81 | 82 | cCommand := &cobra.Command{ 83 | Use: "concurrency", 84 | Short: cShort, 85 | Long: fmt.Sprintf("%s.\n%s", cShort, c.Long), 86 | Example: cExamples, 87 | } 88 | 89 | command.AddCommand(tCommand) 90 | command.AddCommand(cCommand) 91 | 92 | // Throughput subcommands 93 | tfCommand := &cobra.Command{ 94 | Use: "fixed", 95 | Short: tfShort, 96 | Long: fmt.Sprintf("%s.\n%s", tfShort, c.Long), 97 | Example: tfExamples, 98 | Args: fixedArgs, 99 | RunE: RunLoadTestThroughputFixed(p), 100 | } 101 | 102 | tcCommand := &cobra.Command{ 103 | Use: "constraints", 104 | Short: tcShort, 105 | Long: fmt.Sprintf("%s.\n%s\n%s", tcShort, c.Long, ConstraintsHelp), 106 | Example: tcExamples, 107 | Args: constraintsArgs, 108 | RunE: RunLoadTestThroughputConstraints(p), 109 | } 110 | 111 | tcCommand.PersistentFlags().AddFlagSet(ConstraintsFlags) 112 | tCommand.AddCommand(tfCommand) 113 | tCommand.AddCommand(tcCommand) 114 | 115 | // Concurrency subcommands 116 | cfCommand := &cobra.Command{ 117 | Use: "fixed", 118 | Short: cfShort, 119 | Long: fmt.Sprintf("%s.\n%s", cfShort, c.Long), 120 | Example: cfExamples, 121 | Args: fixedArgs, 122 | RunE: RunLoadTestConcurrencyFixed(p), 123 | } 124 | 125 | ccCommand := &cobra.Command{ 126 | Use: "constraints", 127 | Short: ccShort, 128 | Long: fmt.Sprintf("%s.\n%s\n%s", ccShort, c.Long, ConstraintsHelp), 129 | Example: ccExamples, 130 | Args: constraintsArgs, 131 | RunE: RunLoadTestConcurrencyConstraints(p), 132 | } 133 | 134 | ccCommand.PersistentFlags().AddFlagSet(ConstraintsFlags) 135 | cCommand.AddCommand(cfCommand) 136 | cCommand.AddCommand(ccCommand) 137 | 138 | return command 139 | } 140 | 141 | // fixedArgs validates arguments for a fixed QPS test. 142 | func fixedArgs(cmd *cobra.Command, args []string) error { 143 | if len(args) < 1 { 144 | return fmt.Errorf("%w: requires at least one test value", errors.ErrInvalidArgument) 145 | } 146 | 147 | _, err := ExtractTests(args) 148 | 149 | return err 150 | } 151 | 152 | // constraintsArgs validates arguments for a constraints tests. 153 | func constraintsArgs(cmd *cobra.Command, args []string) error { 154 | if len(args) != 1 { 155 | return fmt.Errorf("%w: requires starting test value", errors.ErrInvalidArgument) 156 | } 157 | 158 | _, err := ExtractTests(args) 159 | 160 | return err 161 | } 162 | -------------------------------------------------------------------------------- /cmd/core/errors/errors.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package errors 10 | 11 | import "errors" 12 | 13 | // ErrInvalidFormat is raised when the input does not match the desired format. 14 | var ErrInvalidFormat = errors.New("invalid format") 15 | 16 | // ErrInvalidType is raised when provided object is not of the desired type. 17 | var ErrInvalidType = errors.New("invalid type") 18 | 19 | // ErrInvalidArgument is raised when the command is given invalid arguments. 20 | var ErrInvalidArgument = errors.New("invalid argument") 21 | -------------------------------------------------------------------------------- /cmd/core/exec.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package core 10 | 11 | import ( 12 | "os" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/options" 15 | "github.com/facebookincubator/fbender/cmd/core/runner" 16 | "github.com/facebookincubator/fbender/log" 17 | "github.com/facebookincubator/fbender/tester/run" 18 | "github.com/spf13/cobra" 19 | ) 20 | 21 | // cobraRunE is just an alias for function type which implements cobra.RunE. 22 | type cobraRunE = func(cmd *cobra.Command, args []string) error 23 | 24 | // executor invokes actual test function with proper params. 25 | type executor func(p *runner.Params, o *options.Options) error 26 | 27 | func exec(p CommandParams, e executor, gs ...OptionsGenerator) cobraRunE { 28 | return func(cmd *cobra.Command, args []string) error { 29 | o, err := GenerateOptions(cmd, args, gs...) 30 | if err != nil { 31 | return err 32 | } 33 | 34 | params, err := p(cmd, o) 35 | if err != nil { 36 | return err 37 | } 38 | 39 | // We want runtime errors to be logged and not trigger help message 40 | if err := e(params, o); err != nil { 41 | log.Errorf("Error: %v\n", err) 42 | os.Exit(1) 43 | } 44 | 45 | return nil 46 | } 47 | } 48 | 49 | func setupConstraints(o *options.Options, cmd *cobra.Command, args []string) (*options.Options, error) { 50 | for _, constraint := range o.Constraints { 51 | if err := constraint.Metric.Setup(o); err != nil { 52 | //nolint:wrapcheck 53 | return nil, err 54 | } 55 | } 56 | 57 | return o, nil 58 | } 59 | 60 | //nolint:gochecknoglobals 61 | var fixedOptionsGenerators = []OptionsGenerator{ 62 | ExtractArgs, 63 | ExtractOptions, 64 | } 65 | 66 | //nolint:gochecknoglobals 67 | var constraintsOptionsGenerators = []OptionsGenerator{ 68 | ExtractArgs, 69 | ExtractOptions, 70 | ExtractConstraintsOptions, 71 | setupConstraints, 72 | } 73 | 74 | // RunLoadTestThroughputFixed returns a new cobra RunE method for the load 75 | // tester with fixed QPS tests. 76 | func RunLoadTestThroughputFixed(p CommandParams) cobraRunE { 77 | return exec(p, fixedThroughputExecutor, fixedOptionsGenerators...) 78 | } 79 | 80 | func fixedThroughputExecutor(p *runner.Params, o *options.Options) error { 81 | return run.LoadTestThroughputFixed(runner.NewThroughputRunner(p), o, o.Tests...) 82 | } 83 | 84 | // RunLoadTestThroughputConstraints returns a new cobra RunE method for the QPS 85 | // load tester with constraint checks. 86 | func RunLoadTestThroughputConstraints(p CommandParams) cobraRunE { 87 | return exec(p, constraintsThroughputExecutor, constraintsOptionsGenerators...) 88 | } 89 | 90 | func constraintsThroughputExecutor(p *runner.Params, o *options.Options) error { 91 | return run.LoadTestThroughputConstraints(runner.NewThroughputRunner(p), o, o.Start, o.Growth, o.Constraints...) 92 | } 93 | 94 | // RunLoadTestConcurrencyFixed returns a new cobra RunE method for the load 95 | // tester with fixed concurrent connections count. 96 | func RunLoadTestConcurrencyFixed(p CommandParams) cobraRunE { 97 | return exec(p, fixedConcurrencyExecutor, fixedOptionsGenerators...) 98 | } 99 | 100 | func fixedConcurrencyExecutor(p *runner.Params, o *options.Options) error { 101 | return run.LoadTestConcurrencyFixed(runner.NewConcurrencyRunner(p), o, o.Tests...) 102 | } 103 | 104 | // RunLoadTestConcurrencyConstraints returns a new cobra RunE method for the 105 | // concurrency load tester with constraint checks. 106 | func RunLoadTestConcurrencyConstraints(p CommandParams) cobraRunE { 107 | return exec(p, constraintsConcurrencyExecutor, constraintsOptionsGenerators...) 108 | } 109 | 110 | func constraintsConcurrencyExecutor(p *runner.Params, o *options.Options) error { 111 | return run.LoadTestConcurrencyConstraints(runner.NewConcurrencyRunner(p), o, o.Start, o.Growth, o.Constraints...) 112 | } 113 | -------------------------------------------------------------------------------- /cmd/core/extract.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package core 10 | 11 | import ( 12 | "strconv" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/options" 15 | "github.com/facebookincubator/fbender/flags" 16 | "github.com/spf13/cobra" 17 | ) 18 | 19 | // OptionsGenerator is used to generate options from command line params. 20 | type OptionsGenerator func(o *options.Options, cmd *cobra.Command, args []string) (*options.Options, error) 21 | 22 | // ExtractTests parses list of arguments to a list of int values. 23 | func ExtractTests(args []string) ([]int, error) { 24 | values := make([]int, 0) 25 | 26 | for _, arg := range args { 27 | value, err := strconv.Atoi(arg) 28 | if err != nil { 29 | //nolint:wrapcheck 30 | return nil, err 31 | } 32 | 33 | values = append(values, value) 34 | } 35 | 36 | return values, nil 37 | } 38 | 39 | // ExtractArgs extracts arguments commonly used options across all tests. 40 | func ExtractArgs(o *options.Options, cmd *cobra.Command, args []string) (*options.Options, error) { 41 | var err error 42 | 43 | if o == nil { 44 | o = options.NewOptions() 45 | } 46 | 47 | o.Tests, err = ExtractTests(args) 48 | if err != nil { 49 | return nil, err 50 | } 51 | 52 | o.Start = o.Tests[0] 53 | 54 | return o, nil 55 | } 56 | 57 | // ExtractOptions extracts flags commonly used options across all commands. 58 | func ExtractOptions(o *options.Options, cmd *cobra.Command, _ []string) (*options.Options, error) { 59 | var err error 60 | 61 | if o == nil { 62 | o = options.NewOptions() 63 | } 64 | 65 | o.Target, err = cmd.Flags().GetString("target") 66 | if err != nil { 67 | //nolint:wrapcheck 68 | return nil, err 69 | } 70 | 71 | o.Duration, err = cmd.Flags().GetDuration("duration") 72 | if err != nil { 73 | //nolint:wrapcheck 74 | return nil, err 75 | } 76 | 77 | o.Input, err = cmd.Flags().GetString("input") 78 | if err != nil { 79 | //nolint:wrapcheck 80 | return nil, err 81 | } 82 | 83 | o.BufferSize, err = cmd.Flags().GetInt("buffer") 84 | if err != nil { 85 | //nolint:wrapcheck 86 | return nil, err 87 | } 88 | 89 | o.Timeout, err = cmd.Flags().GetDuration("timeout") 90 | if err != nil { 91 | //nolint:wrapcheck 92 | return nil, err 93 | } 94 | 95 | o.Distribution, err = flags.GetDistribution(cmd.Flags(), "dist") 96 | if err != nil { 97 | //nolint:wrapcheck 98 | return nil, err 99 | } 100 | 101 | o.Unit, err = cmd.Flags().GetDuration("unit") 102 | if err != nil { 103 | //nolint:wrapcheck 104 | return nil, err 105 | } 106 | 107 | o.NoStatistics, err = cmd.Flags().GetBool("nostats") 108 | if err != nil { 109 | //nolint:wrapcheck 110 | return nil, err 111 | } 112 | 113 | return o, nil 114 | } 115 | 116 | // ExtractConstraintsOptions extracts flag commonly used options across constraints test commands. 117 | func ExtractConstraintsOptions(o *options.Options, cmd *cobra.Command, _ []string) (*options.Options, error) { 118 | var err error 119 | 120 | if o == nil { 121 | o = options.NewOptions() 122 | } 123 | 124 | o.Constraints, err = flags.GetConstraints(cmd.Flags(), "constraints") 125 | if err != nil { 126 | //nolint:wrapcheck 127 | return nil, err 128 | } 129 | 130 | o.Growth, err = flags.GetGrowth(cmd.Flags(), "growth") 131 | if err != nil { 132 | //nolint:wrapcheck 133 | return nil, err 134 | } 135 | 136 | return o, nil 137 | } 138 | 139 | // GenerateOptions runs given generators for a command and returns options. 140 | func GenerateOptions(cmd *cobra.Command, args []string, gs ...OptionsGenerator) (*options.Options, error) { 141 | var o *options.Options 142 | 143 | var err error 144 | 145 | for _, g := range gs { 146 | o, err = g(o, cmd, args) 147 | if err != nil { 148 | return nil, err 149 | } 150 | } 151 | 152 | return o, nil 153 | } 154 | -------------------------------------------------------------------------------- /cmd/core/flags.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package core 10 | 11 | import ( 12 | "strings" 13 | 14 | "github.com/facebookincubator/fbender/flags" 15 | "github.com/facebookincubator/fbender/metric" 16 | "github.com/facebookincubator/fbender/tester" 17 | "github.com/spf13/pflag" 18 | ) 19 | 20 | //nolint:gochecknoglobals 21 | var ( 22 | // ConstraintsFlags contains flags for specifying constraints tests options. 23 | ConstraintsFlags = pflag.NewFlagSet("Constraints test flags", pflag.ExitOnError) 24 | // ConstraintsValue is a pflag value for constraints. 25 | ConstraintsValue = flags.NewConstraintSliceValue(metric.Parser) 26 | // ConstraintsHelp is a help message on how to use constraints. 27 | ConstraintsHelp = strings.Join([]string{tester.ConstraintsHelp, metric.Help}, "\n") 28 | ) 29 | 30 | //nolint:gochecknoinits 31 | func init() { 32 | growth := &flags.GrowthValue{Growth: &tester.PercentageGrowth{Increase: 100.}} 33 | 34 | ConstraintsFlags.VarP(ConstraintsValue, "constraints", "c", "constraints to be checked after each test") 35 | ConstraintsFlags.VarP(growth, "growth", "g", "growth used to determinate the next test (+AMOUNT|%PERCENT|^PRECISION)") 36 | } 37 | -------------------------------------------------------------------------------- /cmd/core/input/input.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package input 10 | 11 | import ( 12 | "bufio" 13 | "fmt" 14 | "io" 15 | "os" 16 | 17 | "github.com/facebookincubator/fbender/cmd/core/errors" 18 | "github.com/facebookincubator/fbender/cmd/core/runner" 19 | "github.com/facebookincubator/fbender/log" 20 | ) 21 | 22 | // Transformer converts input line into a request. 23 | type Transformer func(string) (interface{}, error) 24 | 25 | // Modifier changes request right before sending. 26 | type Modifier func(interface{}) (interface{}, error) 27 | 28 | // NewRequestGenerator reads data from the specified input and converts it into 29 | // requests using given transformer. The lines which aren't formatted correctly 30 | // are skipped. The requests are then reused in a round-robin manner inside the 31 | // generator. If modifiers are provided they are applied to the request every 32 | // time just before being returned. 33 | func NewRequestGenerator(filename string, transformer Transformer, mods ...Modifier) (runner.RequestGenerator, error) { 34 | file, err := open(filename) 35 | if err != nil { 36 | return nil, err 37 | } 38 | 39 | defer close(file) 40 | 41 | data := parse(file, transformer) 42 | if len(data) < 1 { 43 | return nil, fmt.Errorf("%w: at least one valid input line is required", errors.ErrInvalidFormat) 44 | } 45 | 46 | return func(i int) interface{} { 47 | var err error 48 | 49 | request := data[i%len(data)] 50 | for _, mod := range mods { 51 | request, err = mod(request) 52 | if err != nil { 53 | panic(err) 54 | } 55 | } 56 | 57 | return request 58 | }, nil 59 | } 60 | 61 | // parse reads data from the specified input and converts it into requests 62 | // using given transformer. The lines which are not formatted correctly are 63 | // skipped and a warning message is printed to stderr. 64 | func parse(input io.Reader, transformer Transformer) []interface{} { 65 | requests := make([]interface{}, 0) 66 | scanner := bufio.NewScanner(input) 67 | 68 | for scanner.Scan() { 69 | line := scanner.Text() 70 | request, err := transformer(line) 71 | 72 | if err != nil { 73 | log.Errorf("Warning: Error parsing input line %q: %v\n", line, err) 74 | } else { 75 | requests = append(requests, request) 76 | } 77 | } 78 | 79 | return requests 80 | } 81 | 82 | func open(filename string) (*os.File, error) { 83 | if len(filename) == 0 { 84 | log.Errorf("Reading input lines until EOF:\n") 85 | 86 | return os.Stdin, nil 87 | } 88 | 89 | return os.Open(filename) 90 | } 91 | 92 | func close(file io.Closer) { 93 | if file == os.Stdin { 94 | return 95 | } 96 | 97 | if err := file.Close(); err != nil { 98 | log.Errorf("Warning: Error closing input file: %v\n", err) 99 | } 100 | } 101 | -------------------------------------------------------------------------------- /cmd/core/options/options.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package options 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/tester" 15 | "github.com/pinterest/bender" 16 | ) 17 | 18 | // Options represents common options for the Commands. 19 | type Options struct { 20 | Target string 21 | Duration time.Duration 22 | Tests []int 23 | Start int 24 | 25 | Input string 26 | 27 | BufferSize int 28 | Timeout time.Duration 29 | Distribution func(float64) bender.IntervalGenerator 30 | Unit time.Duration 31 | NoStatistics bool 32 | 33 | Constraints []*tester.Constraint 34 | Growth tester.Growth 35 | 36 | Recorders []bender.Recorder 37 | } 38 | 39 | // NewOptions returns new options. 40 | func NewOptions() *Options { 41 | return &Options{ 42 | Tests: []int{}, 43 | Constraints: []*tester.Constraint{}, 44 | Recorders: []bender.Recorder{}, 45 | } 46 | } 47 | 48 | // GetUnit returns a unit used in tests. 49 | func (o *Options) GetUnit() time.Duration { 50 | return o.Unit 51 | } 52 | 53 | // AddRecorder adds a recorder to options. 54 | func (o *Options) AddRecorder(recorder bender.Recorder) { 55 | o.Recorders = append(o.Recorders, recorder) 56 | } 57 | -------------------------------------------------------------------------------- /cmd/core/runner/concurrency.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package runner 10 | 11 | import ( 12 | "context" 13 | "time" 14 | 15 | "github.com/facebookincubator/fbender/cmd/core/options" 16 | "github.com/facebookincubator/fbender/recorders" 17 | "github.com/facebookincubator/fbender/tester" 18 | "github.com/facebookincubator/fbender/utils" 19 | "github.com/pinterest/bender" 20 | ) 21 | 22 | // ConcurrencyRunner is a test runner for load test concurrency commands. 23 | type ConcurrencyRunner struct { 24 | runner 25 | workerSem *bender.WorkerSemaphore 26 | spinnerCancel context.CancelFunc 27 | } 28 | 29 | // NewConcurrencyRunner returns new ConcurrencyRunner. 30 | func NewConcurrencyRunner(params *Params) *ConcurrencyRunner { 31 | return &ConcurrencyRunner{ 32 | runner: runner{ 33 | Params: params, 34 | }, 35 | } 36 | } 37 | 38 | // Before prepares requests, recorders and interval generator. 39 | func (r *ConcurrencyRunner) Before(workers tester.Workers, opts interface{}) error { 40 | if err := r.runner.Before(workers, opts); err != nil { 41 | return err 42 | } 43 | 44 | o, ok := opts.(*options.Options) 45 | if !ok { 46 | return tester.ErrInvalidOptions 47 | } 48 | 49 | r.workerSem = bender.NewWorkerSemaphore() 50 | 51 | go func() { r.workerSem.Signal(workers) }() 52 | 53 | r.requests = make(chan interface{}) 54 | 55 | ctx, cancel := context.WithCancel(context.Background()) 56 | 57 | go func() { 58 | for i := 0; ; i++ { 59 | select { 60 | case <-ctx.Done(): 61 | close(r.requests) 62 | 63 | return 64 | default: 65 | r.requests <- r.Params.RequestGenerator(i) 66 | } 67 | } 68 | }() 69 | 70 | // We want tne progressbar to measure the time passed. 71 | const scale = 10 72 | count := int(o.Duration/time.Second) * scale 73 | 74 | r.progress, r.bar = recorders.NewLoadTestProgress(count) 75 | r.progress.Start() 76 | 77 | go func() { 78 | for i := 0; i < count; i++ { 79 | time.Sleep(time.Second / scale) 80 | r.bar.Incr() 81 | } 82 | cancel() 83 | r.progress.Stop() 84 | r.spinnerCancel = utils.NewBackgroundSpinner("Waiting for tests to finish", 0) 85 | }() 86 | 87 | return nil 88 | } 89 | 90 | // After cleans up after the test. 91 | func (r *ConcurrencyRunner) After(test int, options interface{}) { 92 | r.spinnerCancel() 93 | r.runner.After(test, options) 94 | } 95 | 96 | // WorkerSemaphore returns a worker semaphore for concurrency test. 97 | func (r *ConcurrencyRunner) WorkerSemaphore() *bender.WorkerSemaphore { 98 | return r.workerSem 99 | } 100 | -------------------------------------------------------------------------------- /cmd/core/runner/runner.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package runner 10 | 11 | import ( 12 | "runtime" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/options" 15 | "github.com/facebookincubator/fbender/log" 16 | "github.com/facebookincubator/fbender/recorders" 17 | "github.com/facebookincubator/fbender/tester" 18 | "github.com/facebookincubator/fbender/utils" 19 | "github.com/gosuri/uiprogress" 20 | "github.com/pinterest/bender" 21 | "github.com/pinterest/bender/hist" 22 | "github.com/sirupsen/logrus" 23 | ) 24 | 25 | // RequestGenerator is used to generate requests. 26 | type RequestGenerator func(i int) interface{} 27 | 28 | // Params represents test parameters for the runner. 29 | type Params struct { 30 | Tester tester.Tester 31 | RequestGenerator RequestGenerator 32 | } 33 | 34 | // runner groups fields used in both runners. 35 | type runner struct { 36 | requests chan interface{} 37 | recorder chan interface{} 38 | 39 | recorders []bender.Recorder 40 | histogram *hist.Histogram 41 | progress *uiprogress.Progress 42 | bar *uiprogress.Bar 43 | 44 | Params *Params 45 | } 46 | 47 | // reset "frees" all the runner fields. 48 | func (r *runner) reset() { 49 | r.requests = nil 50 | r.recorder = nil 51 | r.recorders = nil 52 | r.histogram = nil 53 | r.progress = nil 54 | r.bar = nil 55 | } 56 | 57 | // Before initializes all common fields. 58 | func (r *runner) Before(test int, opts interface{}) error { 59 | o, ok := opts.(*options.Options) 60 | if !ok { 61 | return tester.ErrInvalidOptions 62 | } 63 | 64 | cancel := utils.NewBackgroundSpinner("Cleaning up the memory", 0) 65 | 66 | r.reset() 67 | runtime.GC() 68 | cancel() 69 | 70 | cancel = utils.NewBackgroundSpinner("Preparing the test", 0) 71 | 72 | r.recorder = make(chan interface{}, o.BufferSize) 73 | r.recorders = []bender.Recorder{ 74 | recorders.NewLogrusRecorder(logrus.StandardLogger(), logrus.Fields{"test": test}), 75 | } 76 | 77 | r.recorders = append(r.recorders, o.Recorders...) 78 | 79 | if !o.NoStatistics { 80 | r.histogram = hist.NewHistogram(2*int(o.Timeout/o.Unit), int(o.Unit)) 81 | r.recorders = append(r.recorders, bender.NewHistogramRecorder(r.histogram)) 82 | } 83 | 84 | cancel() 85 | 86 | log.Printf("Running test: %d\n", test) 87 | 88 | return nil 89 | } 90 | 91 | // After cleans up after the test. 92 | func (r *runner) After(test int, options interface{}) { 93 | if r.histogram != nil { 94 | log.Printf("%s", r.histogram.String()) 95 | } 96 | } 97 | 98 | // Tester returns the protocol tester. 99 | func (r *runner) Tester() tester.Tester { 100 | return r.Params.Tester 101 | } 102 | 103 | // Requests returns the requests channel. 104 | func (r *runner) Requests() chan interface{} { 105 | return r.requests 106 | } 107 | 108 | // Recorder returns the recorder. 109 | func (r *runner) Recorder() chan interface{} { 110 | return r.recorder 111 | } 112 | 113 | // Recorders returns a list of recorders. 114 | func (r *runner) Recorders() []bender.Recorder { 115 | return r.recorders 116 | } 117 | -------------------------------------------------------------------------------- /cmd/core/runner/throughput.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package runner 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/options" 15 | "github.com/facebookincubator/fbender/recorders" 16 | "github.com/facebookincubator/fbender/tester" 17 | "github.com/pinterest/bender" 18 | ) 19 | 20 | // ThroughputRunner is a test runner for load test throughput commands. 21 | type ThroughputRunner struct { 22 | runner 23 | intervals bender.IntervalGenerator 24 | } 25 | 26 | // NewThroughputRunner returns new ThroughputRunner. 27 | func NewThroughputRunner(params *Params) *ThroughputRunner { 28 | return &ThroughputRunner{ 29 | runner: runner{ 30 | Params: params, 31 | }, 32 | } 33 | } 34 | 35 | // Before prepares requests, recorders and interval generator. 36 | func (r *ThroughputRunner) Before(qps tester.QPS, opts interface{}) error { 37 | if err := r.runner.Before(qps, opts); err != nil { 38 | return err 39 | } 40 | 41 | o, ok := opts.(*options.Options) 42 | if !ok { 43 | return tester.ErrInvalidOptions 44 | } 45 | 46 | count := int(float64(qps) * float64(o.Duration/time.Second)) 47 | r.intervals = o.Distribution(float64(qps)) 48 | 49 | r.requests = make(chan interface{}, o.BufferSize) 50 | 51 | go func() { 52 | for i := 0; i < count; i++ { 53 | r.requests <- r.Params.RequestGenerator(i) 54 | } 55 | close(r.requests) 56 | }() 57 | 58 | r.progress, r.bar = recorders.NewLoadTestProgress(count) 59 | r.progress.Start() 60 | r.recorders = append(r.recorders, recorders.NewProgressBarRecorder(r.bar)) 61 | 62 | return nil 63 | } 64 | 65 | // After cleans up after the test. 66 | func (r *ThroughputRunner) After(test int, options interface{}) { 67 | r.progress.Stop() 68 | r.runner.After(test, options) 69 | } 70 | 71 | // Intervals returns the interval generator. 72 | func (r *ThroughputRunner) Intervals() bender.IntervalGenerator { 73 | return r.intervals 74 | } 75 | -------------------------------------------------------------------------------- /cmd/core/sync.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package core 10 | 11 | import ( 12 | "sync" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var ( 17 | postinitLock = make(chan struct{}, 1) 18 | postinitWaitGroup = &sync.WaitGroup{} 19 | ) 20 | 21 | // DeferPostInit defers an execution of postinit function until a StartPostInit 22 | // is called. 23 | func DeferPostInit(postinit func()) { 24 | postinitWaitGroup.Add(1) 25 | 26 | go func() { 27 | <-postinitLock 28 | postinit() 29 | postinitWaitGroup.Done() 30 | postinitLock <- struct{}{} 31 | }() 32 | } 33 | 34 | // StartPostInit starts all postinit functions and blocks execution until all of 35 | // them are finished. 36 | func StartPostInit() { 37 | postinitLock <- struct{}{} 38 | 39 | postinitWaitGroup.Wait() 40 | } 41 | -------------------------------------------------------------------------------- /cmd/dhcpv4/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv4 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "dhcpv4", 18 | Short: "Test DHCPv4", 19 | Long: ` 20 | Target: ipv4, ipv6, hostname, ipv4:port, [ipv6]:port, hostname:port. 21 | Port defaults to 67, unless you know what you're doing you shouldn't change it. 22 | 23 | Input format: "DeviceMAC" 24 | 01:23:45:67:89:ab 25 | E3:63:BD:7B:D2:2C 26 | c8:6c:2c:47:96:fd`, 27 | Fixed: ` fbender dhcpv4 {test} fixed -t $TARGET 10 20 28 | fbender dhcpv4 {test} fixed -t $TARGET -d 5m 50`, 29 | Constraints: ` fbender dhcpv4 {test} constraints -t $TARGET -c "AVG(latency)<10" 20 30 | fbender dhcpv4 {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 31 | } 32 | 33 | // Command is the DHCPv4 subcommand. 34 | //nolint:gochecknoglobals 35 | var Command = core.NewTestCommand(template, params) 36 | 37 | //nolint:gochecknoinits 38 | func init() { 39 | optionCodes := NewOptionCodeSliceValue() 40 | Command.PersistentFlags().VarP(optionCodes, "oro", "r", "dhcpv4 parameter request list") 41 | Command.Aliases = []string{"dhcp4"} 42 | } 43 | -------------------------------------------------------------------------------- /cmd/dhcpv4/dhcpv4.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv4 10 | 11 | import ( 12 | "net" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/input" 15 | "github.com/facebookincubator/fbender/cmd/core/options" 16 | "github.com/facebookincubator/fbender/cmd/core/runner" 17 | tester "github.com/facebookincubator/fbender/tester/dhcpv4" 18 | "github.com/facebookincubator/fbender/utils" 19 | "github.com/insomniacslk/dhcp/dhcpv4" 20 | "github.com/insomniacslk/dhcp/dhcpv4/async" 21 | "github.com/insomniacslk/dhcp/iana" 22 | "github.com/spf13/cobra" 23 | ) 24 | 25 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 26 | optionCodes, err := GetOptionCodes(cmd.Flags(), "oro") 27 | if err != nil { 28 | return nil, err 29 | } 30 | 31 | r, err := input.NewRequestGenerator(o.Input, inputTransformer(optionCodes)) 32 | if err != nil { 33 | //nolint:wrapcheck 34 | return nil, err 35 | } 36 | 37 | t := &tester.Tester{ 38 | Target: utils.WithDefaultPort(o.Target, async.DefaultServerPort), 39 | Timeout: o.Timeout, 40 | BufferSize: o.BufferSize, 41 | } 42 | 43 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 44 | } 45 | 46 | func inputTransformer(optionCodes []dhcpv4.OptionCode) input.Transformer { 47 | defaultCodes := []dhcpv4.OptionCode{ 48 | dhcpv4.OptionSubnetMask, 49 | dhcpv4.OptionRouter, 50 | dhcpv4.OptionDomainName, 51 | dhcpv4.OptionDomainNameServer, 52 | } 53 | 54 | return func(input string) (interface{}, error) { 55 | mac, err := net.ParseMAC(input) 56 | if err != nil { 57 | //nolint:wrapcheck 58 | return nil, err 59 | } 60 | 61 | discover, err := dhcpv4.New() 62 | if err != nil { 63 | //nolint:wrapcheck 64 | return nil, err 65 | } 66 | 67 | discover.HWType = iana.HWTypeEthernet 68 | discover.ClientHWAddr = mac 69 | discover.UpdateOption(dhcpv4.OptMessageType(dhcpv4.MessageTypeDiscover)) 70 | 71 | optionCodes = append(optionCodes, defaultCodes...) 72 | dhcpv4.WithRequestedOptions(optionCodes...)(discover) 73 | 74 | return discover, nil 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /cmd/dhcpv4/flags.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv4 10 | 11 | import ( 12 | "encoding/csv" 13 | "fmt" 14 | "strconv" 15 | "strings" 16 | 17 | "github.com/facebookincubator/fbender/flags" 18 | "github.com/insomniacslk/dhcp/dhcpv4" 19 | "github.com/spf13/pflag" 20 | ) 21 | 22 | type optionCodeSliceValue struct { 23 | value dhcpv4.OptionCodeList 24 | changed bool 25 | } 26 | 27 | // NewOptionCodeSliceValue creates a new option code slice value for pflag. 28 | func NewOptionCodeSliceValue() pflag.Value { 29 | return &optionCodeSliceValue{ 30 | changed: false, 31 | } 32 | } 33 | 34 | func readAsCSV(val string) ([]string, error) { 35 | if val == "" { 36 | return []string{}, nil 37 | } 38 | 39 | stringReader := strings.NewReader(val) 40 | csvReader := csv.NewReader(stringReader) 41 | 42 | return csvReader.Read() 43 | } 44 | 45 | func (s *optionCodeSliceValue) Set(value string) error { 46 | values, err := readAsCSV(value) 47 | if err != nil { 48 | return err 49 | } 50 | 51 | var buf []byte 52 | 53 | var optcodes dhcpv4.OptionCodeList 54 | 55 | for _, v := range values { 56 | var optcode uint64 57 | 58 | optcode, err = strconv.ParseUint(v, 10, 8) 59 | if err != nil { 60 | //nolint:wrapcheck 61 | return err 62 | } 63 | 64 | buf = append(buf, byte(optcode)) 65 | } 66 | 67 | err = optcodes.FromBytes(buf) 68 | if err != nil { 69 | //nolint:wrapcheck 70 | return err 71 | } 72 | 73 | if !s.changed { 74 | s.value = optcodes 75 | } else { 76 | s.value.Add(optcodes...) 77 | } 78 | 79 | s.changed = true 80 | 81 | return nil 82 | } 83 | 84 | func (s *optionCodeSliceValue) Type() string { 85 | return "optioncodes" 86 | } 87 | 88 | func (s *optionCodeSliceValue) String() string { 89 | return s.value.String() 90 | } 91 | 92 | // GetOptionCodes returns an option code slice from a pflag set. 93 | func GetOptionCodes(f *pflag.FlagSet, name string) (dhcpv4.OptionCodeList, error) { 94 | flag := f.Lookup(name) 95 | if flag == nil { 96 | return nil, fmt.Errorf("%w: %q", flags.ErrUndefined, name) 97 | } 98 | 99 | return GetOptionCodesValue(flag.Value) 100 | } 101 | 102 | // GetOptionCodesValue returns an option code slice from a pflag value. 103 | func GetOptionCodesValue(v pflag.Value) (dhcpv4.OptionCodeList, error) { 104 | if optcodes, ok := v.(*optionCodeSliceValue); ok { 105 | return optcodes.value, nil 106 | } 107 | 108 | return nil, fmt.Errorf("%w, want: optioncodes, got: %s", flags.ErrInvalidType, v.Type()) 109 | } 110 | -------------------------------------------------------------------------------- /cmd/dhcpv6/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv6 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "dhcpv6", 18 | Short: "Test DHCPv6", 19 | Long: ` 20 | Target: ipv4, ipv6, hostname, ipv4:port, [ipv6]:port, hostname:port. 21 | Port defaults to 547, unless you know what you're doing you shouldn't change it. 22 | 23 | Input format: "DeviceMAC" 24 | 01:23:45:67:89:ab 25 | E3:63:BD:7B:D2:2C 26 | c8:6c:2c:47:96:fd`, 27 | Fixed: ` fbender dhcpv6 {test} fixed -t $TARGET 10 20 28 | fbender dhcpv6 {test} fixed -t $TARGET -d 5m 50`, 29 | Constraints: ` fbender dhcpv6 {test} constraints -t $TARGET -c "AVG(latency)<10" 20 30 | fbender dhcpv6 {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 31 | } 32 | 33 | // Command is the TFTP subcommand. 34 | //nolint:gochecknoglobals 35 | var Command = core.NewTestCommand(template, params) 36 | 37 | //nolint:gochecknoinits 38 | func init() { 39 | optionCodes := NewOptionCodeSliceValue() 40 | Command.PersistentFlags().VarP(optionCodes, "oro", "r", "dhcpv6 requested options (ORO)") 41 | Command.Aliases = []string{"dhcp6"} 42 | } 43 | -------------------------------------------------------------------------------- /cmd/dhcpv6/dhcpv6.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv6 10 | 11 | import ( 12 | "net" 13 | 14 | "github.com/facebookincubator/fbender/cmd/core/input" 15 | "github.com/facebookincubator/fbender/cmd/core/options" 16 | "github.com/facebookincubator/fbender/cmd/core/runner" 17 | tester "github.com/facebookincubator/fbender/tester/dhcpv6" 18 | "github.com/facebookincubator/fbender/utils" 19 | "github.com/insomniacslk/dhcp/dhcpv6" 20 | "github.com/spf13/cobra" 21 | ) 22 | 23 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 24 | optionCodes, err := GetOptionCodes(cmd.Flags(), "oro") 25 | if err != nil { 26 | return nil, err 27 | } 28 | 29 | r, err := input.NewRequestGenerator(o.Input, inputTransformer(optionCodes)) 30 | if err != nil { 31 | //nolint:wrapcheck 32 | return nil, err 33 | } 34 | 35 | t := &tester.Tester{ 36 | Target: utils.WithDefaultPort(o.Target, dhcpv6.DefaultServerPort), 37 | Timeout: o.Timeout, 38 | BufferSize: o.BufferSize, 39 | } 40 | 41 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 42 | } 43 | 44 | func inputTransformer(optionCodes []dhcpv6.OptionCode) input.Transformer { 45 | return func(input string) (interface{}, error) { 46 | mac, err := net.ParseMAC(input) 47 | if err != nil { 48 | //nolint:wrapcheck 49 | return nil, err 50 | } 51 | 52 | return dhcpv6.NewSolicit(mac, dhcpv6.WithRequestedOptions(optionCodes...)) 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /cmd/dhcpv6/flags.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv6 10 | 11 | import ( 12 | "encoding/csv" 13 | "fmt" 14 | "strconv" 15 | "strings" 16 | 17 | "github.com/facebookincubator/fbender/flags" 18 | "github.com/insomniacslk/dhcp/dhcpv6" 19 | "github.com/spf13/pflag" 20 | ) 21 | 22 | type optionCodeSliceValue struct { 23 | value *[]dhcpv6.OptionCode 24 | changed bool 25 | } 26 | 27 | // NewOptionCodeSliceValue creates a new option code slice value for pflag. 28 | func NewOptionCodeSliceValue() pflag.Value { 29 | v := []dhcpv6.OptionCode{} 30 | 31 | return &optionCodeSliceValue{ 32 | value: &v, 33 | changed: false, 34 | } 35 | } 36 | 37 | func readAsCSV(val string) ([]string, error) { 38 | if val == "" { 39 | return []string{}, nil 40 | } 41 | 42 | stringReader := strings.NewReader(val) 43 | csvReader := csv.NewReader(stringReader) 44 | 45 | return csvReader.Read() 46 | } 47 | 48 | func (s *optionCodeSliceValue) Set(value string) error { 49 | values, err := readAsCSV(value) 50 | if err != nil { 51 | return err 52 | } 53 | 54 | optcodes := []dhcpv6.OptionCode{} 55 | 56 | for _, v := range values { 57 | optcode, err := strconv.ParseUint(v, 10, 8) 58 | if err != nil { 59 | //nolint:wrapcheck 60 | return err 61 | } 62 | 63 | optcodes = append(optcodes, dhcpv6.OptionCode(optcode)) 64 | } 65 | 66 | if !s.changed { 67 | *s.value = optcodes 68 | } else { 69 | *s.value = append(*s.value, optcodes...) 70 | } 71 | 72 | s.changed = true 73 | 74 | return nil 75 | } 76 | 77 | func (s *optionCodeSliceValue) Type() string { 78 | return "optioncodes" 79 | } 80 | 81 | func (s *optionCodeSliceValue) String() string { 82 | out := make([]string, len(*s.value)) 83 | for i, o := range *s.value { 84 | out[i] = o.String() 85 | } 86 | 87 | return strings.Join(out, ", ") 88 | } 89 | 90 | // GetOptionCodes returns an option code slice from a pflag set. 91 | func GetOptionCodes(f *pflag.FlagSet, name string) ([]dhcpv6.OptionCode, error) { 92 | flag := f.Lookup(name) 93 | if flag == nil { 94 | return nil, fmt.Errorf("%w: %q", flags.ErrUndefined, name) 95 | } 96 | 97 | return GetOptionCodesValue(flag.Value) 98 | } 99 | 100 | // GetOptionCodesValue returns an option code slice from a pflag value. 101 | func GetOptionCodesValue(v pflag.Value) ([]dhcpv6.OptionCode, error) { 102 | if optcodes, ok := v.(*optionCodeSliceValue); ok { 103 | return *optcodes.value, nil 104 | } 105 | 106 | return nil, fmt.Errorf("%w, want: optioncodes, got: %s", flags.ErrInvalidType, v.Type()) 107 | } 108 | -------------------------------------------------------------------------------- /cmd/dhcpv6/flags_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv6_test 10 | 11 | import ( 12 | "strconv" 13 | "testing" 14 | 15 | dhcpv6flags "github.com/facebookincubator/fbender/cmd/dhcpv6" 16 | "github.com/facebookincubator/fbender/flags" 17 | "github.com/insomniacslk/dhcp/dhcpv6" 18 | "github.com/spf13/pflag" 19 | "github.com/stretchr/testify/suite" 20 | ) 21 | 22 | type OptionCodeSliceValueTestSuite struct { 23 | suite.Suite 24 | value pflag.Value 25 | } 26 | 27 | func (s *OptionCodeSliceValueTestSuite) SetupTest() { 28 | s.value = dhcpv6flags.NewOptionCodeSliceValue() 29 | s.Require().NotNil(s.value) 30 | } 31 | 32 | func (s *OptionCodeSliceValueTestSuite) TestSet_NoErrors() { 33 | err := s.value.Set("1,2") 34 | s.Require().NoError(err) 35 | 36 | v, err := dhcpv6flags.GetOptionCodesValue(s.value) 37 | s.Require().NoError(err) 38 | 39 | o := []dhcpv6.OptionCode{dhcpv6.OptionClientID, dhcpv6.OptionServerID} 40 | s.Assert().Equal(o, v) 41 | 42 | // Check if consecutive calls append values 43 | err = s.value.Set("3") 44 | s.Require().NoError(err) 45 | 46 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 47 | s.Require().NoError(err) 48 | 49 | o = append(o, dhcpv6.OptionIANA) 50 | s.Assert().Equal(o, v) 51 | 52 | err = s.value.Set("4,5") 53 | s.Require().NoError(err) 54 | 55 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 56 | s.Require().NoError(err) 57 | 58 | o = append(o, dhcpv6.OptionIATA, dhcpv6.OptionIAAddr) 59 | s.Assert().Equal(o, v) 60 | } 61 | 62 | func (s *OptionCodeSliceValueTestSuite) TestSet_Errors() { 63 | // Errors - single value 64 | err := s.value.Set("notanumber") 65 | s.Assert().ErrorIs(err, strconv.ErrSyntax) 66 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"notanumber\": invalid syntax") 67 | 68 | v, err := dhcpv6flags.GetOptionCodesValue(s.value) 69 | s.Require().NoError(err) 70 | s.Assert().Empty(v) 71 | 72 | err = s.value.Set("42.5") 73 | s.Assert().ErrorIs(err, strconv.ErrSyntax) 74 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"42.5\": invalid syntax") 75 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 76 | s.Require().NoError(err) 77 | s.Assert().Empty(v) 78 | 79 | err = s.value.Set("-10") 80 | s.Assert().ErrorIs(err, strconv.ErrSyntax) 81 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"-10\": invalid syntax") 82 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 83 | s.Require().NoError(err) 84 | s.Assert().Empty(v) 85 | 86 | err = s.value.Set("256") 87 | s.Assert().ErrorIs(err, strconv.ErrRange) 88 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"256\": value out of range") 89 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 90 | s.Require().NoError(err) 91 | s.Assert().Empty(v) 92 | 93 | // Errors - multiple values 94 | err = s.value.Set("42,notanumber") 95 | s.Assert().ErrorIs(err, strconv.ErrSyntax) 96 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"notanumber\": invalid syntax") 97 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 98 | s.Require().NoError(err) 99 | s.Assert().Empty(v) 100 | 101 | err = s.value.Set("42,42.5,notanumber") 102 | s.Assert().ErrorIs(err, strconv.ErrSyntax) 103 | s.Assert().EqualError(err, "strconv.ParseUint: parsing \"42.5\": invalid syntax") 104 | v, err = dhcpv6flags.GetOptionCodesValue(s.value) 105 | s.Require().NoError(err) 106 | s.Assert().Empty(v) 107 | } 108 | 109 | func (s *OptionCodeSliceValueTestSuite) TestType() { 110 | s.Assert().Equal("optioncodes", s.value.Type()) 111 | } 112 | 113 | func (s *OptionCodeSliceValueTestSuite) TestString_Known() { 114 | // No options 115 | s.Assert().Equal("", s.value.String()) 116 | 117 | // Single option 118 | err := s.value.Set("1") 119 | s.Require().NoError(err) 120 | v := s.value.String() 121 | s.Assert().Equal("Client Identifier", v) 122 | 123 | // Multiple options 124 | err = s.value.Set("2,3") 125 | s.Require().NoError(err) 126 | v = s.value.String() 127 | s.Assert().Equal("Client Identifier, Server Identifier, IA_NA", v) 128 | } 129 | 130 | func (s *OptionCodeSliceValueTestSuite) TestString_Unknown() { 131 | // No options 132 | s.Assert().Equal("", s.value.String()) 133 | 134 | // Single option 135 | err := s.value.Set("10") 136 | s.Require().NoError(err) 137 | v := s.value.String() 138 | s.Assert().Equal("unknown (10)", v) 139 | 140 | // Multiple options 141 | err = s.value.Set("35,255") 142 | s.Require().NoError(err) 143 | v = s.value.String() 144 | s.Assert().Equal("unknown (10), unknown (35), unknown (255)", v) 145 | } 146 | 147 | func (s *OptionCodeSliceValueTestSuite) TestGetOptionCodes() { 148 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 149 | f.Var(s.value, "optioncodes", "set option codes") 150 | 151 | err := s.value.Set("39") 152 | s.Require().NoError(err) 153 | 154 | v, err := dhcpv6flags.GetOptionCodes(f, "optioncodes") 155 | s.Require().NoError(err) 156 | s.Assert().Equal([]dhcpv6.OptionCode{dhcpv6.OptionFQDN}, v) 157 | 158 | // Check error when flag does not exist 159 | _, err = dhcpv6flags.GetOptionCodes(f, "nonexistent") 160 | s.Assert().ErrorIs(err, flags.ErrUndefined) 161 | s.Assert().EqualError(err, "flag accessed but not defined: \"nonexistent\"") 162 | 163 | // Check error when value is of different type 164 | f.Int("myint", 0, "set myint") 165 | _, err = dhcpv6flags.GetOptionCodes(f, "myint") 166 | s.Assert().ErrorIs(err, flags.ErrInvalidType) 167 | s.Assert().EqualError(err, "accessed flag type does not match, want: optioncodes, got: int") 168 | } 169 | 170 | func (s *OptionCodeSliceValueTestSuite) TestGetOptionCodesValue() { 171 | err := s.value.Set("39") 172 | s.Require().NoError(err) 173 | 174 | v, err := dhcpv6flags.GetOptionCodesValue(s.value) 175 | s.Require().NoError(err) 176 | s.Assert().Equal([]dhcpv6.OptionCode{dhcpv6.OptionFQDN}, v) 177 | 178 | // Check error when value is of different type 179 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 180 | f.Int("myint", 0, "set myint") 181 | flag := f.Lookup("myint") 182 | s.Require().NotNil(flag) 183 | _, err = dhcpv6flags.GetOptionCodesValue(flag.Value) 184 | s.Assert().ErrorIs(err, flags.ErrInvalidType) 185 | s.Assert().EqualError(err, "accessed flag type does not match, want: optioncodes, got: int") 186 | } 187 | 188 | func TestOptionCodeSliceValueTestSuite(t *testing.T) { 189 | suite.Run(t, new(OptionCodeSliceValueTestSuite)) 190 | } 191 | -------------------------------------------------------------------------------- /cmd/dns/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dns 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "dns", 18 | Short: "Test DNS", 19 | Long: ` 20 | Queries may be prefixed with timestamp and a random 16-character hex to avoid 21 | hitting the cache. In bash this could have been achieved by running: 22 | $(date +%s).$(openssl rand -hex 16).domain 23 | 24 | Target: ipv4, ipv6, hostname, ipv4:port, [ipv6]:port, hostname:port. 25 | The port defaults to 53. 26 | 27 | Input format: "Domain QType [Rcode]" 28 | example.com AAAA 29 | other.example.com TXT NOERROR 30 | mail.example.com MX 31 | www.doesnotexist.co.uk NXDOMAIN`, 32 | Fixed: ` fbender dns {test} fixed -t $TARGET 10 20 33 | fbender dns {test} fixed -t $TARGET -r -d 5m 50`, 34 | Constraints: ` fbender dns {test} constraints -t $TARGET -r -c "AVG(latency)<10" 20 35 | fbender dns {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 36 | } 37 | 38 | // Command is the DNS subcommand. 39 | //nolint:gochecknoglobals 40 | var Command = core.NewTestCommand(template, params) 41 | 42 | //nolint:gochecknoinits 43 | func init() { 44 | Command.PersistentFlags().BoolP("randomize", "r", false, "randomize queries with timestamp and a random hex") 45 | core.DeferPostInit(postinit) 46 | } 47 | 48 | func postinit() { 49 | protocol := NewProtocolValue() 50 | 51 | Command.PersistentFlags().VarP(protocol, "protocol", "p", "protocol used for DNS queries (udp|tcp)") 52 | 53 | if err := BashCompletionProtocol(Command, Command.PersistentFlags(), "protocol"); err != nil { 54 | panic(err) 55 | } 56 | } 57 | -------------------------------------------------------------------------------- /cmd/dns/dns.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dns 10 | 11 | import ( 12 | "fmt" 13 | "strings" 14 | "time" 15 | 16 | "github.com/facebookincubator/fbender/cmd/core/errors" 17 | "github.com/facebookincubator/fbender/cmd/core/input" 18 | "github.com/facebookincubator/fbender/cmd/core/options" 19 | "github.com/facebookincubator/fbender/cmd/core/runner" 20 | tester "github.com/facebookincubator/fbender/tester/dns" 21 | "github.com/facebookincubator/fbender/utils" 22 | "github.com/miekg/dns" 23 | "github.com/spf13/cobra" 24 | ) 25 | 26 | // DefaultServerPort is a default dns server port. 27 | const DefaultServerPort = 53 28 | 29 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 30 | randomize, err := cmd.Flags().GetBool("randomize") 31 | if err != nil { 32 | //nolint:wrapcheck 33 | return nil, err 34 | } 35 | 36 | protocol, err := GetProtocol(cmd.Flags(), "protocol") 37 | if err != nil { 38 | return nil, err 39 | } 40 | 41 | r, err := input.NewRequestGenerator(o.Input, inputTransformer, getModifiers(randomize)...) 42 | if err != nil { 43 | //nolint:wrapcheck 44 | return nil, err 45 | } 46 | 47 | t := &tester.Tester{ 48 | Target: utils.WithDefaultPort(o.Target, DefaultServerPort), 49 | Timeout: o.Timeout, 50 | Protocol: protocol, 51 | } 52 | 53 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 54 | } 55 | 56 | func inputTransformer(input string) (interface{}, error) { 57 | var domain, typeString, rcodeString string 58 | 59 | n, err := fmt.Sscanf(input, "%s %s %s", &domain, &typeString, &rcodeString) 60 | if err != nil && n < 2 { 61 | return nil, fmt.Errorf("%w, want: \"Domain QType [RCode]\", got: %q", errors.ErrInvalidFormat, input) 62 | } 63 | 64 | msgTyp, ok := dns.StringToType[strings.ToUpper(typeString)] 65 | if !ok { 66 | return nil, fmt.Errorf("%w, invalid QType: %q", errors.ErrInvalidFormat, typeString) 67 | } 68 | 69 | msg := new(tester.ExtendedMsg) 70 | 71 | msg.SetQuestion(dns.Fqdn(domain), msgTyp) 72 | msg.Rcode = -1 73 | 74 | if n == 3 { 75 | rcode, ok := dns.StringToRcode[rcodeString] 76 | if !ok { 77 | return nil, fmt.Errorf("%w, invalid RCode: %q", errors.ErrInvalidFormat, rcodeString) 78 | } 79 | 80 | msg.Rcode = rcode 81 | } 82 | 83 | return msg, nil 84 | } 85 | 86 | func getModifiers(randomize bool) []input.Modifier { 87 | if randomize { 88 | return []input.Modifier{randomPrefixModifier} 89 | } 90 | 91 | return []input.Modifier{} 92 | } 93 | 94 | const prefixLength = 16 95 | 96 | func randomPrefixModifier(request interface{}) (interface{}, error) { 97 | msg, ok := request.(*tester.ExtendedMsg) 98 | if !ok { 99 | return nil, fmt.Errorf("%w, want: *dns.ExtendedMsg, got: %T", errors.ErrInvalidType, request) 100 | } 101 | 102 | hex, err := utils.RandomHex(prefixLength) 103 | if err != nil { 104 | //nolint:wrapcheck 105 | return nil, err 106 | } 107 | 108 | // Create a new message so we don't destroy the original to avoid recursive prefixing 109 | modified := new(tester.ExtendedMsg) 110 | domain := fmt.Sprintf("%d.%s.%s", time.Now().Unix(), hex, msg.Question[0].Name) 111 | msgTyp := msg.Question[0].Qtype 112 | 113 | modified.SetQuestion(dns.Fqdn(domain), msgTyp) 114 | modified.Rcode = msg.Rcode 115 | 116 | return modified, nil 117 | } 118 | -------------------------------------------------------------------------------- /cmd/dns/flags.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dns 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | 15 | "github.com/facebookincubator/fbender/flags" 16 | "github.com/facebookincubator/fbender/utils" 17 | "github.com/spf13/cobra" 18 | "github.com/spf13/pflag" 19 | ) 20 | 21 | // protocols is a set of available protocols. 22 | //nolint:gochecknoglobals 23 | var protocols = map[string]struct{}{ 24 | "udp": {}, 25 | "tcp": {}, 26 | } 27 | 28 | // ErrInvalidProtocol is raised when an unknown protocol is set. 29 | var ErrInvalidProtocol = errors.New("invalid protocol") 30 | 31 | type protocolValue struct { 32 | value string 33 | } 34 | 35 | // NewProtocolValue returns new Protocol flag with a default value. 36 | func NewProtocolValue() pflag.Value { 37 | return &protocolValue{value: "udp"} 38 | } 39 | 40 | func (s *protocolValue) Set(value string) error { 41 | if _, ok := protocols[value]; ok { 42 | s.value = value 43 | 44 | return nil 45 | } 46 | 47 | return fmt.Errorf("%w, want: \"udp\" or \"tcp\", got: %q", ErrInvalidProtocol, value) 48 | } 49 | 50 | func (s *protocolValue) Type() string { 51 | return "protocol" 52 | } 53 | 54 | func (s *protocolValue) String() string { 55 | return s.value 56 | } 57 | 58 | // GetProtocol returns a protocol from a pflag set. 59 | func GetProtocol(f *pflag.FlagSet, name string) (string, error) { 60 | flag := f.Lookup(name) 61 | if flag == nil { 62 | return "", fmt.Errorf("%w: %q", flags.ErrUndefined, name) 63 | } 64 | 65 | return GetProtocolValue(flag.Value) 66 | } 67 | 68 | // GetProtocolValue returns a protocol from a pflag value. 69 | func GetProtocolValue(v pflag.Value) (string, error) { 70 | if protocol, ok := v.(*protocolValue); ok { 71 | return protocol.value, nil 72 | } 73 | 74 | return "", fmt.Errorf("%w, want: protocol, got: %s", flags.ErrInvalidType, v.Type()) 75 | } 76 | 77 | // Bash completion function constants. 78 | const ( 79 | fname = "__fbender_handle_dns_protocol_flag" 80 | fbody = `COMPREPLY=($(compgen -W "udp tcp" -- "${cur}"))` 81 | ) 82 | 83 | // BashCompletionProtocol adds bash completion to a protocol flag. 84 | func BashCompletionProtocol(cmd *cobra.Command, flags *pflag.FlagSet, name string) error { 85 | return utils.BashCompletion(cmd, flags, name, fname, fbody) 86 | } 87 | -------------------------------------------------------------------------------- /cmd/dns/flags_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dns_test 10 | 11 | import ( 12 | "strings" 13 | "testing" 14 | 15 | dnsflags "github.com/facebookincubator/fbender/cmd/dns" 16 | "github.com/facebookincubator/fbender/flags" 17 | "github.com/spf13/cobra" 18 | "github.com/spf13/pflag" 19 | "github.com/stretchr/testify/suite" 20 | ) 21 | 22 | type ProtocolValueTestSuite struct { 23 | suite.Suite 24 | value pflag.Value 25 | } 26 | 27 | func (s *ProtocolValueTestSuite) SetupTest() { 28 | s.value = dnsflags.NewProtocolValue() 29 | s.Require().NotNil(s.value) 30 | } 31 | 32 | func (s *ProtocolValueTestSuite) TestSet_NoErrors() { 33 | err := s.value.Set("udp") 34 | s.Require().NoError(err) 35 | 36 | v, err := dnsflags.GetProtocolValue(s.value) 37 | s.Require().NoError(err) 38 | s.Assert().Equal("udp", v) 39 | 40 | err = s.value.Set("tcp") 41 | s.Require().NoError(err) 42 | 43 | v, err = dnsflags.GetProtocolValue(s.value) 44 | s.Require().NoError(err) 45 | s.Assert().Equal("tcp", v) 46 | } 47 | 48 | func (s *ProtocolValueTestSuite) TestSet_Errors() { 49 | // Save original flag value 50 | o, err := dnsflags.GetProtocolValue(s.value) 51 | s.Require().NoError(err) 52 | 53 | // Try invalid value 54 | err = s.value.Set("unknown") 55 | s.Assert().ErrorIs(err, dnsflags.ErrInvalidProtocol) 56 | s.Assert().EqualError(err, "invalid protocol, want: \"udp\" or \"tcp\", got: \"unknown\"") 57 | 58 | // The value shouldn't change 59 | v, err := dnsflags.GetProtocolValue(s.value) 60 | s.Require().NoError(err) 61 | s.Assert().Equal(o, v) 62 | } 63 | 64 | func (s *ProtocolValueTestSuite) TestType() { 65 | s.Assert().Equal("protocol", s.value.Type()) 66 | } 67 | 68 | func (s *ProtocolValueTestSuite) TestGetProtocol() { 69 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 70 | f.Var(s.value, "protocol", "set protocol") 71 | 72 | err := s.value.Set("tcp") 73 | s.Require().NoError(err) 74 | 75 | v, err := dnsflags.GetProtocol(f, "protocol") 76 | s.Require().NoError(err) 77 | s.Assert().Equal("tcp", v) 78 | 79 | // Check error when flag does not exist 80 | _, err = dnsflags.GetProtocol(f, "nonexistent") 81 | s.Assert().ErrorIs(err, flags.ErrUndefined) 82 | s.Assert().EqualError(err, "flag accessed but not defined: \"nonexistent\"") 83 | 84 | // Check error when value is of different type 85 | f.Int("myint", 0, "set myint") 86 | _, err = dnsflags.GetProtocol(f, "myint") 87 | s.Assert().ErrorIs(err, flags.ErrInvalidType) 88 | s.Assert().EqualError(err, "accessed flag type does not match, want: protocol, got: int") 89 | } 90 | 91 | func (s *ProtocolValueTestSuite) TestGetProtocolValue() { 92 | err := s.value.Set("tcp") 93 | s.Require().NoError(err) 94 | 95 | v, err := dnsflags.GetProtocolValue(s.value) 96 | s.Require().NoError(err) 97 | s.Assert().Equal("tcp", v) 98 | 99 | // Check error when value is of different type 100 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 101 | f.Int("myint", 0, "set myint") 102 | 103 | flag := f.Lookup("myint") 104 | s.Require().NotNil(flag) 105 | 106 | _, err = dnsflags.GetProtocolValue(flag.Value) 107 | s.Assert().ErrorIs(err, flags.ErrInvalidType) 108 | s.Assert().EqualError(err, "accessed flag type does not match, want: protocol, got: int") 109 | } 110 | 111 | func (s *ProtocolValueTestSuite) TestBashCompletionProtocol() { 112 | c := &cobra.Command{} 113 | f := c.Flags().VarPF(s.value, "protocol", "p", "set protocol") 114 | 115 | // Check if the complete function is appended 116 | err := dnsflags.BashCompletionProtocol(c, c.Flags(), "protocol") 117 | s.Require().NoError(err) 118 | s.Assert().Contains(c.BashCompletionFunction, "__fbender_handle_dns_protocol_flag") 119 | 120 | // Check if the flag has the bash 121 | s.Require().Contains(f.Annotations, "cobra_annotation_bash_completion_custom") 122 | s.Assert().Equal([]string{"__fbender_handle_dns_protocol_flag"}, 123 | f.Annotations["cobra_annotation_bash_completion_custom"]) 124 | 125 | // Check if the function is appended only once 126 | err = dnsflags.BashCompletionProtocol(c, c.Flags(), "protocol") 127 | s.Require().NoError(err) 128 | 129 | count := strings.Count(c.BashCompletionFunction, "__fbender_handle_dns_protocol_flag") 130 | s.Assert().Equal(1, count, "Completion function should be added only once") 131 | } 132 | 133 | func TestProtocolValueTestSuite(t *testing.T) { 134 | suite.Run(t, new(ProtocolValueTestSuite)) 135 | } 136 | -------------------------------------------------------------------------------- /cmd/http/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package http 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "http", 18 | Short: "Test HTTP", 19 | Long: ` 20 | Target: ipv4, ipv4:port, ipv6, [ipv6]:port, hostname, hostname:port. 21 | 22 | Input format: "GET RelativeURL" or "POST RelativeURL FormData" 23 | GET index.html 24 | GET / 25 | POST echo message=Hello 26 | POST echo/ message=Hello&name=Mikolaj`, 27 | Fixed: ` fbender http {test} fixed -t $TARGET 10 20 28 | fbender http {test} fixed -t $TARGET -s -d 5m 50`, 29 | Constraints: ` fbender http {test} constraints -t $TARGET -s -c "AVG(latency)<10" 20 30 | fbender http {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 31 | } 32 | 33 | // Command is the HTTP subcommand. 34 | //nolint:gochecknoglobals 35 | var Command = core.NewTestCommand(template, params) 36 | 37 | //nolint:gochecknoinits 38 | func init() { 39 | Command.PersistentFlags().BoolP("ssl", "s", false, "enable ssl (use HTTPS)") 40 | } 41 | -------------------------------------------------------------------------------- /cmd/http/http.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package http 10 | 11 | import ( 12 | "fmt" 13 | "net/http" 14 | "net/url" 15 | "strconv" 16 | "strings" 17 | 18 | "github.com/facebookincubator/fbender/cmd/core/errors" 19 | "github.com/facebookincubator/fbender/cmd/core/input" 20 | "github.com/facebookincubator/fbender/cmd/core/options" 21 | "github.com/facebookincubator/fbender/cmd/core/runner" 22 | tester "github.com/facebookincubator/fbender/tester/http" 23 | "github.com/spf13/cobra" 24 | ) 25 | 26 | const formats = "'GET RelativeURL' or 'POST RelativeURL FormData'" 27 | 28 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 29 | ssl, err := cmd.Flags().GetBool("ssl") 30 | if err != nil { 31 | //nolint:wrapcheck 32 | return nil, err 33 | } 34 | 35 | r, err := input.NewRequestGenerator(o.Input, inputTransformer(ssl, o.Target), requestCreator) 36 | if err != nil { 37 | //nolint:wrapcheck 38 | return nil, err 39 | } 40 | 41 | t := &tester.Tester{ 42 | Timeout: o.Timeout, 43 | } 44 | 45 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 46 | } 47 | 48 | func inputTransformer(ssl bool, target string) input.Transformer { 49 | protocol := "http" 50 | if ssl { 51 | protocol = "https" 52 | } 53 | 54 | return func(input string) (interface{}, error) { 55 | i := strings.Index(input, " ") 56 | if i < 0 { 57 | return nil, fmt.Errorf("%w, want: %s, got: %q", errors.ErrInvalidFormat, formats, input) 58 | } 59 | 60 | method, data := input[:i], input[i+1:] 61 | switch method { 62 | case "GET": 63 | return parseGetRequest(protocol, target, data) 64 | case "POST": 65 | return parsePostRequest(protocol, target, data) 66 | } 67 | 68 | return nil, fmt.Errorf("%w, want: (GET|POST), got: %q", errors.ErrInvalidFormat, method) 69 | } 70 | } 71 | 72 | type request interface { 73 | Create() (*http.Request, error) 74 | } 75 | 76 | type getRequest struct { 77 | url string 78 | } 79 | 80 | func (r *getRequest) Create() (*http.Request, error) { 81 | //nolint:noctx 82 | return http.NewRequest("GET", r.url, nil) 83 | } 84 | 85 | func parseGetRequest(protocol, target, data string) (interface{}, error) { 86 | rawurl, err := joinURL(protocol, target, data) 87 | if err != nil { 88 | return nil, err 89 | } 90 | 91 | return &getRequest{url: rawurl}, nil 92 | } 93 | 94 | type postRequest struct { 95 | url string 96 | body string 97 | } 98 | 99 | func (r *postRequest) Create() (*http.Request, error) { 100 | //nolint:noctx 101 | req, err := http.NewRequest("POST", r.url, strings.NewReader(r.body)) 102 | if err != nil { 103 | //nolint:wrapcheck 104 | return nil, err 105 | } 106 | 107 | req.Header.Add("Content-Type", "application/x-www-form-urlencoded") 108 | req.Header.Add("Content-Length", strconv.Itoa(len(r.body))) 109 | 110 | return req, nil 111 | } 112 | 113 | func parsePostRequest(protocol, target, data string) (interface{}, error) { 114 | i := strings.Index(data, " ") 115 | if i < 0 { 116 | return nil, fmt.Errorf("%w, want: %s, got: \"POST %s\"", errors.ErrInvalidFormat, formats, data) 117 | } 118 | 119 | form, err := url.ParseQuery(data[i+1:]) 120 | if err != nil { 121 | //nolint:wrapcheck 122 | return nil, err 123 | } 124 | 125 | rawurl, err := joinURL(protocol, target, data[:i]) 126 | if err != nil { 127 | return nil, err 128 | } 129 | 130 | return &postRequest{url: rawurl, body: form.Encode()}, nil 131 | } 132 | 133 | // NewRequest uses reader interface for message body, which is being used up. 134 | // Therefore we cannot reuse a once created request and need to invoke 135 | // NewRequest everytime before sending it. 136 | func requestCreator(r interface{}) (interface{}, error) { 137 | if r, ok := r.(request); ok { 138 | return r.Create() 139 | } 140 | 141 | return nil, fmt.Errorf("%w, want: request, got: %T", errors.ErrInvalidType, r) 142 | } 143 | 144 | func joinURL(protocol, target, path string) (string, error) { 145 | path = strings.TrimPrefix(path, "/") 146 | rawurl := fmt.Sprintf("%s://%s/%s", protocol, target, path) 147 | _, err := url.Parse(rawurl) 148 | 149 | //nolint:wrapcheck 150 | return rawurl, err 151 | } 152 | -------------------------------------------------------------------------------- /cmd/tftp/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tftp 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "tftp", 18 | Short: "Test TFTP", 19 | Long: ` 20 | The specified timeout applies to a single datagram in a tftp transfer rather 21 | than to the whole session. 22 | 23 | Target: ipv4:port, [ipv6]:port, hostname:port. 24 | 25 | Input format: "Filename octet" or "Filename netascii" 26 | /my/file octet 27 | /my/otherfile octet 28 | /another netascii`, 29 | Fixed: ` fbender tftp {test} fixed -t $TARGET 10 20 30 | fbender tftp {test} fixed -t $TARGET -d 5m 50`, 31 | Constraints: ` fbender tftp {test} constraints -t $TARGET -b 1500 -c "AVG(latency)<10" 20 32 | fbender tftp {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 33 | } 34 | 35 | // Command is the TFTP subcommand. 36 | //nolint:gochecknoglobals 37 | var Command = core.NewTestCommand(template, params) 38 | 39 | //nolint:gochecknoinits 40 | func init() { 41 | Command.PersistentFlags().IntP("blocksize", "s", 512, "blocksize option as in RFC2348") 42 | } 43 | -------------------------------------------------------------------------------- /cmd/tftp/tftp.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tftp 10 | 11 | import ( 12 | "fmt" 13 | "strings" 14 | 15 | "github.com/facebookincubator/fbender/cmd/core/errors" 16 | "github.com/facebookincubator/fbender/cmd/core/input" 17 | "github.com/facebookincubator/fbender/cmd/core/options" 18 | "github.com/facebookincubator/fbender/cmd/core/runner" 19 | tester "github.com/facebookincubator/fbender/tester/tftp" 20 | "github.com/facebookincubator/fbender/utils" 21 | "github.com/pinterest/bender/tftp" 22 | "github.com/spf13/cobra" 23 | ) 24 | 25 | // DefaultServerPort is a default tftp server port. 26 | const DefaultServerPort = 69 27 | 28 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 29 | blocksize, err := cmd.Flags().GetInt("blocksize") 30 | if err != nil { 31 | //nolint:wrapcheck 32 | return nil, err 33 | } 34 | 35 | r, err := input.NewRequestGenerator(o.Input, inputTransformer) 36 | if err != nil { 37 | //nolint:wrapcheck 38 | return nil, err 39 | } 40 | 41 | t := &tester.Tester{ 42 | Target: utils.WithDefaultPort(o.Target, DefaultServerPort), 43 | Timeout: o.Timeout, 44 | BlockSize: blocksize, 45 | } 46 | 47 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 48 | } 49 | 50 | func inputTransformer(input string) (interface{}, error) { 51 | i := strings.Index(input, " ") 52 | if i < 0 { 53 | return nil, fmt.Errorf("%w, want: \"File Mode\" got %q", errors.ErrInvalidFormat, input) 54 | } 55 | 56 | filename, mode := input[:i], input[i+1:] 57 | if mode != "octet" && mode != "netascii" { 58 | return nil, fmt.Errorf("%w, want: (octet|netascii), got: %q", errors.ErrInvalidFormat, mode) 59 | } 60 | 61 | return &tftp.Request{ 62 | Filename: filename, 63 | Mode: tftp.RequestMode(mode), 64 | }, nil 65 | } 66 | -------------------------------------------------------------------------------- /cmd/udp/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package udp 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd/core" 13 | ) 14 | 15 | //nolint:gochecknoglobals 16 | var template = &core.CommandTemplate{ 17 | Name: "udp", 18 | Short: "Test UDP", 19 | Long: ` 20 | Target: ipv4, ipv6, hostname. 21 | 22 | Input format: "DstPort Base64EncodedeData" 23 | 2545 TG9yZW0= 24 | 7346 aXBzdW0gZG9sb3Igc2l0 25 | 5012 YW1ldCBpbg==`, 26 | Fixed: ` fbender udp {test} fixed -t $TARGET 10 20 27 | fbender udp {test} fixed -t $TARGET -d 5m 50`, 28 | Constraints: ` fbender udp {test} constraints -t $TARGET -c "AVG(latency)<10" 20 29 | fbender udp {test} constraints -t $TARGET -g ^10 -c "MAX(errors)<10" 40`, 30 | } 31 | 32 | // Command is the UDP subcommand. 33 | //nolint:gochecknoglobals 34 | var Command = core.NewTestCommand(template, params) 35 | -------------------------------------------------------------------------------- /cmd/udp/udp.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package udp 10 | 11 | import ( 12 | "encoding/base64" 13 | "fmt" 14 | 15 | "github.com/facebookincubator/fbender/cmd/core/errors" 16 | "github.com/facebookincubator/fbender/cmd/core/input" 17 | "github.com/facebookincubator/fbender/cmd/core/options" 18 | "github.com/facebookincubator/fbender/cmd/core/runner" 19 | "github.com/facebookincubator/fbender/protocols/udp" 20 | tester "github.com/facebookincubator/fbender/tester/udp" 21 | "github.com/spf13/cobra" 22 | ) 23 | 24 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 25 | r, err := input.NewRequestGenerator(o.Input, inputTransformer) 26 | if err != nil { 27 | //nolint:wrapcheck 28 | return nil, err 29 | } 30 | 31 | t := &tester.Tester{ 32 | Target: o.Target, 33 | Timeout: o.Timeout, 34 | } 35 | 36 | return &runner.Params{Tester: t, RequestGenerator: r}, nil 37 | } 38 | 39 | func inputTransformer(input string) (interface{}, error) { 40 | var encodedData string 41 | 42 | datagram := new(udp.Datagram) 43 | 44 | n, err := fmt.Sscanf(input, "%d %s", &datagram.Port, &encodedData) 45 | if err != nil || n < 2 { 46 | return nil, fmt.Errorf("%w, want: \"Port Base64Payload\", got: %q", errors.ErrInvalidFormat, input) 47 | } 48 | 49 | datagram.Data, err = base64.StdEncoding.DecodeString(encodedData) 50 | if err != nil { 51 | //nolint:wrapcheck 52 | return nil, err 53 | } 54 | 55 | return datagram, nil 56 | } 57 | -------------------------------------------------------------------------------- /docs/EXTENDING.md: -------------------------------------------------------------------------------- 1 | # Extending FBender 2 | 3 | ## Adding new protocol 4 | 5 | ### Implement protocol executor 6 | 7 | If the protocol executor implementation can be used other than in FBender try 8 | contributing to [Bender](https://github.com/pinterest/bender) and use the 9 | upstreamed version in the tester. If you think that it's not beneficial to send 10 | a Pull Request with it to Bender create a subdirectory in `protocols` directory. 11 | 12 | ### Create tester 13 | 14 | Create a subdirectory in `tester` with the name of your protocol and a single 15 | `tester.go` file. The tester file should contain a definition of `struct Tester` 16 | and optional definition of `interface Options`. Tester needs to implement: 17 | 18 | ```go 19 | // Tester is used to setup the test for a specific endpoint. 20 | type Tester interface { 21 | // Before is called once, before any tests. 22 | Before(options interface{}) error 23 | // After is called once, after all tests (or after some of them if a test fails). 24 | // This should be used to cleanup everything that was set up in the Before. 25 | After(options interface{}) 26 | // BeforeEach is called before every test. 27 | BeforeEach(options interface{}) error 28 | // AfterEach is called after every test, even if the test fails. This should 29 | // be used to cleanup everything that was set up in the BeforeEach. 30 | AfterEach(options interface{}) 31 | // RequestExecutor is called every time a test is to be ran to get an executor. 32 | RequestExecutor(options interface{}) (bender.RequestExecutor, error) 33 | } 34 | ``` 35 | 36 | For compatibility with different testers options are passed as an `interface{}`, 37 | when asserting options type your tester should return `tester.ErrInvalidOptions` 38 | in case of failure. 39 | 40 | ### Create a command 41 | 42 | Create a subdirectory in `cmd` with the name of your protocol and a files named 43 | `cmd.go` and `${PROTOCOL}.go`. File `cmd.go` should contain a definition of a 44 | command and an optional `init` function adding additional flags or subcommands. 45 | You can copy the following template and replace `${VARIABLES}` with appropriate 46 | values. 47 | 48 | ```go 49 | var template = &core.CommandTemplate{ 50 | Name: "${PROTOCOL}", 51 | Short: "Test ${PROTOCOL}", 52 | Long: ` 53 | Input format: "${INPUT_FORMAT}" 54 | ${INPUT_EXAMPLE_1} 55 | ${INPUT_EXAMPLE_2}`, 56 | Fixed: ` fbender ${PROTOCOL} {test} fixed -t $TARGET 10 20 57 | ${ANOTHER_FIXED_TEST_EXAMPLE}`, 58 | Constraints: ` fbender ${PROTOCOL} {test} constraints -t $TARGET -c "AVG(latency)<10" 20 59 | ${ANOTHER_CONSTRAINTS_TEST_EXAMPLE}`, 60 | } 61 | 62 | var Command = core.NewTestCommand(template, params) 63 | ``` 64 | 65 | The `${PROTOCOL}.go` should implement all protocol specific features. The 66 | simplest one could look like this: 67 | 68 | ```go 69 | func params(cmd *cobra.Command, o *options.Options) (*runner.Params, error) { 70 | // create input based request generator 71 | requests, err := input.NewRequestGenerator(o.Input, inputTransformer) 72 | if err != nil { 73 | return nil, err 74 | } 75 | // create tester for your protocol 76 | tester := &protocol.Tester{ 77 | Target: o.Target, 78 | } 79 | return &runner.Params{Tester: tester, RequestGenerator: requests}, nil 80 | } 81 | 82 | // inputTransformer accepts any string as an input and returns it as a request 83 | func inputTransformer(input string) (interface{}, error) { 84 | return input, nil 85 | } 86 | ``` 87 | 88 | Take a look at already implemented protocols for better overview. 89 | 90 | ### Register a subcommand 91 | 92 | Add import line and a subcommand to the main `cmd/cmd.go` file. 93 | 94 | ```go 95 | import ( 96 | // ... 97 | "github.com/facebookincubator/fbender/cmd/dhcpv4" 98 | "github.com/facebookincubator/fbender/cmd/dhcpv6" 99 | "github.com/facebookincubator/fbender/cmd/dns" 100 | "github.com/facebookincubator/fbender/cmd/http" 101 | // ... , add you command package here (alphabetically) 102 | "github.com/facebookincubator/fbender/cmd/tftp" 103 | "github.com/facebookincubator/fbender/cmd/udp" 104 | ) 105 | 106 | var Subcommands = []*cobra.Command{ 107 | dhcpv4.Command, 108 | dhcpv6.Command, 109 | dns.Command, 110 | http.Command, 111 | // ... , add you command here (alphabetically) 112 | tftp.Command, 113 | udp.Command, 114 | } 115 | ``` 116 | 117 | ## Adding internal features 118 | 119 | To adjust FBender to your needs without the necessity to create your own forks 120 | of the repository and always keep up to date with the newest version we 121 | recommend creating custom `main` function and "patching" your features into the 122 | main command. For example at Facebook we use it to add load tests for internal 123 | services and integrate our metric system for constraints tests. 124 | 125 | ```go 126 | package main 127 | 128 | import ( 129 | "github.com/facebookincubator/fbender/cmd" 130 | "github.com/facebookincubator/fbender/cmd/core" 131 | 132 | "fbender/internal/cmd" 133 | "fbender/internal/metric" 134 | ) 135 | 136 | func main() { 137 | // Add internal commands 138 | cmd.Command.AddCommand(cmd.Command) 139 | // Patch internal metrics 140 | core.ConstraintsValue.Parsers = append(core.ConstraintsValue.Parsers, metric.MetricParser) 141 | // Execute command 142 | cmd.Execute() 143 | } 144 | ``` 145 | -------------------------------------------------------------------------------- /flags/constraint.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "encoding/csv" 13 | "fmt" 14 | "strings" 15 | 16 | "github.com/facebookincubator/fbender/tester" 17 | "github.com/spf13/pflag" 18 | ) 19 | 20 | // ConstraintSliceValue is a pflag value storing constraints. 21 | type ConstraintSliceValue struct { 22 | Parsers []tester.MetricParser 23 | 24 | value *[]*tester.Constraint 25 | changed bool 26 | } 27 | 28 | // NewConstraintSliceValue creates a new constraint slice value for pflag. 29 | func NewConstraintSliceValue(parsers ...tester.MetricParser) *ConstraintSliceValue { 30 | v := []*tester.Constraint{} 31 | 32 | return &ConstraintSliceValue{ 33 | Parsers: parsers, 34 | value: &v, 35 | changed: false, 36 | } 37 | } 38 | 39 | func readAsCSV(val string) ([]string, error) { 40 | if val == "" { 41 | return []string{}, nil 42 | } 43 | 44 | stringReader := strings.NewReader(val) 45 | csvReader := csv.NewReader(stringReader) 46 | 47 | return csvReader.Read() 48 | } 49 | 50 | // Set validates given string given constraints and parses them to constraint 51 | // structures using metric parsers. 52 | func (c *ConstraintSliceValue) Set(value string) error { 53 | values, err := readAsCSV(value) 54 | if err != nil { 55 | return err 56 | } 57 | 58 | constraints := []*tester.Constraint{} 59 | 60 | for _, v := range values { 61 | constraint, err := tester.ParseConstraint(v, c.Parsers...) 62 | if err != nil { 63 | return fmt.Errorf("error parsing constraint %q: %w", v, err) 64 | } 65 | 66 | constraints = append(constraints, constraint) 67 | } 68 | 69 | if !c.changed { 70 | *c.value = constraints 71 | } else { 72 | *c.value = append(*c.value, constraints...) 73 | } 74 | 75 | c.changed = true 76 | 77 | return nil 78 | } 79 | 80 | // Type returns the ConstraintSliceValue Type. 81 | func (c *ConstraintSliceValue) Type() string { 82 | return "constraints" 83 | } 84 | 85 | func (c *ConstraintSliceValue) String() string { 86 | return fmt.Sprintf("%+v", *c.value) 87 | } 88 | 89 | // GetConstraints returns a constraints from a pflag set. 90 | func GetConstraints(f *pflag.FlagSet, name string) ([]*tester.Constraint, error) { 91 | flag := f.Lookup(name) 92 | if flag == nil { 93 | return nil, fmt.Errorf("%w: %q", ErrUndefined, name) 94 | } 95 | 96 | return GetConstraintsValue(flag.Value) 97 | } 98 | 99 | // GetConstraintsValue returns a constraints from a pflag value. 100 | func GetConstraintsValue(v pflag.Value) ([]*tester.Constraint, error) { 101 | if constraints, ok := v.(*ConstraintSliceValue); ok { 102 | return *constraints.value, nil 103 | } 104 | 105 | return nil, fmt.Errorf("%w, want: constraints, got: %s", ErrInvalidType, v.Type()) 106 | } 107 | -------------------------------------------------------------------------------- /flags/distribution.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "sort" 15 | "strings" 16 | 17 | "github.com/facebookincubator/fbender/utils" 18 | "github.com/pinterest/bender" 19 | "github.com/spf13/cobra" 20 | "github.com/spf13/pflag" 21 | ) 22 | 23 | // DistributionGenerator represents distribution generator function. 24 | type DistributionGenerator = func(float64) bender.IntervalGenerator 25 | 26 | const ( 27 | uniformGenerator = "uniform" 28 | exponentialGenerator = "exponential" 29 | ) 30 | 31 | //nolint:gochecknoglobals 32 | var generators = map[string]DistributionGenerator{ 33 | uniformGenerator: bender.UniformIntervalGenerator, 34 | exponentialGenerator: bender.ExponentialIntervalGenerator, 35 | } 36 | 37 | // Distribution represents a interval generator flag value. 38 | type Distribution struct { 39 | Name string 40 | generator DistributionGenerator 41 | } 42 | 43 | // NewDefaultDistribution returns new distribution flag with default values. 44 | func NewDefaultDistribution() *Distribution { 45 | return &Distribution{ 46 | Name: uniformGenerator, 47 | generator: generators[uniformGenerator], 48 | } 49 | } 50 | 51 | // ErrInvalidGenerator is raised when an unknown generator is set. 52 | var ErrInvalidGenerator = errors.New("invalid generator") 53 | 54 | // DistributionChoices returns a string representation of available generators. 55 | func DistributionChoices() []string { 56 | choices := []string{} 57 | 58 | for key := range generators { 59 | choices = append(choices, key) 60 | } 61 | 62 | sort.Strings(choices) 63 | 64 | return choices 65 | } 66 | 67 | func (d *Distribution) String() string { 68 | return d.Name 69 | } 70 | 71 | // Set validates a given value and sets distribution (allows prefix matching). 72 | func (d *Distribution) Set(value string) error { 73 | matches := []string{} 74 | 75 | for key := range generators { 76 | if strings.HasPrefix(key, value) { 77 | matches = append(matches, key) 78 | } 79 | } 80 | 81 | if len(matches) == 0 { 82 | choices := ChoicesString(DistributionChoices()) 83 | 84 | return fmt.Errorf("%w, want: %s, got: %q", ErrInvalidGenerator, choices, value) 85 | } else if len(matches) > 1 { 86 | sort.Strings(matches) 87 | 88 | return fmt.Errorf("%w, ambiguous prefix %q matches: %s", ErrInvalidGenerator, value, ChoicesString(matches)) 89 | } 90 | 91 | generator := matches[0] 92 | d.Name = generator 93 | d.generator = generators[generator] 94 | 95 | return nil 96 | } 97 | 98 | // Type returns a distribution type. 99 | func (d *Distribution) Type() string { 100 | return "distribution" 101 | } 102 | 103 | // Get returns a distibution generator. 104 | func (d *Distribution) Get() DistributionGenerator { 105 | return d.generator 106 | } 107 | 108 | // GetDistribution returns a distribution from a pflag set. 109 | func GetDistribution(f *pflag.FlagSet, name string) (DistributionGenerator, error) { 110 | flag := f.Lookup(name) 111 | if flag == nil { 112 | return nil, fmt.Errorf("%w: %q", ErrUndefined, name) 113 | } 114 | 115 | return GetDistributionValue(flag.Value) 116 | } 117 | 118 | // GetDistributionValue returns a distribution from a pflag value. 119 | func GetDistributionValue(v pflag.Value) (DistributionGenerator, error) { 120 | if distribution, ok := v.(*Distribution); ok { 121 | return distribution.Get(), nil 122 | } 123 | 124 | return nil, fmt.Errorf("%w, want: distribution, got: %s", ErrInvalidType, v.Type()) 125 | } 126 | 127 | // Bash completion function constants. 128 | const ( 129 | fnameDistribution = "__fbender_handle_distribution_flag" 130 | fbodyDistribution = `COMPREPLY=($(compgen -W "uniform exponential" -- "${cur}"))` 131 | ) 132 | 133 | // BashCompletionDistribution adds bash completion to a distribution flag. 134 | func BashCompletionDistribution(cmd *cobra.Command, f *pflag.FlagSet, name string) error { 135 | flag := f.Lookup(name) 136 | if flag == nil { 137 | return fmt.Errorf("%w: %q", ErrUndefined, name) 138 | } 139 | 140 | if _, ok := flag.Value.(*Distribution); !ok { 141 | return fmt.Errorf("%w, want: distribution, got: %s", ErrInvalidType, flag.Value.Type()) 142 | } 143 | 144 | return utils.BashCompletion(cmd, f, name, fnameDistribution, fbodyDistribution) 145 | } 146 | -------------------------------------------------------------------------------- /flags/errors.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "errors" 13 | ) 14 | 15 | // ErrUndefined is raised when user tries to access an undefined flag. 16 | var ErrUndefined = errors.New("flag accessed but not defined") 17 | 18 | // ErrInvalidType is raised when user tries to get a value from a flag of a different type. 19 | var ErrInvalidType = errors.New("accessed flag type does not match") 20 | -------------------------------------------------------------------------------- /flags/flags_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags_test 10 | 11 | import ( 12 | "reflect" 13 | "testing" 14 | 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | // assertPointerEqual checks whether two pointers are equal. 19 | func assertPointerEqual(t *testing.T, expected, actual interface{}, args ...interface{}) { 20 | t.Helper() 21 | 22 | expectedPointer := reflect.ValueOf(expected).Pointer() 23 | actualPointer := reflect.ValueOf(actual).Pointer() 24 | assert.Equal(t, expectedPointer, actualPointer, args...) 25 | } 26 | -------------------------------------------------------------------------------- /flags/growth.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "fmt" 13 | 14 | "github.com/facebookincubator/fbender/tester" 15 | "github.com/spf13/pflag" 16 | ) 17 | 18 | // GrowthValue represents growth flag value. 19 | type GrowthValue struct { 20 | Growth tester.Growth 21 | } 22 | 23 | func (g *GrowthValue) String() string { 24 | return g.Growth.String() 25 | } 26 | 27 | // Set validates a given growth and saves it. 28 | func (g *GrowthValue) Set(value string) error { 29 | var err error 30 | 31 | g.Growth, err = tester.ParseGrowth(value) 32 | if err != nil { 33 | return fmt.Errorf("error parsing growth %q: %w", value, err) 34 | } 35 | 36 | return nil 37 | } 38 | 39 | // Type returns a growth value type. 40 | func (g *GrowthValue) Type() string { 41 | return "growth" 42 | } 43 | 44 | // GetGrowth returns a growth from a pflag set. 45 | func GetGrowth(f *pflag.FlagSet, name string) (tester.Growth, error) { 46 | flag := f.Lookup(name) 47 | if flag == nil { 48 | return nil, fmt.Errorf("%w: %q", ErrUndefined, name) 49 | } 50 | 51 | return GetGrowthValue(flag.Value) 52 | } 53 | 54 | // GetGrowthValue returns a growth from a pflag value. 55 | func GetGrowthValue(v pflag.Value) (tester.Growth, error) { 56 | if growth, ok := v.(*GrowthValue); ok { 57 | return growth.Growth, nil 58 | } 59 | 60 | return nil, fmt.Errorf("%w, want: *GrowthValue, got: %T", ErrInvalidType, v) 61 | } 62 | -------------------------------------------------------------------------------- /flags/growth_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/flags" 15 | "github.com/facebookincubator/fbender/tester" 16 | "github.com/spf13/pflag" 17 | "github.com/stretchr/testify/assert" 18 | "github.com/stretchr/testify/require" 19 | ) 20 | 21 | func TestGrowth__String(t *testing.T) { 22 | linGrowth := &tester.LinearGrowth{Increase: 100} 23 | value := &flags.GrowthValue{Growth: linGrowth} 24 | assert.Equal(t, linGrowth.String(), value.String()) 25 | // Change growth values 26 | linGrowth.Increase = 200 27 | assert.Equal(t, linGrowth.String(), value.String()) 28 | // Change growth type 29 | perGrowth := &tester.PercentageGrowth{Increase: 100.} 30 | value.Growth = perGrowth 31 | assert.Equal(t, perGrowth.String(), value.String()) 32 | } 33 | 34 | func TestGrowth__Set(t *testing.T) { 35 | value := new(flags.GrowthValue) 36 | // Valid LinearGrowth string 37 | err := value.Set("+200") 38 | require.NoError(t, err) 39 | assert.IsType(t, new(tester.LinearGrowth), value.Growth) 40 | // Valid PercentageGrowth string 41 | err = value.Set("%100.0") 42 | require.NoError(t, err) 43 | assert.IsType(t, new(tester.PercentageGrowth), value.Growth) 44 | // Valid ExponentialGrowth false 45 | err = value.Set("^25") 46 | require.NoError(t, err) 47 | assert.IsType(t, new(tester.ExponentialGrowth), value.Growth) 48 | // Unknown prefix 49 | err = value.Set("@200") 50 | assert.ErrorIs(t, err, tester.ErrInvalidGrowth) 51 | assert.EqualError(t, err, "error parsing growth \"@200\": unknown growth, want +int, %%flaot, ^int") 52 | } 53 | 54 | func TestGrowth__Type(t *testing.T) { 55 | value := &flags.GrowthValue{Growth: &tester.LinearGrowth{Increase: 100}} 56 | assert.Equal(t, "growth", value.Type()) 57 | } 58 | 59 | func TestGetGrowth(t *testing.T) { 60 | value := &flags.GrowthValue{Growth: &tester.LinearGrowth{Increase: 100}} 61 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 62 | f.Var(value, "growth", "set growth") 63 | err := value.Set("+200") 64 | require.NoError(t, err) 65 | growth, err := flags.GetGrowth(f, "growth") 66 | require.NoError(t, err) 67 | assert.IsType(t, new(tester.LinearGrowth), growth) 68 | // Check if value changes 69 | err = value.Set("%100") 70 | require.NoError(t, err) 71 | growth, err = flags.GetGrowth(f, "growth") 72 | require.NoError(t, err) 73 | assert.IsType(t, new(tester.PercentageGrowth), growth) 74 | // Check if error when flag does not exist 75 | _, err = flags.GetGrowth(f, "nonexistent") 76 | assert.ErrorIs(t, err, flags.ErrUndefined) 77 | assert.EqualError(t, err, "flag accessed but not defined: \"nonexistent\"") 78 | // Check if error when value is of different type 79 | f.Int("myint", 0, "set myint") 80 | _, err = flags.GetGrowth(f, "myint") 81 | assert.ErrorIs(t, err, flags.ErrInvalidType) 82 | assert.EqualError(t, err, "accessed flag type does not match, want: *GrowthValue, got: *pflag.intValue") 83 | } 84 | 85 | func TestGetGrowthValue(t *testing.T) { 86 | value := new(flags.GrowthValue) 87 | err := value.Set("+200") 88 | require.NoError(t, err) 89 | growth, err := flags.GetGrowthValue(value) 90 | require.NoError(t, err) 91 | assert.IsType(t, new(tester.LinearGrowth), growth) 92 | // Check if value changes 93 | err = value.Set("%100") 94 | require.NoError(t, err) 95 | growth, err = flags.GetGrowthValue(value) 96 | require.NoError(t, err) 97 | assert.IsType(t, new(tester.PercentageGrowth), growth) 98 | // Check if error when value is of different type 99 | f := pflag.NewFlagSet("Test FlagSet", pflag.ExitOnError) 100 | f.Int("myint", 0, "set myint") 101 | flag := f.Lookup("myint") 102 | require.NotNil(t, flag) 103 | _, err = flags.GetGrowthValue(flag.Value) 104 | assert.ErrorIs(t, err, flags.ErrInvalidType) 105 | assert.EqualError(t, err, "accessed flag type does not match, want: *GrowthValue, got: *pflag.intValue") 106 | } 107 | -------------------------------------------------------------------------------- /flags/log.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "os" 15 | "sort" 16 | "strconv" 17 | "strings" 18 | 19 | "github.com/facebookincubator/fbender/utils" 20 | "github.com/sirupsen/logrus" 21 | "github.com/spf13/cobra" 22 | "github.com/spf13/pflag" 23 | ) 24 | 25 | // LogLevel represents a log level flag value. 26 | type LogLevel struct { 27 | Logger *logrus.Logger 28 | } 29 | 30 | // LogLevelChoices returns a string representation of available levels. 31 | func LogLevelChoices() []string { 32 | choices := []string{} 33 | for _, level := range logrus.AllLevels { 34 | choices = append(choices, level.String()) 35 | } 36 | 37 | return choices 38 | } 39 | 40 | // ErrInvalidLogLevel is raised when an unknown level is set. 41 | var ErrInvalidLogLevel = errors.New("invalid log level") 42 | 43 | func (l *LogLevel) String() string { 44 | return l.Logger.Level.String() 45 | } 46 | 47 | // Set validates a given value and sets log level. 48 | func (l *LogLevel) Set(value string) error { 49 | level, err := logrus.ParseLevel(value) 50 | if err != nil { 51 | v, err := strconv.ParseUint(value, 10, 32) 52 | if err != nil || v >= uint64(len(logrus.AllLevels)) { 53 | choices := ChoicesString(LogLevelChoices()) 54 | 55 | return fmt.Errorf("%w, want: integer [0..%d] or %s, got: %s", 56 | ErrInvalidLogLevel, len(logrus.AllLevels)-1, choices, value) 57 | } 58 | 59 | level = logrus.Level(v) 60 | } 61 | 62 | l.Logger.SetLevel(level) 63 | 64 | return nil 65 | } 66 | 67 | // Type returns a log level value type. 68 | func (l *LogLevel) Type() string { 69 | return "level" 70 | } 71 | 72 | // Bash completion function constants. 73 | const ( 74 | fnameLogLevel = "__fbender_handle_loglevel_flag" 75 | fbodyLogLevel = `COMPREPLY=($(compgen -W "%s" -- "${cur}"))` 76 | ) 77 | 78 | // BashCompletionLogLevel adds bash completion to a distribution flag. 79 | func BashCompletionLogLevel(cmd *cobra.Command, f *pflag.FlagSet, name string) error { 80 | flag := f.Lookup(name) 81 | if flag == nil { 82 | return fmt.Errorf("%w: %q", ErrUndefined, name) 83 | } 84 | 85 | if _, ok := flag.Value.(*LogLevel); !ok { 86 | return fmt.Errorf("%w, want: level, got: %s", ErrInvalidType, flag.Value.Type()) 87 | } 88 | 89 | fbody := fmt.Sprintf(fbodyLogLevel, strings.Join(LogLevelChoices(), " ")) 90 | 91 | return utils.BashCompletion(cmd, f, name, fnameLogLevel, fbody) 92 | } 93 | 94 | // LogFormat represents a log format flag value. 95 | type LogFormat struct { 96 | Logger *logrus.Logger 97 | Format string 98 | } 99 | 100 | const ( 101 | testFormatter = "text" 102 | jsonFormatter = "json" 103 | ) 104 | 105 | // We need to store a lambda function which generates formatters so we don't 106 | // assign the same static formatter to different Loggers. 107 | //nolint:gochecknoglobals 108 | var formatters = map[string](func() logrus.Formatter){ 109 | testFormatter: func() logrus.Formatter { return new(logrus.TextFormatter) }, 110 | jsonFormatter: func() logrus.Formatter { return new(logrus.JSONFormatter) }, 111 | } 112 | 113 | // LogFormatChoices returns a string representation of available formats. 114 | func LogFormatChoices() []string { 115 | choices := []string{} 116 | for key := range formatters { 117 | choices = append(choices, key) 118 | } 119 | 120 | sort.Strings(choices) 121 | 122 | return choices 123 | } 124 | 125 | // ErrInvalidLogFormat is raised when an unknown format is set. 126 | var ErrInvalidLogFormat = errors.New("invalid log format") 127 | 128 | func (l *LogFormat) String() string { 129 | return l.Format 130 | } 131 | 132 | // Set validates a given value and sets log format. 133 | func (l *LogFormat) Set(value string) error { 134 | if formatter, ok := formatters[value]; ok { 135 | l.Logger.Formatter = formatter() 136 | l.Format = value 137 | 138 | return nil 139 | } 140 | 141 | return fmt.Errorf("%w, want: %s, got: %q", 142 | ErrInvalidLogFormat, ChoicesString(LogFormatChoices()), value) 143 | } 144 | 145 | // Type returns a log format value type. 146 | func (l *LogFormat) Type() string { 147 | return "format" 148 | } 149 | 150 | // Bash completion function constants. 151 | const ( 152 | fnameLogFormat = "__fbender_handle_logformat_flag" 153 | fbodyLogFormat = `COMPREPLY=($(compgen -W "%s" -- "${cur}"))` 154 | ) 155 | 156 | // BashCompletionLogFormat adds bash completion to a distribution flag. 157 | func BashCompletionLogFormat(cmd *cobra.Command, f *pflag.FlagSet, name string) error { 158 | flag := f.Lookup(name) 159 | if flag == nil { 160 | return fmt.Errorf("%w: %q", ErrUndefined, name) 161 | } 162 | 163 | if _, ok := flag.Value.(*LogFormat); !ok { 164 | return fmt.Errorf("%w, want: format, got: %s", ErrInvalidType, flag.Value.Type()) 165 | } 166 | 167 | fbody := fmt.Sprintf(fbodyLogFormat, strings.Join(LogFormatChoices(), " ")) 168 | 169 | return utils.BashCompletion(cmd, f, name, fnameLogFormat, fbody) 170 | } 171 | 172 | // LogOutput represents a log output flag value. 173 | type LogOutput struct { 174 | Logger *logrus.Logger 175 | Filename string 176 | Out *os.File 177 | } 178 | 179 | // NewLogOutput returns new log output flag Value. 180 | func NewLogOutput(logger *logrus.Logger) *LogOutput { 181 | logger.Out = os.Stdout 182 | 183 | return &LogOutput{ 184 | Logger: logger, 185 | Out: os.Stdout, 186 | } 187 | } 188 | 189 | func (l *LogOutput) String() string { 190 | if l.Out == os.Stdout { 191 | return "" 192 | } else if l.Out == os.Stderr { 193 | return "" 194 | } 195 | 196 | return l.Out.Name() 197 | } 198 | 199 | // Set opens the file and sets the Out of the Logger. 200 | func (l *LogOutput) Set(value string) error { 201 | // Close any file we have been writing to. 202 | if l.Out != nil && l.Out != os.Stdout && l.Out != os.Stderr { 203 | if err := l.Out.Close(); err != nil { 204 | return fmt.Errorf("unable to close output %q: %w", l.Out.Name(), err) 205 | } 206 | } 207 | 208 | // If the value is not an empty string it is a file name. 209 | if len(value) > 0 { 210 | file, err := os.Create(value) 211 | if err != nil { 212 | return fmt.Errorf("unable to create output %q: %w", value, err) 213 | } 214 | 215 | l.Out = file 216 | } else { 217 | l.Out = os.Stdout 218 | } 219 | 220 | l.Logger.Out = l.Out 221 | 222 | return nil 223 | } 224 | 225 | // Type returns a log output value type. 226 | func (l *LogOutput) Type() string { 227 | return "output" 228 | } 229 | -------------------------------------------------------------------------------- /flags/utils.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags 10 | 11 | import ( 12 | "fmt" 13 | "strings" 14 | ) 15 | 16 | // ChoicesString converts choices list into printable text. 17 | func ChoicesString(choices []string) string { 18 | return fmt.Sprintf("(%s)", strings.Join(choices, "|")) 19 | } 20 | -------------------------------------------------------------------------------- /flags/utils_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package flags_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/flags" 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | func TestChoicesString(t *testing.T) { 19 | // Empty list 20 | assert.Equal(t, "()", flags.ChoicesString([]string{})) 21 | // Single element list 22 | assert.Equal(t, "(one)", flags.ChoicesString([]string{"one"})) 23 | // Multiple elements list 24 | assert.Equal(t, "(one|two)", flags.ChoicesString([]string{"one", "two"})) 25 | assert.Equal(t, "(one|two|three)", flags.ChoicesString([]string{"one", "two", "three"})) 26 | } 27 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/facebookincubator/fbender 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/gosuri/uilive v0.0.4 // indirect 7 | github.com/gosuri/uiprogress v0.0.1 8 | github.com/insomniacslk/dhcp v0.0.0-20201112113307-4de412bc85d8 9 | github.com/mdlayher/netx v0.0.0-20200512211805-669a06fde734 // indirect 10 | github.com/miekg/dns v1.1.35 11 | github.com/pin/tftp v0.0.0-20200229063000-e4f073737eb2 12 | github.com/pinterest/bender v0.0.0-20201102205149-897b051c8257 13 | github.com/sirupsen/logrus v1.7.0 14 | github.com/spf13/cobra v1.1.1 15 | github.com/spf13/pflag v1.0.5 16 | github.com/stretchr/objx v0.3.0 // indirect 17 | github.com/stretchr/testify v1.6.2-0.20201103103935-92707c0b2d50 18 | github.com/tj/go-spin v1.1.0 19 | gopkg.in/yaml.v3 v3.0.0 // indirect 20 | ) 21 | -------------------------------------------------------------------------------- /log/log.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | /* 10 | Package log implements standard logging. 11 | */ 12 | package log 13 | 14 | import ( 15 | "fmt" 16 | "io" 17 | "os" 18 | ) 19 | 20 | // Standard writers for the log package. 21 | //nolint:gochecknoglobals 22 | var ( 23 | Stderr io.Writer = os.Stderr 24 | Stdout io.Writer = os.Stdout 25 | ) 26 | 27 | // Fprintf formats according to a format specifier and writes to w. 28 | // It panics if any write error encountered. 29 | func Fprintf(w io.Writer, format string, args ...interface{}) { 30 | if _, err := fmt.Fprintf(w, format, args...); err != nil { 31 | panic(err) 32 | } 33 | } 34 | 35 | // Printf formats according to a format specifier and writes to standard output. 36 | // It panics if any write error encountered. 37 | func Printf(format string, args ...interface{}) { 38 | Fprintf(Stdout, format, args...) 39 | } 40 | 41 | // Errorf formats according to a format specifier and writes to standard error. 42 | // It panics if any write error encountered. 43 | func Errorf(format string, args ...interface{}) { 44 | Fprintf(Stderr, format, args...) 45 | } 46 | -------------------------------------------------------------------------------- /log/log_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package log_test 10 | 11 | import ( 12 | "bytes" 13 | "testing" 14 | 15 | "github.com/facebookincubator/fbender/log" 16 | "github.com/stretchr/testify/assert" 17 | "github.com/stretchr/testify/mock" 18 | ) 19 | 20 | type MockedWriter struct { 21 | mock.Mock 22 | } 23 | 24 | func (m *MockedWriter) Write(p []byte) (int, error) { 25 | args := m.Called(p) 26 | 27 | return args.Int(0), args.Error(1) 28 | } 29 | 30 | func TestFprintf_NoError(t *testing.T) { 31 | w := new(bytes.Buffer) 32 | log.Fprintf(w, "Hello %s", "World") 33 | assert.Equal(t, "Hello World", w.String()) 34 | } 35 | 36 | func TestFprintf_Error(t *testing.T) { 37 | w := new(MockedWriter) 38 | w.On("Write", "Hello World").Return(0, assert.AnError) 39 | 40 | defer func() { 41 | if r := recover(); r == nil { 42 | assert.Fail(t, "Expected Fprintf to panic") 43 | } 44 | }() 45 | 46 | log.Fprintf(w, "Hello %s", "World") 47 | } 48 | 49 | func TestPrintf_NoError(t *testing.T) { 50 | w := new(bytes.Buffer) 51 | log.Stdout = w 52 | log.Printf("Hello %s", "World") 53 | assert.Equal(t, "Hello World", w.String()) 54 | } 55 | 56 | func TestPrintf_Error(t *testing.T) { 57 | w := new(MockedWriter) 58 | log.Stdout = w 59 | w.On("Write", "Hello World").Return(0, assert.AnError) 60 | 61 | defer func() { 62 | if r := recover(); r == nil { 63 | assert.Fail(t, "Expected Fprintf to panic") 64 | } 65 | }() 66 | 67 | log.Printf("Hello %s", "World") 68 | } 69 | 70 | func TestErrorf_NoError(t *testing.T) { 71 | w := new(bytes.Buffer) 72 | log.Stderr = w 73 | log.Errorf("Hello %s", "World") 74 | assert.Equal(t, "Hello World", w.String()) 75 | } 76 | 77 | func TestErrorf_Error(t *testing.T) { 78 | w := new(MockedWriter) 79 | log.Stderr = w 80 | w.On("Write", "Hello World").Return(0, assert.AnError) 81 | 82 | defer func() { 83 | if r := recover(); r == nil { 84 | assert.Fail(t, "Expected Fprintf to panic") 85 | } 86 | }() 87 | 88 | log.Errorf("Hello %s", "World") 89 | } 90 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package main 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/cmd" 13 | ) 14 | 15 | func main() { 16 | cmd.Execute() 17 | } 18 | -------------------------------------------------------------------------------- /metric/errors.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package metric 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/recorders" 15 | "github.com/facebookincubator/fbender/tester" 16 | "github.com/pinterest/bender" 17 | ) 18 | 19 | // ErrorsMetric fetches data from statistics. 20 | type ErrorsMetric struct { 21 | Statistics recorders.Statistics 22 | } 23 | 24 | // ErrorsMetricOptions represents errors metric options. 25 | type ErrorsMetricOptions interface { 26 | AddRecorder(bender.Recorder) 27 | } 28 | 29 | // Setup prepares errors metric. 30 | func (m *ErrorsMetric) Setup(options interface{}) error { 31 | opts, ok := options.(ErrorsMetricOptions) 32 | if !ok { 33 | return tester.ErrInvalidOptions 34 | } 35 | 36 | opts.AddRecorder(recorders.NewStatisticsRecorder(&m.Statistics)) 37 | 38 | return nil 39 | } 40 | 41 | // Fetch calculates the errors percentage for the statistics. 42 | func (m *ErrorsMetric) Fetch(start time.Time, duration time.Duration) ([]tester.DataPoint, error) { 43 | errorsPct := float64(m.Statistics.Errors) / float64(m.Statistics.Requests) * 100.0 44 | 45 | // return a single point with time equal to end of the test 46 | return []tester.DataPoint{ 47 | {Time: start.Add(duration), Value: errorsPct}, 48 | }, nil 49 | } 50 | 51 | // Name returns the name of the errors statistic. 52 | func (m *ErrorsMetric) Name() string { 53 | return "errors" 54 | } 55 | -------------------------------------------------------------------------------- /metric/latency.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package metric 10 | 11 | import ( 12 | "sync" 13 | "time" 14 | 15 | "github.com/facebookincubator/fbender/tester" 16 | "github.com/pinterest/bender" 17 | ) 18 | 19 | // LatencyMetric fetches data from statistics. 20 | type LatencyMetric struct { 21 | mutex sync.Mutex 22 | points []tester.DataPoint 23 | } 24 | 25 | // LatencyMetricOptions represents errors metric options. 26 | type LatencyMetricOptions interface { 27 | AddRecorder(bender.Recorder) 28 | GetUnit() time.Duration 29 | } 30 | 31 | // Setup prepares errors metric. 32 | func (m *LatencyMetric) Setup(options interface{}) error { 33 | opts, ok := options.(LatencyMetricOptions) 34 | if !ok { 35 | return tester.ErrInvalidOptions 36 | } 37 | 38 | unit := opts.GetUnit() 39 | opts.AddRecorder(func(msg interface{}) { 40 | switch msg := msg.(type) { 41 | case *bender.StartEvent: 42 | m.points = make([]tester.DataPoint, 0) 43 | case *bender.EndRequestEvent: 44 | m.mutex.Lock() 45 | m.points = append(m.points, tester.DataPoint{ 46 | Time: time.Unix(msg.Start, 0), 47 | Value: float64(msg.End-msg.Start) / float64(unit), 48 | }) 49 | m.mutex.Unlock() 50 | } 51 | }) 52 | 53 | return nil 54 | } 55 | 56 | // Fetch calculates the errors percentage for the statistics. 57 | func (m *LatencyMetric) Fetch(start time.Time, duration time.Duration) ([]tester.DataPoint, error) { 58 | return m.points, nil 59 | } 60 | 61 | // Name returns the name of the errors statistic. 62 | func (m *LatencyMetric) Name() string { 63 | return "latency" 64 | } 65 | -------------------------------------------------------------------------------- /metric/parser.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package metric 10 | 11 | import ( 12 | "github.com/facebookincubator/fbender/tester" 13 | ) 14 | 15 | // Help is a help message on available metrics. 16 | const Help = ` 17 | Basic Metrics: 18 | * errors - errors percentage of all requests, ignores aggregator 19 | * latency - latency of the packets (in unit specified by --unit) 20 | MAX(errors) < 10.0 21 | MIN(errors) < 42.0 22 | AVG(latency) < 30` 23 | 24 | // Parser is a parser for standard metrics. 25 | func Parser(value string) (tester.Metric, error) { 26 | switch value { 27 | case "errors": 28 | return new(ErrorsMetric), nil 29 | case "latency": 30 | return new(LatencyMetric), nil 31 | default: 32 | return nil, tester.ErrNotParsed 33 | } 34 | } 35 | -------------------------------------------------------------------------------- /protocols/udp/udp.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package udp 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "net" 15 | "strconv" 16 | "time" 17 | 18 | "github.com/facebookincubator/fbender/log" 19 | "github.com/pinterest/bender" 20 | ) 21 | 22 | // MaxResponseSize is the max response size for UDP test. 23 | const MaxResponseSize = 2048 24 | 25 | // Datagram represents a udp datagram to be sent. 26 | type Datagram struct { 27 | Port int 28 | Data []byte 29 | } 30 | 31 | // ErrInvalidType is raised when object type mismatch. 32 | var ErrInvalidType = errors.New("invalid type") 33 | 34 | // ResponseValidator validates a udp response. 35 | type ResponseValidator func(request *Datagram, response []byte) error 36 | 37 | // CreateExecutor creates a new UDP RequestExecutor. 38 | func CreateExecutor(timeout time.Duration, validator ResponseValidator, hosts ...string) bender.RequestExecutor { 39 | var i int 40 | 41 | return func(_ int64, request interface{}) (interface{}, error) { 42 | datagram, ok := request.(*Datagram) 43 | if !ok { 44 | return nil, fmt.Errorf("%w, want: *Datagram, got: %T", ErrInvalidType, request) 45 | } 46 | 47 | addr := net.JoinHostPort(hosts[i], strconv.Itoa(datagram.Port)) 48 | i = (i + 1) % len(hosts) 49 | 50 | // Setup connection 51 | conn, err := net.Dial("udp", addr) 52 | if err != nil { 53 | //nolint:wrapcheck 54 | return nil, err 55 | } 56 | 57 | defer func() { 58 | if err = conn.Close(); err != nil { 59 | log.Errorf("Error closing connection: %v\n", err) 60 | } 61 | }() 62 | 63 | if err = conn.SetWriteDeadline(time.Now().Add(timeout)); err != nil { 64 | //nolint:wrapcheck 65 | return nil, err 66 | } 67 | 68 | _, err = conn.Write(datagram.Data) 69 | if err != nil { 70 | //nolint:wrapcheck 71 | return nil, err 72 | } 73 | 74 | buffer := make([]byte, MaxResponseSize) 75 | 76 | if err = conn.SetReadDeadline(time.Now().Add(timeout)); err != nil { 77 | //nolint:wrapcheck 78 | return nil, err 79 | } 80 | 81 | n, err := conn.Read(buffer) 82 | if err != nil { 83 | //nolint:wrapcheck 84 | return nil, err 85 | } 86 | 87 | if err = validator(datagram, buffer[:n]); err != nil { 88 | return nil, err 89 | } 90 | 91 | return buffer[:n], nil 92 | } 93 | } 94 | -------------------------------------------------------------------------------- /recorders/logrus.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders 10 | 11 | import ( 12 | "github.com/pinterest/bender" 13 | "github.com/sirupsen/logrus" 14 | ) 15 | 16 | // NewLogrusRecorder creates a new logrus.Logger-based recorder. 17 | func NewLogrusRecorder(l *logrus.Logger, defaults ...logrus.Fields) bender.Recorder { 18 | return func(msg interface{}) { 19 | log := logrus.NewEntry(l) 20 | 21 | for _, fields := range defaults { 22 | for key, value := range fields { 23 | log = log.WithField(key, value) 24 | } 25 | } 26 | 27 | switch msg := msg.(type) { 28 | case *bender.StartRequestEvent: 29 | logStartRequestEvent(log, msg) 30 | case *bender.EndRequestEvent: 31 | logEndRequestEvent(log, msg) 32 | } 33 | } 34 | } 35 | 36 | func logStartRequestEvent(log *logrus.Entry, msg *bender.StartRequestEvent) { 37 | log.WithFields(logrus.Fields{ 38 | "start": msg.Time, 39 | "request": msg.Request, 40 | }).Debug("Start") 41 | } 42 | 43 | func logEndRequestEvent(log *logrus.Entry, msg *bender.EndRequestEvent) { 44 | log = log.WithFields(logrus.Fields{ 45 | "start": msg.Start, 46 | "end": msg.End, 47 | "elapsed": int(msg.End - msg.Start), 48 | "response": msg.Response, 49 | }) 50 | 51 | if msg.Err != nil { 52 | log.WithError(msg.Err).Warn("Fail") 53 | } else { 54 | log.Info("Success") 55 | } 56 | } 57 | -------------------------------------------------------------------------------- /recorders/logrus_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/recorders" 15 | "github.com/pinterest/bender" 16 | "github.com/sirupsen/logrus" 17 | "github.com/sirupsen/logrus/hooks/test" 18 | "github.com/stretchr/testify/assert" 19 | "github.com/stretchr/testify/suite" 20 | ) 21 | 22 | type LogrusRecorderTestSuite struct { 23 | suite.Suite 24 | logger *logrus.Logger 25 | hook *test.Hook 26 | recorder chan interface{} 27 | logrusRecorder bender.Recorder 28 | } 29 | 30 | func (s *LogrusRecorderTestSuite) SetupSuite() { 31 | s.logger, s.hook = test.NewNullLogger() 32 | s.logger.Level = logrus.DebugLevel 33 | } 34 | 35 | func (s *LogrusRecorderTestSuite) SetupTest() { 36 | s.hook.Reset() 37 | s.recorder = make(chan interface{}, 1) 38 | s.logrusRecorder = recorders.NewLogrusRecorder(s.logger) 39 | } 40 | 41 | func (s *LogrusRecorderTestSuite) recordSingleEvent(event interface{}) { 42 | s.recorder <- event 43 | close(s.recorder) 44 | bender.Record(s.recorder, s.logrusRecorder) 45 | } 46 | 47 | func (s *LogrusRecorderTestSuite) TestStartEvent() { 48 | s.recordSingleEvent(new(bender.StartEvent)) 49 | s.Assert().Len(s.hook.Entries, 0) 50 | } 51 | 52 | func (s *LogrusRecorderTestSuite) TestEndEvent() { 53 | s.recordSingleEvent(new(bender.EndEvent)) 54 | s.Assert().Len(s.hook.Entries, 0) 55 | } 56 | 57 | func (s *LogrusRecorderTestSuite) TestWaitEvent() { 58 | s.recordSingleEvent(new(bender.WaitEvent)) 59 | s.Assert().Len(s.hook.Entries, 0) 60 | } 61 | 62 | func (s *LogrusRecorderTestSuite) TestStartRequestEvent() { 63 | event := &bender.StartRequestEvent{ 64 | Time: 420, 65 | Request: "my request", 66 | } 67 | s.recordSingleEvent(event) 68 | s.Require().Len(s.hook.Entries, 1) 69 | s.Assert().Equal(logrus.DebugLevel, s.hook.LastEntry().Level) 70 | s.Assert().Equal("Start", s.hook.LastEntry().Message) 71 | s.Assert().Equal(logrus.Fields{ 72 | "start": int64(420), 73 | "request": "my request", 74 | }, s.hook.LastEntry().Data) 75 | 76 | s.hook.Reset() 77 | s.recorder = make(chan interface{}, 1) 78 | s.logrusRecorder = recorders.NewLogrusRecorder(s.logger, logrus.Fields{ 79 | "keyA0": "valueA0", 80 | "keyA1": "valueA1", 81 | }, logrus.Fields{ 82 | "keyB0": "valueB0", 83 | "keyB1": "valueB1", 84 | }) 85 | s.recordSingleEvent(event) 86 | s.Require().Len(s.hook.Entries, 1) 87 | s.Assert().Equal(logrus.DebugLevel, s.hook.LastEntry().Level) 88 | s.Assert().Equal("Start", s.hook.LastEntry().Message) 89 | s.Assert().Equal(logrus.Fields{ 90 | "keyA0": "valueA0", 91 | "keyA1": "valueA1", 92 | "keyB0": "valueB0", 93 | "keyB1": "valueB1", 94 | "start": int64(420), 95 | "request": "my request", 96 | }, s.hook.LastEntry().Data) 97 | } 98 | 99 | func (s *LogrusRecorderTestSuite) TestEndRequestEvent_NoError() { 100 | event := &bender.EndRequestEvent{ 101 | Start: 420, 102 | End: 4200, 103 | Response: "my response", 104 | Err: nil, 105 | } 106 | s.recordSingleEvent(event) 107 | s.Require().Len(s.hook.Entries, 1) 108 | s.Assert().Equal(logrus.InfoLevel, s.hook.LastEntry().Level) 109 | s.Assert().Equal("Success", s.hook.LastEntry().Message) 110 | s.Assert().Equal(logrus.Fields{ 111 | "start": int64(420), 112 | "end": int64(4200), 113 | "elapsed": 3780, 114 | "response": "my response", 115 | }, s.hook.LastEntry().Data) 116 | 117 | s.hook.Reset() 118 | s.recorder = make(chan interface{}, 1) 119 | s.logrusRecorder = recorders.NewLogrusRecorder(s.logger, logrus.Fields{ 120 | "keyA0": "valueA0", 121 | "keyA1": "valueA1", 122 | }, logrus.Fields{ 123 | "keyB0": "valueB0", 124 | "keyB1": "valueB1", 125 | }) 126 | s.recordSingleEvent(event) 127 | s.Require().Len(s.hook.Entries, 1) 128 | s.Assert().Equal(logrus.InfoLevel, s.hook.LastEntry().Level) 129 | s.Assert().Equal("Success", s.hook.LastEntry().Message) 130 | s.Assert().Equal(logrus.Fields{ 131 | "keyA0": "valueA0", 132 | "keyA1": "valueA1", 133 | "keyB0": "valueB0", 134 | "keyB1": "valueB1", 135 | "start": int64(420), 136 | "end": int64(4200), 137 | "elapsed": 3780, 138 | "response": "my response", 139 | }, s.hook.LastEntry().Data) 140 | } 141 | 142 | func (s *LogrusRecorderTestSuite) TestEndRequestEvent_Error() { 143 | event := &bender.EndRequestEvent{ 144 | Start: 420, 145 | End: 4200, 146 | Response: "invalid response", 147 | Err: assert.AnError, 148 | } 149 | s.recordSingleEvent(event) 150 | s.Require().Len(s.hook.Entries, 1) 151 | s.Assert().Equal(logrus.WarnLevel, s.hook.LastEntry().Level) 152 | s.Assert().Equal("Fail", s.hook.LastEntry().Message) 153 | s.Assert().Equal(logrus.Fields{ 154 | "start": int64(420), 155 | "end": int64(4200), 156 | "elapsed": 3780, 157 | "response": "invalid response", 158 | "error": assert.AnError, 159 | }, s.hook.LastEntry().Data) 160 | 161 | s.hook.Reset() 162 | s.recorder = make(chan interface{}, 1) 163 | s.logrusRecorder = recorders.NewLogrusRecorder(s.logger, logrus.Fields{ 164 | "keyA0": "valueA0", 165 | "keyA1": "valueA1", 166 | }, logrus.Fields{ 167 | "keyB0": "valueB0", 168 | "keyB1": "valueB1", 169 | }) 170 | s.recordSingleEvent(event) 171 | s.Require().Len(s.hook.Entries, 1) 172 | s.Assert().Equal(logrus.WarnLevel, s.hook.LastEntry().Level) 173 | s.Assert().Equal("Fail", s.hook.LastEntry().Message) 174 | s.Assert().Equal(logrus.Fields{ 175 | "keyA0": "valueA0", 176 | "keyA1": "valueA1", 177 | "keyB0": "valueB0", 178 | "keyB1": "valueB1", 179 | "start": int64(420), 180 | "end": int64(4200), 181 | "elapsed": 3780, 182 | "response": "invalid response", 183 | "error": assert.AnError, 184 | }, s.hook.LastEntry().Data) 185 | } 186 | 187 | func TestLogrusRecorderTestSuite(t *testing.T) { 188 | suite.Run(t, new(LogrusRecorderTestSuite)) 189 | } 190 | -------------------------------------------------------------------------------- /recorders/progress.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders 10 | 11 | import ( 12 | "fmt" 13 | "os" 14 | 15 | "github.com/gosuri/uiprogress" 16 | "github.com/pinterest/bender" 17 | ) 18 | 19 | // NewLoadTestProgress returns progress and bar ready to be used in a recoreder. 20 | func NewLoadTestProgress(count int) (*uiprogress.Progress, *uiprogress.Bar) { 21 | progress := uiprogress.New() 22 | 23 | // We want to print progress on stderr so results can be easily redirected 24 | progress.SetOut(os.Stderr) 25 | 26 | // Create new progress bar displaying ELAPSED, CURRENT/MAX and COMPLETED 27 | bar := progress.AddBar(count) 28 | 29 | bar.PrependElapsed() 30 | bar.AppendFunc(func(b *uiprogress.Bar) string { 31 | return fmt.Sprintf("%d / %d", b.Current(), count) 32 | }) 33 | bar.AppendCompleted() 34 | 35 | return progress, bar 36 | } 37 | 38 | // NewProgressBarRecorder creates a new progress bar recorder. 39 | func NewProgressBarRecorder(bar *uiprogress.Bar) bender.Recorder { 40 | return func(msg interface{}) { 41 | //nolint:gocritic 42 | switch msg.(type) { 43 | case *bender.EndRequestEvent: 44 | bar.Incr() 45 | } 46 | } 47 | } 48 | -------------------------------------------------------------------------------- /recorders/progress_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders_test 10 | 11 | import ( 12 | "io/ioutil" 13 | "testing" 14 | 15 | "github.com/facebookincubator/fbender/recorders" 16 | "github.com/gosuri/uiprogress" 17 | "github.com/pinterest/bender" 18 | "github.com/stretchr/testify/suite" 19 | ) 20 | 21 | type ProgressBarRecorderTestSuite struct { 22 | suite.Suite 23 | progress *uiprogress.Progress 24 | bar *uiprogress.Bar 25 | recorder chan interface{} 26 | progressRecorder bender.Recorder 27 | } 28 | 29 | func (s *ProgressBarRecorderTestSuite) SetupTest() { 30 | s.progress = uiprogress.New() 31 | s.progress.SetOut(ioutil.Discard) 32 | s.progress.Start() 33 | s.bar = s.progress.AddBar(10) 34 | s.recorder = make(chan interface{}, 1) 35 | s.progressRecorder = recorders.NewProgressBarRecorder(s.bar) 36 | } 37 | 38 | func (s *ProgressBarRecorderTestSuite) TearDownTest() { 39 | s.progress.Stop() 40 | } 41 | 42 | func (s *ProgressBarRecorderTestSuite) recordSingleEvent(event interface{}) { 43 | s.recorder <- event 44 | close(s.recorder) 45 | bender.Record(s.recorder, s.progressRecorder) 46 | } 47 | 48 | func (s *ProgressBarRecorderTestSuite) TestStartEvent() { 49 | s.recordSingleEvent(new(bender.StartEvent)) 50 | s.Equal(0, s.bar.Current()) 51 | } 52 | 53 | func (s *ProgressBarRecorderTestSuite) TestEndEvent() { 54 | s.recordSingleEvent(new(bender.EndEvent)) 55 | s.Equal(0, s.bar.Current()) 56 | } 57 | 58 | func (s *ProgressBarRecorderTestSuite) TestWaitEvent() { 59 | s.recordSingleEvent(new(bender.WaitEvent)) 60 | s.Equal(0, s.bar.Current()) 61 | } 62 | 63 | func (s *ProgressBarRecorderTestSuite) TestStartRequestEvent() { 64 | s.recordSingleEvent(new(bender.StartRequestEvent)) 65 | s.Equal(0, s.bar.Current()) 66 | } 67 | 68 | func (s *ProgressBarRecorderTestSuite) TestEndRequestEvent() { 69 | s.recordSingleEvent(new(bender.EndRequestEvent)) 70 | s.Equal(1, s.bar.Current()) 71 | } 72 | 73 | func TestProgressBarRecorderTestSuite(t *testing.T) { 74 | suite.Run(t, new(ProgressBarRecorderTestSuite)) 75 | } 76 | -------------------------------------------------------------------------------- /recorders/statistics.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders 10 | 11 | import ( 12 | "sync/atomic" 13 | 14 | "github.com/pinterest/bender" 15 | ) 16 | 17 | // Statistics groups statistics gathered by statistics recoreder. 18 | type Statistics struct { 19 | Requests int64 20 | Errors int64 21 | } 22 | 23 | // Reset zeroes statistics. 24 | func (s *Statistics) Reset() { 25 | atomic.StoreInt64(&s.Requests, 0) 26 | atomic.StoreInt64(&s.Errors, 0) 27 | } 28 | 29 | // NewStatisticsRecorder creates new recorder which gathers statistics. 30 | func NewStatisticsRecorder(statistics *Statistics) bender.Recorder { 31 | return func(msg interface{}) { 32 | switch msg := msg.(type) { 33 | case *bender.StartEvent: 34 | statistics.Reset() 35 | case *bender.EndRequestEvent: 36 | atomic.AddInt64(&statistics.Requests, 1) 37 | 38 | if msg.Err != nil { 39 | atomic.AddInt64(&statistics.Errors, 1) 40 | } 41 | } 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /recorders/statistics_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package recorders_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/recorders" 15 | "github.com/pinterest/bender" 16 | "github.com/stretchr/testify/assert" 17 | "github.com/stretchr/testify/suite" 18 | ) 19 | 20 | type StatisticsRecorderTestSuite struct { 21 | suite.Suite 22 | statistics *recorders.Statistics 23 | recorder chan interface{} 24 | statisticsRecorder bender.Recorder 25 | } 26 | 27 | func (s *StatisticsRecorderTestSuite) SetupTest() { 28 | s.recorder = make(chan interface{}, 1) 29 | s.statistics = new(recorders.Statistics) 30 | s.statisticsRecorder = recorders.NewStatisticsRecorder(s.statistics) 31 | } 32 | 33 | func (s *StatisticsRecorderTestSuite) recordSingleEvent(event interface{}) { 34 | s.recorder <- event 35 | close(s.recorder) 36 | bender.Record(s.recorder, s.statisticsRecorder) 37 | } 38 | 39 | func (s *StatisticsRecorderTestSuite) TestStartEvent() { 40 | s.statistics.Requests = 42 41 | s.statistics.Errors = 6 42 | s.recordSingleEvent(new(bender.StartEvent)) 43 | s.Equal(int64(0), s.statistics.Requests) 44 | s.Equal(int64(0), s.statistics.Errors) 45 | } 46 | 47 | func (s *StatisticsRecorderTestSuite) TestEndEvent() { 48 | s.recordSingleEvent(new(bender.EndEvent)) 49 | s.Equal(int64(0), s.statistics.Requests) 50 | s.Equal(int64(0), s.statistics.Errors) 51 | } 52 | 53 | func (s *StatisticsRecorderTestSuite) TestWaitEvent() { 54 | s.recordSingleEvent(new(bender.WaitEvent)) 55 | s.Equal(int64(0), s.statistics.Requests) 56 | s.Equal(int64(0), s.statistics.Errors) 57 | } 58 | 59 | func (s *StatisticsRecorderTestSuite) TestStartRequestEvent() { 60 | s.recordSingleEvent(new(bender.StartRequestEvent)) 61 | s.Equal(int64(0), s.statistics.Requests) 62 | s.Equal(int64(0), s.statistics.Errors) 63 | } 64 | 65 | func (s *StatisticsRecorderTestSuite) TestEndRequestEvent_NoError() { 66 | s.recordSingleEvent(&bender.EndRequestEvent{Err: nil}) 67 | s.Equal(int64(1), s.statistics.Requests) 68 | s.Equal(int64(0), s.statistics.Errors) 69 | } 70 | 71 | func (s *StatisticsRecorderTestSuite) TestEndRequestEvent_Error() { 72 | s.recordSingleEvent(&bender.EndRequestEvent{Err: assert.AnError}) 73 | s.Equal(int64(1), s.statistics.Requests) 74 | s.Equal(int64(1), s.statistics.Errors) 75 | } 76 | 77 | func TestStatisticsRecorderTestSuite(t *testing.T) { 78 | suite.Run(t, new(StatisticsRecorderTestSuite)) 79 | } 80 | -------------------------------------------------------------------------------- /tester/comparator.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester 10 | 11 | import ( 12 | "errors" 13 | ) 14 | 15 | // Comparator allows to compare given values. 16 | type Comparator interface { 17 | Compare(x, y float64) bool 18 | Name() string 19 | } 20 | 21 | type comparator struct { 22 | repr string 23 | cmp func(x, y float64) bool 24 | } 25 | 26 | func (c *comparator) Compare(x, y float64) bool { 27 | return c.cmp(x, y) 28 | } 29 | 30 | func (c *comparator) Name() string { 31 | return c.repr 32 | } 33 | 34 | // Available comparators. 35 | //nolint:gochecknoglobals 36 | var ( 37 | LessThan Comparator = &comparator{repr: "<", cmp: func(x, y float64) bool { return x < y }} 38 | GreaterThan = &comparator{repr: ">", cmp: func(x, y float64) bool { return x > y }} 39 | ) 40 | 41 | // Comparators is a map of comparators representation to the actual comparator. 42 | //nolint:gochecknoglobals 43 | var Comparators = map[string]Comparator{ 44 | LessThan.Name(): LessThan, 45 | GreaterThan.Name(): GreaterThan, 46 | } 47 | 48 | // ErrInvalidComparator is returned when a comparator cannot be found. 49 | var ErrInvalidComparator = errors.New("invalid comparator") 50 | 51 | // ParseComparator returns a comparator based on the given string. 52 | func ParseComparator(name string) (Comparator, error) { 53 | if cmp, ok := Comparators[name]; ok { 54 | return cmp, nil 55 | } 56 | 57 | return nil, ErrInvalidComparator 58 | } 59 | -------------------------------------------------------------------------------- /tester/comparator_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester_test 10 | 11 | import ( 12 | "math/rand" 13 | "testing" 14 | "time" 15 | 16 | "github.com/facebookincubator/fbender/tester" 17 | "github.com/stretchr/testify/assert" 18 | ) 19 | 20 | func TestLessThan__Compare(t *testing.T) { 21 | s := rand.NewSource(time.Now().UnixNano()) 22 | r := rand.New(s) 23 | 24 | for i := 0; i < 4096; i++ { 25 | x, y := r.Float64(), r.Float64() 26 | assert.Equal(t, x < y, tester.LessThan.Compare(x, y), "%f < %f", x, y) 27 | } 28 | } 29 | 30 | func TestLessThan__Name(t *testing.T) { 31 | assert.Equal(t, "<", tester.LessThan.Name()) 32 | } 33 | 34 | func TestParseComparator_LessThan(t *testing.T) { 35 | cmp, err := tester.ParseComparator("<") 36 | assert.NoError(t, err) 37 | assertPointerEqual(t, tester.LessThan, cmp) 38 | } 39 | 40 | func TestGreaterThan__Compare(t *testing.T) { 41 | s := rand.NewSource(time.Now().UnixNano()) 42 | r := rand.New(s) 43 | 44 | for i := 0; i < 4096; i++ { 45 | x, y := r.Float64(), r.Float64() 46 | assert.Equal(t, x > y, tester.GreaterThan.Compare(x, y), "%f > %f", x, y) 47 | } 48 | } 49 | 50 | func TestGreaterThan__Name(t *testing.T) { 51 | assert.Equal(t, ">", tester.GreaterThan.Name()) 52 | } 53 | 54 | func TestParseComparator_GreaterThan(t *testing.T) { 55 | cmp, err := tester.ParseComparator(">") 56 | assert.NoError(t, err) 57 | assertPointerEqual(t, tester.GreaterThan, cmp) 58 | } 59 | 60 | func TestParseComparator(t *testing.T) { 61 | cmp, err := tester.ParseComparator("!") 62 | assert.Nil(t, cmp) 63 | assert.Equal(t, tester.ErrInvalidComparator, err) 64 | } 65 | -------------------------------------------------------------------------------- /tester/constraint.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "strconv" 15 | "time" 16 | 17 | "github.com/facebookincubator/fbender/utils" 18 | ) 19 | 20 | // Constraint represents a constraint tests should meet to be considered 21 | // successful. 22 | type Constraint struct { 23 | Metric Metric 24 | Aggregator Aggregator 25 | Comparator Comparator 26 | Threshold float64 27 | } 28 | 29 | func (c *Constraint) String() string { 30 | return fmt.Sprintf("%s(%s) %s %.2f", 31 | c.Aggregator.Name(), c.Metric.Name(), c.Comparator.Name(), c.Threshold) 32 | } 33 | 34 | // ErrNoDataPoints is raised when no data points are found. 35 | var ErrNoDataPoints = errors.New("no data points") 36 | 37 | // ErrNotSatisfied is raised when a condition is not met. 38 | var ErrNotSatisfied = errors.New("unsatisfied condition") 39 | 40 | // Check fetches metric and checks if the constraint has been satisfied. 41 | func (c *Constraint) Check(start time.Time, duration time.Duration) error { 42 | points, err := c.Metric.Fetch(start, duration) 43 | if err != nil { 44 | //nolint:wrapcheck 45 | return err 46 | } 47 | 48 | if points == nil { 49 | return ErrNoDataPoints 50 | } 51 | 52 | value := c.Aggregator.Aggregate(points) 53 | if !c.Comparator.Compare(value, c.Threshold) { 54 | return fmt.Errorf("%w: %.4f %s %.4f", ErrNotSatisfied, value, c.Comparator.Name(), c.Threshold) 55 | } 56 | 57 | return nil 58 | } 59 | 60 | // ConstraintsHelp is an help message on how to use constraints. 61 | const ConstraintsHelp = ` 62 | Constraints follow the syntax: 63 | Constraint ::= () 64 | Aggregator ::= "MIN" | "MAX" 65 | Metric ::= 66 | Cmp ::= "<" | ">" 67 | Threshold ::= 68 | 69 | Constraints examples: 70 | MIN(metric) < 20.5 71 | MAX(metric) > 0.45 72 | MIN(metric) < 123 73 | 74 | ` + GrowthHelp 75 | 76 | // ErrNotParsed should be returned when a parser did not parse a constraint. 77 | var ErrNotParsed = errors.New("constraint could not be parsed") 78 | 79 | // ErrInvalidFormat is raised when the constraint format is not correct. 80 | var ErrInvalidFormat = errors.New("invalid constraint format") 81 | 82 | // MetricParser is used to parse string values to a metric. 83 | // Parsers should return a metric and error if it successfully parsed 84 | // a metric string, or a fatal error occurred. Otherwise it should return 85 | // ErrNotParsed which will result in trying next parser from the list. 86 | type MetricParser func(string) (Metric, error) 87 | 88 | // Named capture groups of the constraints matching regexp. 89 | const ( 90 | aggregatorMatch = `(?P\w+)` 91 | metricMatch = `(?P\S+)` 92 | comparatorMatch = `(?P[<>=~!@#$%^&?]+)` 93 | thresholdMatch = `(?P[-+]?\d*\.?\d+)` 94 | ) 95 | 96 | //nolint:gochecknoglobals 97 | var constraintRegexp = utils.MustCompile( 98 | fmt.Sprintf( 99 | `^\s*%s\(%s\)\s*%s\s*%s\s*$`, 100 | aggregatorMatch, metricMatch, comparatorMatch, thresholdMatch, 101 | ), 102 | ) 103 | 104 | // ParseConstraint creates a constraint from a string representation. 105 | func ParseConstraint(s string, parsers ...MetricParser) (*Constraint, error) { 106 | if !constraintRegexp.MatchString(s) { 107 | return nil, ErrInvalidFormat 108 | } 109 | 110 | match := constraintRegexp.FindStringSubmatchMap(s) 111 | 112 | aggregator, err := ParseAggregator(match["aggregator"]) 113 | if err != nil { 114 | return nil, err 115 | } 116 | 117 | metric, err := parseMetric(match["metric"], parsers...) 118 | if err != nil { 119 | return nil, err 120 | } 121 | 122 | comparator, err := ParseComparator(match["comparator"]) 123 | if err != nil { 124 | return nil, err 125 | } 126 | 127 | threshold, err := strconv.ParseFloat(match["threshold"], 64) 128 | if err != nil { 129 | //nolint:wrapcheck 130 | return nil, err 131 | } 132 | 133 | return &Constraint{ 134 | Metric: metric, 135 | Aggregator: aggregator, 136 | Comparator: comparator, 137 | Threshold: threshold, 138 | }, nil 139 | } 140 | 141 | func parseMetric(name string, parsers ...MetricParser) (Metric, error) { 142 | for _, parser := range parsers { 143 | metric, err := parser(name) 144 | if err == nil { 145 | return metric, nil 146 | } else if !errors.Is(err, ErrNotParsed) { 147 | return nil, err 148 | } 149 | } 150 | 151 | return nil, ErrNotParsed 152 | } 153 | -------------------------------------------------------------------------------- /tester/constraint_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester_test 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "strings" 15 | "testing" 16 | "time" 17 | 18 | "github.com/facebookincubator/fbender/metric" 19 | "github.com/facebookincubator/fbender/tester" 20 | "github.com/stretchr/testify/mock" 21 | "github.com/stretchr/testify/suite" 22 | ) 23 | 24 | type MockedMetric struct { 25 | mock.Mock 26 | } 27 | 28 | func (m *MockedMetric) Setup(options interface{}) error { 29 | args := m.Called(options) 30 | 31 | return args.Error(0) 32 | } 33 | 34 | func (m *MockedMetric) Fetch(start time.Time, duration time.Duration) ([]tester.DataPoint, error) { 35 | args := m.Called(start, duration) 36 | 37 | return args.Get(0).([]tester.DataPoint), args.Error(1) 38 | } 39 | 40 | func (m *MockedMetric) Name() string { 41 | args := m.Called() 42 | 43 | return args.String(0) 44 | } 45 | 46 | type ParseConstraintTestSuite struct { 47 | suite.Suite 48 | 49 | metric *MockedMetric 50 | parsers []tester.MetricParser 51 | } 52 | 53 | func (s *ParseConstraintTestSuite) SetupTest() { 54 | s.metric = new(MockedMetric) 55 | s.parsers = []tester.MetricParser{ 56 | metric.Parser, 57 | } 58 | } 59 | 60 | //nolint:funlen 61 | func (s *ParseConstraintTestSuite) TestConstructor() { 62 | // Fork bomb is not a valid constraint. 63 | c, err := tester.ParseConstraint("💣(){ 💣|💣& };💣", s.parsers...) 64 | s.Assert().Nil(c) 65 | s.Assert().Error(err) 66 | // Invalid aggregator 'PWN'. 67 | c, err = tester.ParseConstraint("PWN(time) < 24", s.parsers...) 68 | s.Assert().Nil(c) 69 | s.Assert().Equal(err, tester.ErrInvalidAggregator) 70 | // Valid constraint - AVG. 71 | c, err = tester.ParseConstraint("AVG(errors) < 10", s.parsers...) 72 | s.Assert().NotNil(c) 73 | s.Assert().NoError(err) 74 | // Valid constraint - MIN. 75 | c, err = tester.ParseConstraint("MIN(errors) < 5", s.parsers...) 76 | s.Assert().NotNil(c) 77 | s.Assert().NoError(err) 78 | // Valid constraint - MAX. 79 | c, err = tester.ParseConstraint("MAX(errors) < 20", s.parsers...) 80 | s.Assert().NotNil(c) 81 | s.Assert().NoError(err) 82 | // Invalid metric '0xdeadbeef'. 83 | c, err = tester.ParseConstraint("MAX(0xdeadbeef) < 10", s.parsers...) 84 | s.Assert().Nil(c) 85 | s.Assert().Equal(err, tester.ErrNotParsed) 86 | // Should return error if metric parser returned error. 87 | c, err = tester.ParseConstraint("MAX(mymetrics) > 1", func(_ string) (tester.Metric, error) { 88 | //nolint:goerr113 89 | return nil, errors.New("error") 90 | }) 91 | s.Assert().Nil(c) 92 | s.Assert().Error(err) 93 | // Valid metric - errors. 94 | c, err = tester.ParseConstraint("MAX(errors) < 10", s.parsers...) 95 | s.Assert().NotNil(c) 96 | s.Assert().NoError(err) 97 | // Valid metric - latency. 98 | c, err = tester.ParseConstraint("MAX(latency) < 100", s.parsers...) 99 | s.Assert().NotNil(c) 100 | s.Assert().NoError(err) 101 | // Invalid comparator '@@'. 102 | c, err = tester.ParseConstraint("MAX(errors) @@ 10", s.parsers...) 103 | s.Assert().Nil(c) 104 | s.Assert().Equal(err, tester.ErrInvalidComparator) 105 | // Valid threshold. 106 | c, err = tester.ParseConstraint("MIN(latency) < 1.337", s.parsers...) 107 | s.Assert().NotNil(c) 108 | s.Assert().NoError(err) 109 | // Invalid threshold - comma. 110 | c, err = tester.ParseConstraint("MAX(latency) < 1,337", s.parsers...) 111 | s.Assert().Nil(c) 112 | s.Assert().Error(err) 113 | // Invalid threshold - hex. 114 | c, err = tester.ParseConstraint("MAX(latency) < 0x414141", s.parsers...) 115 | s.Assert().Nil(c) 116 | s.Assert().Error(err) 117 | // Invalid threshold - overflow. 118 | number := `179769313486231570814527423731704356798070567525844996598917476803157260780028538760589558632766878171540 119 | 458953514382464234321326889464182768467546703537516986049910576551282076245490090389328944075868508455133942304583236903 120 | 2229481658085593321233482747978262041447231687381771809192998812504040261841248583681337` 121 | c, err = tester.ParseConstraint( 122 | fmt.Sprintf( 123 | "MAX(latency) < %s", 124 | strings.ReplaceAll(number, "\n", ""), 125 | ), 126 | s.parsers..., 127 | ) 128 | s.Assert().Nil(c) 129 | s.Require().Error(err) 130 | s.Assert().Contains(err.Error(), "value out of range") 131 | } 132 | 133 | func (s *ParseConstraintTestSuite) TestCheck() { 134 | now := time.Now() 135 | // No datapoints should result in error. 136 | s.metric.On("Fetch", now, time.Second). 137 | Return([]tester.DataPoint(nil), nil).Once() 138 | 139 | c := &tester.Constraint{Metric: s.metric} 140 | err := c.Check(now, time.Second) 141 | 142 | s.Assert().Error(err) 143 | s.metric.AssertExpectations(s.T()) 144 | // If metric.Fetch returned error, Check should result in error. 145 | //nolint:goerr113 146 | s.metric.On("Fetch", now, time.Second). 147 | Return([]tester.DataPoint(nil), errors.New("error")).Once() 148 | 149 | c = &tester.Constraint{Metric: s.metric} 150 | err = c.Check(now, time.Second) 151 | 152 | s.Assert().Error(err) 153 | s.metric.AssertExpectations(s.T()) 154 | // Should pass - (10 < 10 + 1). 155 | s.metric.On("Fetch", now, time.Second). 156 | Return( 157 | []tester.DataPoint{{Value: 10}}, 158 | nil, 159 | ).Once() 160 | 161 | c = &tester.Constraint{ 162 | Metric: s.metric, 163 | Aggregator: tester.MinimumAggregator, 164 | Comparator: tester.LessThan, 165 | Threshold: 10 + 1, 166 | } 167 | err = c.Check(now, time.Second) 168 | 169 | s.Assert().NoError(err) 170 | s.metric.AssertExpectations(s.T()) 171 | // Should not pass - (10 > 10 + 1). 172 | s.metric.On("Fetch", now, time.Second). 173 | Return( 174 | []tester.DataPoint{ 175 | {Value: 10}, 176 | }, 177 | nil).Once() 178 | 179 | c = &tester.Constraint{ 180 | Metric: s.metric, 181 | Aggregator: tester.MinimumAggregator, 182 | Comparator: tester.GreaterThan, 183 | Threshold: 10 + 1, 184 | } 185 | err = c.Check(now, time.Second) 186 | 187 | s.Assert().Error(err) 188 | s.metric.AssertExpectations(s.T()) 189 | } 190 | 191 | func TestParseConstraintTestSuite(t *testing.T) { 192 | suite.Run(t, new(ParseConstraintTestSuite)) 193 | } 194 | -------------------------------------------------------------------------------- /tester/dhcpv4/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv4 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "net" 15 | "time" 16 | 17 | "github.com/insomniacslk/dhcp/dhcpv4" 18 | "github.com/insomniacslk/dhcp/dhcpv4/async" 19 | "github.com/pinterest/bender" 20 | protocol "github.com/pinterest/bender/dhcpv4" 21 | ) 22 | 23 | // Tester is a load tester for DHCPv4. 24 | type Tester struct { 25 | Target string 26 | Timeout time.Duration 27 | BufferSize int 28 | client *async.Client 29 | } 30 | 31 | // Before is called before the first test. 32 | func (t *Tester) Before(options interface{}) error { 33 | target, err := net.ResolveUDPAddr("udp4", t.Target) 34 | if err != nil { 35 | return fmt.Errorf("unable to set up the tester: %w", err) 36 | } 37 | 38 | addr, err := getLocalIPv4("eth0") 39 | if err != nil { 40 | return fmt.Errorf("unable to set up the tester: %w", err) 41 | } 42 | 43 | t.client = &async.Client{ 44 | ReadTimeout: t.Timeout, 45 | WriteTimeout: t.Timeout, 46 | RemoteAddr: target, 47 | LocalAddr: &net.UDPAddr{IP: addr, Port: async.DefaultServerPort}, 48 | IgnoreErrors: true, 49 | } 50 | 51 | return nil 52 | } 53 | 54 | // After is called after all tests are finished. 55 | func (t *Tester) After(_ interface{}) {} 56 | 57 | // BeforeEach is called before every test. 58 | func (t *Tester) BeforeEach(options interface{}) error { 59 | return t.client.Open(t.BufferSize) 60 | } 61 | 62 | // AfterEach is called after every test. 63 | func (t *Tester) AfterEach(_ interface{}) { 64 | t.client.Close() 65 | } 66 | 67 | func validator(req, res *dhcpv4.DHCPv4) error { 68 | return nil 69 | } 70 | 71 | // RequestExecutor returns a request executor. 72 | func (t *Tester) RequestExecutor(_ interface{}) (bender.RequestExecutor, error) { 73 | return protocol.CreateExecutor(t.client, validator) 74 | } 75 | 76 | // ErrNoAddress is raised when an interface has no ipv4 addresses assigned. 77 | var ErrNoAddress = errors.New("no ipv4 address found") 78 | 79 | // getLocalIPv4 returns the interface local IPv4 address. 80 | func getLocalIPv4(ifname string) (net.IP, error) { 81 | iface, err := net.InterfaceByName(ifname) 82 | if err != nil { 83 | //nolint:wrapcheck 84 | return nil, err 85 | } 86 | 87 | ifaddrs, err := iface.Addrs() 88 | if err != nil { 89 | //nolint:wrapcheck 90 | return nil, err 91 | } 92 | 93 | for _, ifaddr := range ifaddrs { 94 | if ipnet, ok := ifaddr.(*net.IPNet); ok { 95 | if ipnet.IP.To4() != nil && !ipnet.IP.IsLoopback() && !ipnet.IP.IsLinkLocalUnicast() { 96 | return ipnet.IP, nil 97 | } 98 | } 99 | } 100 | 101 | return nil, fmt.Errorf("%w, interface: %s", ErrNoAddress, ifname) 102 | } 103 | -------------------------------------------------------------------------------- /tester/dhcpv6/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dhcpv6 10 | 11 | import ( 12 | "fmt" 13 | "net" 14 | "time" 15 | 16 | "github.com/insomniacslk/dhcp/dhcpv6" 17 | "github.com/insomniacslk/dhcp/dhcpv6/async" 18 | "github.com/pinterest/bender" 19 | protocol "github.com/pinterest/bender/dhcpv6" 20 | ) 21 | 22 | // Tester is a load tester for DHCPv6. 23 | type Tester struct { 24 | Target string 25 | Timeout time.Duration 26 | BufferSize int 27 | client *async.Client 28 | } 29 | 30 | // Before is called before the first test. 31 | func (t *Tester) Before(options interface{}) error { 32 | target, err := net.ResolveUDPAddr("udp6", t.Target) 33 | if err != nil { 34 | return fmt.Errorf("unable to set up the tester: %w", err) 35 | } 36 | 37 | ip, err := dhcpv6.GetGlobalAddr("eth0") 38 | if err != nil { 39 | return fmt.Errorf("unable to set up the tester: %w", err) 40 | } 41 | 42 | t.client = &async.Client{ 43 | ReadTimeout: t.Timeout, 44 | WriteTimeout: t.Timeout, 45 | LocalAddr: &net.UDPAddr{IP: ip, Port: dhcpv6.DefaultServerPort, Zone: ""}, 46 | RemoteAddr: target, 47 | IgnoreErrors: true, 48 | } 49 | 50 | return nil 51 | } 52 | 53 | // After is called after all tests are finished. 54 | func (t *Tester) After(_ interface{}) {} 55 | 56 | // BeforeEach is called before every test. 57 | func (t *Tester) BeforeEach(options interface{}) error { 58 | return t.client.Open(t.BufferSize) 59 | } 60 | 61 | // AfterEach is called after every test. 62 | func (t *Tester) AfterEach(_ interface{}) { 63 | t.client.Close() 64 | } 65 | 66 | func validator(req, res dhcpv6.DHCPv6) error { 67 | return nil 68 | } 69 | 70 | // RequestExecutor returns a request executor. 71 | func (t *Tester) RequestExecutor(_ interface{}) (bender.RequestExecutor, error) { 72 | return protocol.CreateExecutor(t.client, validator), nil 73 | } 74 | -------------------------------------------------------------------------------- /tester/dns/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package dns 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "time" 15 | 16 | "github.com/miekg/dns" 17 | "github.com/pinterest/bender" 18 | protocol "github.com/pinterest/bender/dns" 19 | ) 20 | 21 | // ExtendedMsg wraps a dns.Msg with expectations. 22 | type ExtendedMsg struct { 23 | dns.Msg 24 | Rcode int 25 | } 26 | 27 | // Tester is a load tester for DNS. 28 | type Tester struct { 29 | Target string 30 | Timeout time.Duration 31 | Protocol string 32 | client *dns.Client 33 | } 34 | 35 | // ErrInvalidRequest is an error raised when the request is invalid. 36 | var ErrInvalidRequest = errors.New("invalid request") 37 | 38 | // ErrInvalidResponse is raised when the response is invalid. 39 | var ErrInvalidResponse = errors.New("invalid response") 40 | 41 | // Before is called before the first test. 42 | func (t *Tester) Before(options interface{}) error { 43 | //nolint:exhaustivestruct 44 | t.client = &dns.Client{ 45 | ReadTimeout: t.Timeout, 46 | DialTimeout: t.Timeout, 47 | WriteTimeout: t.Timeout, 48 | Net: t.Protocol, 49 | } 50 | 51 | return nil 52 | } 53 | 54 | // After is called after all tests are finished. 55 | func (t *Tester) After(_ interface{}) {} 56 | 57 | // BeforeEach is called before every test. 58 | func (t *Tester) BeforeEach(_ interface{}) error { 59 | return nil 60 | } 61 | 62 | // AfterEach is called after every test. 63 | func (t *Tester) AfterEach(_ interface{}) {} 64 | 65 | func validator(request, response *dns.Msg) error { 66 | if request.Id != response.Id { 67 | return fmt.Errorf("%w: %d, want: %d", ErrInvalidResponse, request.Id, response.Id) 68 | } 69 | 70 | return nil 71 | } 72 | 73 | // RequestExecutor returns a request executor. 74 | func (t *Tester) RequestExecutor(options interface{}) (bender.RequestExecutor, error) { 75 | innerExecutor := protocol.CreateExecutor(t.client, validator, t.Target) 76 | 77 | return func(n int64, request interface{}) (interface{}, error) { 78 | asExtended, ok := request.(*ExtendedMsg) 79 | if !ok { 80 | return nil, fmt.Errorf("%w: invalid type, want: *ExtendedMsg, got: %T", ErrInvalidRequest, request) 81 | } 82 | 83 | resp, err := innerExecutor(n, &asExtended.Msg) 84 | if err != nil { 85 | return resp, err 86 | } 87 | 88 | asMsg, ok := resp.(*dns.Msg) 89 | if !ok { 90 | return nil, fmt.Errorf("%w: invalid type, want: *dns.Msg, got: %T", ErrInvalidResponse, resp) 91 | } 92 | 93 | if asExtended.Rcode != -1 && asExtended.Rcode != asMsg.Rcode { 94 | return resp, fmt.Errorf( 95 | "%w: invalid rcode want: %q, got: %q", ErrInvalidResponse, 96 | dns.RcodeToString[asExtended.Rcode], dns.RcodeToString[asMsg.Rcode]) 97 | } 98 | 99 | return resp, nil 100 | }, nil 101 | } 102 | -------------------------------------------------------------------------------- /tester/growth.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "strconv" 15 | "strings" 16 | ) 17 | 18 | // Growth is used to determine what test should be ran next. 19 | type Growth interface { 20 | OnSuccess(test int) int 21 | OnFail(test int) int 22 | String() string 23 | } 24 | 25 | // LinearGrowth increases test by a specified amount with every successful test. 26 | type LinearGrowth struct { 27 | Increase int 28 | } 29 | 30 | // LinearGrowthPrefix prefix used in linear growth string representation. 31 | const LinearGrowthPrefix = "+" 32 | 33 | func (g *LinearGrowth) String() string { 34 | return fmt.Sprintf("%s%d", LinearGrowthPrefix, g.Increase) 35 | } 36 | 37 | // OnSuccess increases test by a specified amount. 38 | func (g *LinearGrowth) OnSuccess(test int) int { 39 | return test + g.Increase 40 | } 41 | 42 | // OnFail stops the tests. 43 | func (g *LinearGrowth) OnFail(test int) int { 44 | return 0 45 | } 46 | 47 | // PercentageGrowth increases test by a specified percentage with every successful test. 48 | type PercentageGrowth struct { 49 | Increase float64 50 | } 51 | 52 | // PercentageGrowthPrefix prefix used in percentage growth string representation. 53 | const PercentageGrowthPrefix = "%" 54 | 55 | func (g *PercentageGrowth) String() string { 56 | return fmt.Sprintf("%s%.2f", PercentageGrowthPrefix, g.Increase) 57 | } 58 | 59 | // OnSuccess increases test by a specified percentage. 60 | func (g *PercentageGrowth) OnSuccess(test int) int { 61 | return int((100. + g.Increase) / 100. * float64(test)) 62 | } 63 | 64 | // OnFail stops the tests. 65 | func (g *PercentageGrowth) OnFail(test int) int { 66 | return 0 67 | } 68 | 69 | // ExponentialGrowth performs binary search up to a given precision. 70 | type ExponentialGrowth struct { 71 | Precision int 72 | 73 | left, right int 74 | bound bool 75 | } 76 | 77 | // ExponentialGrowthPrefix prefix used in exponential growth string representation. 78 | const ExponentialGrowthPrefix = "^" 79 | 80 | func (g *ExponentialGrowth) String() string { 81 | return fmt.Sprintf("%s%d", ExponentialGrowthPrefix, g.Precision) 82 | } 83 | 84 | // OnSuccess sets the lower bound to the last test and returns (left+right) / 2 85 | // unless the precision has been achieved. 86 | func (g *ExponentialGrowth) OnSuccess(test int) int { 87 | g.left = test 88 | if !g.bound { 89 | return test * 2 90 | } 91 | 92 | if g.right-g.left <= g.Precision { 93 | return 0 94 | } 95 | 96 | return int(float64(g.right+g.left) / 2) 97 | } 98 | 99 | // OnFail sets the upper bound to the last test and returns (left+right) / 2 100 | // unless the precision has been achieved. 101 | func (g *ExponentialGrowth) OnFail(test int) int { 102 | g.right = test 103 | g.bound = true 104 | 105 | if g.right-g.left <= g.Precision { 106 | return 0 107 | } 108 | 109 | return int(float64(g.right+g.left) / 2) 110 | } 111 | 112 | // GrowthHelp provides usage help about the growth. 113 | const GrowthHelp = `Growth determines what will be the next value used for a test. 114 | * linear growth (+int) increases test value by a fixed amount after each success, 115 | stops immediately after the first failure 116 | * percentage growth (%float) increases test value by a fixed percentage after 117 | each success, stops immediately after the first failure 118 | * exponential growth (^int) first doubles the test value after each success to 119 | find an upper bound, then performs a binary search up to a given precision` 120 | 121 | // ErrInvalidGrowth is returned when a growth cannot be found. 122 | var ErrInvalidGrowth = errors.New("unknown growth, want +int, %%flaot, ^int") 123 | 124 | // ParseGrowth creates a growth from its string representation. 125 | func ParseGrowth(value string) (Growth, error) { 126 | switch { 127 | case strings.HasPrefix(value, LinearGrowthPrefix): 128 | inc, err := strconv.Atoi(strings.TrimPrefix(value, LinearGrowthPrefix)) 129 | if err != nil { 130 | //nolint:wrapcheck 131 | return nil, err 132 | } 133 | 134 | return &LinearGrowth{Increase: inc}, nil 135 | 136 | case strings.HasPrefix(value, PercentageGrowthPrefix): 137 | inc, err := strconv.ParseFloat(strings.TrimPrefix(value, PercentageGrowthPrefix), 64) 138 | if err != nil { 139 | //nolint:wrapcheck 140 | return nil, err 141 | } 142 | 143 | return &PercentageGrowth{Increase: inc}, nil 144 | 145 | case strings.HasPrefix(value, ExponentialGrowthPrefix): 146 | prec, err := strconv.Atoi(strings.TrimPrefix(value, ExponentialGrowthPrefix)) 147 | if err != nil { 148 | //nolint:wrapcheck 149 | return nil, err 150 | } 151 | 152 | return &ExponentialGrowth{Precision: prec}, nil 153 | 154 | default: 155 | return nil, ErrInvalidGrowth 156 | } 157 | } 158 | -------------------------------------------------------------------------------- /tester/growth_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester_test 10 | 11 | import ( 12 | "math" 13 | "testing" 14 | 15 | "github.com/facebookincubator/fbender/tester" 16 | "github.com/stretchr/testify/assert" 17 | "github.com/stretchr/testify/require" 18 | ) 19 | 20 | func TestLinearGrowth__OnSuccess(t *testing.T) { 21 | // test checks if the linear growth follows the arithmetic sequence starting 22 | // at s, growing by r for n steps 23 | test := func(s, r, n int) { 24 | g := &tester.LinearGrowth{Increase: r} 25 | a := s 26 | 27 | for i := 1; i < n; i++ { 28 | a = g.OnSuccess(a) 29 | assert.Equal(t, s+r*i, a) 30 | } 31 | } 32 | 33 | test(10, 3, 100) 34 | test(42, 42, 42) 35 | test(100, 1, 100) 36 | } 37 | 38 | func TestLinearGrowth__OnFail(t *testing.T) { 39 | test := func(s, r, n int) { 40 | g := &tester.LinearGrowth{Increase: r} 41 | a := s 42 | 43 | for i := 1; i < n; i++ { 44 | a = g.OnSuccess(a) 45 | } 46 | 47 | a = g.OnFail(a) 48 | assert.Equal(t, 0, a) 49 | } 50 | 51 | test(10, 3, 100) 52 | test(42, 42, 42) 53 | test(100, 1, 100) 54 | } 55 | 56 | func TestLinearGrowth__String(t *testing.T) { 57 | g := &tester.LinearGrowth{Increase: 3} 58 | assert.Equal(t, g.String(), "+3") 59 | 60 | g = &tester.LinearGrowth{Increase: 42} 61 | assert.Equal(t, g.String(), "+42") 62 | 63 | g = &tester.LinearGrowth{Increase: 100} 64 | assert.Equal(t, g.String(), "+100") 65 | } 66 | 67 | func TestParseGrowth_LinearGrowth(t *testing.T) { 68 | // Valid linear growth 69 | g, err := tester.ParseGrowth("+200") 70 | require.NoError(t, err) 71 | assert.IsType(t, new(tester.LinearGrowth), g) 72 | assert.Equal(t, 200, g.(*tester.LinearGrowth).Increase) 73 | // Invalid value 74 | _, err = tester.ParseGrowth("+abcdef") 75 | assert.EqualError(t, err, "strconv.Atoi: parsing \"abcdef\": invalid syntax") 76 | 77 | _, err = tester.ParseGrowth("+99.9") 78 | assert.EqualError(t, err, "strconv.Atoi: parsing \"99.9\": invalid syntax") 79 | } 80 | 81 | func TestPercentageGrowth__OnSuccess(t *testing.T) { 82 | // test checks if the linear growth follows the arithmetic sequence starting 83 | // at s, growing by r for n steps 84 | test := func(s int, r float64, n int) { 85 | g := &tester.PercentageGrowth{Increase: r} 86 | a := s 87 | 88 | for i := 1; i < n; i++ { 89 | a = g.OnSuccess(a) 90 | expected := int(float64(s) * math.Pow((100+r)/100., float64(i))) 91 | // We're reounding every result to int so it may not be equal 92 | assert.InDelta(t, expected, a, float64(s)*r/100.) 93 | } 94 | } 95 | 96 | test(10, 100, 100) 97 | test(2, 200, 10) 98 | test(100, 100, 10) 99 | } 100 | 101 | func TestPercentageGrowth__OnFail(t *testing.T) { 102 | test := func(s int, r float64, n int) { 103 | g := &tester.PercentageGrowth{Increase: r} 104 | a := s 105 | 106 | for i := 1; i < n; i++ { 107 | a = g.OnSuccess(a) 108 | } 109 | 110 | a = g.OnFail(a) 111 | assert.Equal(t, 0, a) 112 | } 113 | 114 | test(10, 3, 100) 115 | test(42, 42, 42) 116 | test(100, 1, 100) 117 | } 118 | 119 | func TestPercentageGrowth__String(t *testing.T) { 120 | g := &tester.PercentageGrowth{Increase: 102.5} 121 | assert.Equal(t, g.String(), "%102.50") 122 | 123 | g = &tester.PercentageGrowth{Increase: 66.66} 124 | assert.Equal(t, g.String(), "%66.66") 125 | 126 | g = &tester.PercentageGrowth{Increase: 100} 127 | assert.Equal(t, g.String(), "%100.00") 128 | } 129 | 130 | func TestParseGrowth_PercentageGrowth(t *testing.T) { 131 | // Valid linear growth 132 | g, err := tester.ParseGrowth("%100.50") 133 | require.NoError(t, err) 134 | 135 | assert.IsType(t, new(tester.PercentageGrowth), g) 136 | assert.Equal(t, 100.50, g.(*tester.PercentageGrowth).Increase) 137 | 138 | // Invalid value 139 | _, err = tester.ParseGrowth("%abcdef") 140 | assert.EqualError(t, err, "strconv.ParseFloat: parsing \"abcdef\": invalid syntax") 141 | } 142 | 143 | func TestExponentialGrowth__OnSuccess(t *testing.T) { 144 | g := &tester.ExponentialGrowth{Precision: 10} 145 | 146 | // When not bound it should double the test value 147 | assert.Equal(t, 40, g.OnSuccess(20)) 148 | assert.Equal(t, 80, g.OnSuccess(40)) 149 | assert.Equal(t, 160, g.OnSuccess(80)) 150 | // When bound it should return a middle value 151 | assert.Equal(t, 120, g.OnFail(160)) 152 | assert.Equal(t, 140, g.OnSuccess(120)) 153 | assert.Equal(t, 150, g.OnSuccess(140)) 154 | // Finally return 0 when precision is met 155 | assert.Equal(t, 0, g.OnSuccess(150)) 156 | } 157 | 158 | func TestExponentialGrowth__OnFail(t *testing.T) { 159 | g := &tester.ExponentialGrowth{Precision: 10} 160 | 161 | assert.Equal(t, 40, g.OnSuccess(20)) 162 | assert.Equal(t, 80, g.OnSuccess(40)) 163 | assert.Equal(t, 160, g.OnSuccess(80)) 164 | // It should set the upper bound and return middle value 165 | assert.Equal(t, 120, g.OnFail(160)) 166 | assert.Equal(t, 100, g.OnFail(120)) 167 | assert.Equal(t, 90, g.OnFail(100)) 168 | // Finally return 0 when precision is met 169 | assert.Equal(t, 0, g.OnSuccess(90)) 170 | } 171 | 172 | func TestExponentialGrowth__String(t *testing.T) { 173 | g := &tester.ExponentialGrowth{Precision: 3} 174 | assert.Equal(t, g.String(), "^3") 175 | 176 | g = &tester.ExponentialGrowth{Precision: 42} 177 | assert.Equal(t, g.String(), "^42") 178 | 179 | g = &tester.ExponentialGrowth{Precision: 100} 180 | assert.Equal(t, g.String(), "^100") 181 | } 182 | 183 | func TestParseGrowth_ExponentialGrowth(t *testing.T) { 184 | // Valid linear growth 185 | g, err := tester.ParseGrowth("^10") 186 | require.NoError(t, err) 187 | 188 | assert.IsType(t, new(tester.ExponentialGrowth), g) 189 | assert.Equal(t, 10, g.(*tester.ExponentialGrowth).Precision) 190 | 191 | // Invalid value 192 | _, err = tester.ParseGrowth("^abcdef") 193 | assert.EqualError(t, err, "strconv.Atoi: parsing \"abcdef\": invalid syntax") 194 | 195 | _, err = tester.ParseGrowth("^99.9") 196 | assert.EqualError(t, err, "strconv.Atoi: parsing \"99.9\": invalid syntax") 197 | } 198 | 199 | func TestParseGrowth(t *testing.T) { 200 | g, err := tester.ParseGrowth("@200") 201 | assert.Nil(t, g) 202 | assert.Equal(t, tester.ErrInvalidGrowth, err) 203 | } 204 | -------------------------------------------------------------------------------- /tester/http/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package http 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "io" 15 | "io/ioutil" 16 | "net/http" 17 | "time" 18 | 19 | "github.com/pinterest/bender" 20 | protocol "github.com/pinterest/bender/http" 21 | ) 22 | 23 | // Tester is a load tester for HTTP. 24 | type Tester struct { 25 | Timeout time.Duration 26 | Validator protocol.ResponseValidator 27 | client *http.Client 28 | } 29 | 30 | // httpStatusOK is the HTTP correct response status code. 31 | const httpStatusOK = 200 32 | 33 | // ErrInvalidResponse is raised when a request returns a code different than 200. 34 | var ErrInvalidResponse = errors.New("invalid response status") 35 | 36 | // Before is called before the first test. 37 | func (t *Tester) Before(options interface{}) error { 38 | t.client = &http.Client{ 39 | Timeout: t.Timeout, 40 | } 41 | 42 | return nil 43 | } 44 | 45 | // After is called after all tests are finished. 46 | func (t *Tester) After(_ interface{}) {} 47 | 48 | // BeforeEach is called before every test. 49 | func (t *Tester) BeforeEach(_ interface{}) error { 50 | return nil 51 | } 52 | 53 | // AfterEach is called after every test. 54 | func (t *Tester) AfterEach(_ interface{}) {} 55 | 56 | // A default validator checks if response type is 200 OK, reads the whole body 57 | // to force download. 58 | func validator(request interface{}, response *http.Response) error { 59 | if response.StatusCode != httpStatusOK { 60 | return fmt.Errorf("%w, want: \"200 OK\", got: \"%s\"", ErrInvalidResponse, response.Status) 61 | } 62 | 63 | _, err := io.Copy(ioutil.Discard, response.Body) 64 | 65 | //nolint:wrapcheck 66 | return err 67 | } 68 | 69 | // RequestExecutor returns a request executor. 70 | func (t *Tester) RequestExecutor(options interface{}) (bender.RequestExecutor, error) { 71 | if t.Validator == nil { 72 | return protocol.CreateExecutor(nil, t.client, validator), nil 73 | } 74 | 75 | return protocol.CreateExecutor(nil, t.client, t.Validator), nil 76 | } 77 | -------------------------------------------------------------------------------- /tester/metric.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester 10 | 11 | import ( 12 | "errors" 13 | "time" 14 | ) 15 | 16 | // DataPoint represents a sample of data. 17 | type DataPoint struct { 18 | Time time.Time 19 | Value float64 20 | } 21 | 22 | // Metric provides a function to get data points of this metric. 23 | type Metric interface { 24 | // Setup is used to setup the metric before the test 25 | Setup(options interface{}) error 26 | // Fetch is used to get metric measures. 27 | Fetch(start time.Time, duration time.Duration) ([]DataPoint, error) 28 | // Name is used to serialize metrics. 29 | Name() string 30 | } 31 | 32 | // Aggregator provides a function to aggregate data points. 33 | type Aggregator interface { 34 | // Aggregate calculates a single value for a metric datapoints. 35 | Aggregate([]DataPoint) float64 36 | // Name is used to serialize metrics. 37 | Name() string 38 | } 39 | 40 | type metricAggregator struct { 41 | repr string 42 | aggr func([]DataPoint) float64 43 | } 44 | 45 | func (m *metricAggregator) Aggregate(points []DataPoint) float64 { 46 | return m.aggr(points) 47 | } 48 | 49 | func (m *metricAggregator) Name() string { 50 | return m.repr 51 | } 52 | 53 | // MinimumAggregator returns the smallest datapoint value. 54 | //nolint:gochecknoglobals 55 | var MinimumAggregator Aggregator = &metricAggregator{ 56 | repr: "MIN", 57 | aggr: func(points []DataPoint) float64 { 58 | if len(points) == 0 { 59 | return 0. 60 | } 61 | 62 | x := points[0].Value 63 | for _, point := range points { 64 | if point.Value < x { 65 | x = point.Value 66 | } 67 | } 68 | 69 | return x 70 | }, 71 | } 72 | 73 | // MaximumAggregator returns the smallest datapoint value. 74 | //nolint:gochecknoglobals 75 | var MaximumAggregator Aggregator = &metricAggregator{ 76 | repr: "MAX", 77 | aggr: func(points []DataPoint) float64 { 78 | if len(points) == 0 { 79 | return 0. 80 | } 81 | 82 | x := points[0].Value 83 | for _, point := range points { 84 | if point.Value > x { 85 | x = point.Value 86 | } 87 | } 88 | 89 | return x 90 | }, 91 | } 92 | 93 | // AverageAggregator returns the average data point value. 94 | //nolint:gochecknoglobals 95 | var AverageAggregator Aggregator = &metricAggregator{ 96 | repr: "AVG", 97 | aggr: func(points []DataPoint) float64 { 98 | if len(points) == 0 { 99 | return 0. 100 | } 101 | 102 | sum := 0. 103 | for _, point := range points { 104 | sum = sum + point.Value 105 | } 106 | 107 | return sum / float64(len(points)) 108 | }, 109 | } 110 | 111 | // Aggregators is a map of aggregators representation to the actual aggregator. 112 | //nolint:gochecknoglobals 113 | var Aggregators = map[string]Aggregator{ 114 | MinimumAggregator.Name(): MinimumAggregator, 115 | MaximumAggregator.Name(): MaximumAggregator, 116 | AverageAggregator.Name(): AverageAggregator, 117 | } 118 | 119 | // ErrInvalidAggregator is returned when a metric aggregator cannot be found. 120 | var ErrInvalidAggregator = errors.New("invalid aggregator") 121 | 122 | // ParseAggregator returns metric aggregator from its name. 123 | func ParseAggregator(name string) (Aggregator, error) { 124 | if aggregator, ok := Aggregators[name]; ok { 125 | return aggregator, nil 126 | } 127 | 128 | return nil, ErrInvalidAggregator 129 | } 130 | -------------------------------------------------------------------------------- /tester/metric_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester_test 10 | 11 | import ( 12 | "math/rand" 13 | "testing" 14 | "time" 15 | 16 | "github.com/facebookincubator/fbender/tester" 17 | "github.com/stretchr/testify/assert" 18 | "github.com/stretchr/testify/suite" 19 | ) 20 | 21 | // randomValues generates a slice filled with random values datapoints. 22 | func randomDataPoints(n int) []tester.DataPoint { 23 | s := rand.NewSource(time.Now().UnixNano()) 24 | r := rand.New(s) 25 | 26 | values := make([]tester.DataPoint, n) 27 | for i := 0; i < 10; i++ { 28 | values[i] = tester.DataPoint{Value: r.Float64()} 29 | } 30 | 31 | return values 32 | } 33 | 34 | type MinimumAggregatorTestSuite struct { 35 | suite.Suite 36 | aggregator tester.Aggregator 37 | } 38 | 39 | func (s *MinimumAggregatorTestSuite) SetupTest() { 40 | s.aggregator = tester.MinimumAggregator 41 | } 42 | 43 | func (s *MinimumAggregatorTestSuite) testAggregateRandomSlice(n int) { 44 | points := randomDataPoints(n) 45 | v := s.aggregator.Aggregate(points) 46 | min := points[0].Value 47 | 48 | for _, point := range points { 49 | s.Assert().True(v <= point.Value, "expected %f to be less or equal to %f", v, point.Value) 50 | 51 | if point.Value < min { 52 | min = point.Value 53 | } 54 | } 55 | 56 | s.Assert().Equal(min, v) 57 | } 58 | 59 | func (s *MinimumAggregatorTestSuite) TestName() { 60 | s.Assert().Equal("MIN", s.aggregator.Name()) 61 | } 62 | 63 | func (s *MinimumAggregatorTestSuite) TestAggregate() { 64 | // For nil and empty list we should get 0. 65 | v := s.aggregator.Aggregate(nil) 66 | s.Assert().Equal(0., v) 67 | 68 | v = s.aggregator.Aggregate([]tester.DataPoint{}) 69 | s.Assert().Equal(0., v) 70 | 71 | // Single value is the minimum 72 | v = s.aggregator.Aggregate([]tester.DataPoint{ 73 | {Value: 42.}, 74 | }) 75 | s.Assert().Equal(42., v) 76 | 77 | // Random tests to make sure we get minimum 78 | for i := 0; i < 128; i++ { 79 | s.testAggregateRandomSlice(4096) 80 | } 81 | } 82 | 83 | type MaximumAggregatorTestSuite struct { 84 | suite.Suite 85 | aggregator tester.Aggregator 86 | } 87 | 88 | func (s *MaximumAggregatorTestSuite) SetupTest() { 89 | s.aggregator = tester.MaximumAggregator 90 | } 91 | 92 | func (s *MaximumAggregatorTestSuite) testAggregateRandomSlice(n int) { 93 | points := randomDataPoints(n) 94 | v := s.aggregator.Aggregate(points) 95 | max := points[0].Value 96 | 97 | for _, point := range points { 98 | s.Assert().True(v >= point.Value, "expected %f to be greater or equal to %f", v, point.Value) 99 | 100 | if point.Value > max { 101 | max = point.Value 102 | } 103 | } 104 | 105 | s.Assert().Equal(max, v) 106 | } 107 | 108 | func (s *MaximumAggregatorTestSuite) TestName() { 109 | s.Assert().Equal("MAX", s.aggregator.Name()) 110 | } 111 | 112 | func (s *MaximumAggregatorTestSuite) TestAggregate() { 113 | // For nil and empty list we should get 0. 114 | v := s.aggregator.Aggregate(nil) 115 | s.Assert().Equal(0., v) 116 | 117 | v = s.aggregator.Aggregate([]tester.DataPoint{}) 118 | s.Assert().Equal(0., v) 119 | 120 | // Single value is the maximum 121 | v = s.aggregator.Aggregate([]tester.DataPoint{ 122 | {Value: 42.}, 123 | }) 124 | s.Assert().Equal(42., v) 125 | 126 | // Random tests to make sure we get maximum 127 | for i := 0; i < 128; i++ { 128 | s.testAggregateRandomSlice(4096) 129 | } 130 | } 131 | 132 | type AverageAggregatorTestSuite struct { 133 | suite.Suite 134 | aggregator tester.Aggregator 135 | } 136 | 137 | func (s *AverageAggregatorTestSuite) SetupTest() { 138 | s.aggregator = tester.AverageAggregator 139 | } 140 | 141 | func (s *AverageAggregatorTestSuite) testAggregateRandomSlice(n int) { 142 | points := randomDataPoints(n) 143 | v := s.aggregator.Aggregate(points) 144 | sum := 0. 145 | 146 | for _, point := range points { 147 | sum += point.Value 148 | } 149 | 150 | s.Assert().Equal(sum/float64(n), v) 151 | } 152 | 153 | func (s *AverageAggregatorTestSuite) TestName() { 154 | s.Assert().Equal("AVG", s.aggregator.Name()) 155 | } 156 | 157 | func (s *AverageAggregatorTestSuite) TestAggregate() { 158 | // For nil and empty list we should get 0. 159 | v := s.aggregator.Aggregate(nil) 160 | s.Assert().Equal(0., v) 161 | 162 | v = s.aggregator.Aggregate([]tester.DataPoint{}) 163 | s.Assert().Equal(0., v) 164 | 165 | // Single value is the average 166 | v = s.aggregator.Aggregate([]tester.DataPoint{ 167 | {Value: 42.}, 168 | }) 169 | s.Assert().Equal(42., v) 170 | 171 | // Few manual tests 172 | points := []tester.DataPoint{ 173 | {Value: 10.}, 174 | {Value: 20.}, 175 | {Value: 30.}, 176 | } 177 | 178 | v = s.aggregator.Aggregate(points) 179 | s.Assert().Equal(20., v) 180 | 181 | points = []tester.DataPoint{ 182 | {Value: 1.}, 183 | {Value: 2.}, 184 | {Value: 3.}, 185 | {Value: 4.}, 186 | {Value: 5.}, 187 | {Value: 6.}, 188 | {Value: 7.}, 189 | {Value: 8.}, 190 | {Value: 9.}, 191 | } 192 | 193 | v = s.aggregator.Aggregate(points) 194 | s.Assert().Equal(5., v) 195 | 196 | // Random tests to make sure we get maximum 197 | for i := 0; i < 128; i++ { 198 | s.testAggregateRandomSlice(4096) 199 | } 200 | } 201 | 202 | func TestParseAggregator(t *testing.T) { 203 | a, err := tester.ParseAggregator("MIN") 204 | assert.NoError(t, err) 205 | assertPointerEqual(t, tester.MinimumAggregator, a, "Expected minimum aggregator") 206 | 207 | a, err = tester.ParseAggregator("MAX") 208 | assert.NoError(t, err) 209 | assertPointerEqual(t, tester.MaximumAggregator, a, "Expected maximum aggregator") 210 | 211 | a, err = tester.ParseAggregator("AVG") 212 | assert.NoError(t, err) 213 | assertPointerEqual(t, tester.AverageAggregator, a, "Expected average aggregator") 214 | 215 | a, err = tester.ParseAggregator("Nonexistent") 216 | assert.Error(t, err) 217 | assert.Nil(t, a) 218 | assert.Equal(t, tester.ErrInvalidAggregator, err) 219 | } 220 | 221 | func TestMinimumAggregatorTestSuite(t *testing.T) { 222 | suite.Run(t, new(MinimumAggregatorTestSuite)) 223 | } 224 | 225 | func TestMaximumAggregatorTestSuite(t *testing.T) { 226 | suite.Run(t, new(MaximumAggregatorTestSuite)) 227 | } 228 | 229 | func TestAverageAggregatorTestSuite(t *testing.T) { 230 | suite.Run(t, new(AverageAggregatorTestSuite)) 231 | } 232 | -------------------------------------------------------------------------------- /tester/run/common.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package run 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/log" 15 | "github.com/facebookincubator/fbender/tester" 16 | ) 17 | 18 | // checkConstraints loops through given constraints and returns whether all of 19 | // them have been met. 20 | func checkConstraints(start time.Time, duration time.Duration, constraints ...*tester.Constraint) bool { 21 | for _, constraint := range constraints { 22 | if err := constraint.Check(start, duration); err != nil { 23 | log.Errorf("Error checking %q: %v\n", constraint.String(), err) 24 | 25 | return false 26 | } 27 | } 28 | 29 | return true 30 | } 31 | -------------------------------------------------------------------------------- /tester/run/common_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package run_test 10 | 11 | import ( 12 | "errors" 13 | "fmt" 14 | "testing" 15 | "time" 16 | 17 | "github.com/facebookincubator/fbender/tester" 18 | "github.com/pinterest/bender" 19 | "github.com/stretchr/testify/mock" 20 | ) 21 | 22 | type MockedTester struct { 23 | mock.Mock 24 | } 25 | 26 | func (m *MockedTester) Before(options interface{}) error { 27 | args := m.Called(options) 28 | 29 | return args.Error(0) 30 | } 31 | 32 | func (m *MockedTester) After(options interface{}) { 33 | m.Called(options) 34 | } 35 | 36 | func (m *MockedTester) BeforeEach(options interface{}) error { 37 | args := m.Called(options) 38 | 39 | return args.Error(0) 40 | } 41 | 42 | func (m *MockedTester) AfterEach(options interface{}) { 43 | m.Called(options) 44 | } 45 | 46 | func (m *MockedTester) RequestExecutor(options interface{}) (bender.RequestExecutor, error) { 47 | args := m.Called(options) 48 | 49 | return m.DummyExecutor, args.Error(0) 50 | } 51 | 52 | func (m *MockedTester) DummyExecutor(timestamp int64, request interface{}) (interface{}, error) { 53 | args := m.Called(timestamp, request) 54 | 55 | return args.Get(0), args.Error(1) 56 | } 57 | 58 | var ErrDummy = errors.New("dummy error") 59 | 60 | type MockedGrowth struct { 61 | mock.Mock 62 | } 63 | 64 | func (m *MockedGrowth) String() string { 65 | args := m.Called() 66 | 67 | return args.String(0) 68 | } 69 | 70 | func (m *MockedGrowth) OnSuccess(test int) int { 71 | args := m.Called(test) 72 | 73 | return args.Int(0) 74 | } 75 | 76 | func (m *MockedGrowth) OnFail(test int) int { 77 | args := m.Called(test) 78 | 79 | return args.Int(0) 80 | } 81 | 82 | type MockedMetric struct { 83 | mock.Mock 84 | } 85 | 86 | func (m *MockedMetric) Setup(options interface{}) error { 87 | args := m.Called(options) 88 | 89 | return args.Error(0) 90 | } 91 | 92 | func (m *MockedMetric) Fetch(start time.Time, duration time.Duration) ([]tester.DataPoint, error) { 93 | args := m.Called(start, duration) 94 | if points, ok := args.Get(0).([]tester.DataPoint); ok { 95 | return points, args.Error(1) 96 | } 97 | 98 | panic(fmt.Sprintf("assert: arguments: DataPointSlice(0) failed because object wasn't correct type: %v", args.Get(0))) 99 | } 100 | 101 | func (m *MockedMetric) Name() string { 102 | args := m.Called() 103 | 104 | return args.String(0) 105 | } 106 | 107 | type MockedAggregator struct { 108 | mock.Mock 109 | } 110 | 111 | func (m *MockedAggregator) Aggregate(points []tester.DataPoint) float64 { 112 | args := m.Called(points) 113 | 114 | return args.Get(0).(float64) 115 | } 116 | 117 | func (m *MockedAggregator) Name() string { 118 | args := m.Called() 119 | 120 | return args.String(0) 121 | } 122 | 123 | type MockedComparator struct { 124 | mock.Mock 125 | } 126 | 127 | func (m *MockedComparator) Compare(x, y float64) bool { 128 | args := m.Called(x, y) 129 | 130 | return args.Bool(0) 131 | } 132 | 133 | func (m *MockedComparator) Name() string { 134 | args := m.Called() 135 | 136 | return args.String(0) 137 | } 138 | 139 | type MockedConstraint struct { 140 | Metric *MockedMetric 141 | Aggregator *MockedAggregator 142 | Comparator *MockedComparator 143 | Threshold float64 144 | } 145 | 146 | // NewMockedConstraint returns a new mocked constraint with already mocked 147 | // calls for a proper Constraint.Check function. Each call will return a result 148 | // from the results list. 149 | func NewMockedConstraint(results ...bool) *MockedConstraint { 150 | p := []tester.DataPoint{} 151 | n := len(results) 152 | c := &MockedConstraint{ 153 | Metric: new(MockedMetric), 154 | Aggregator: new(MockedAggregator), 155 | Comparator: new(MockedComparator), 156 | Threshold: float64(100), 157 | } 158 | 159 | c.Metric.On("Fetch", mock.Anything, mock.Anything).Return(p, nil).Times(n) 160 | c.Aggregator.On("Aggregate", p).Return(float64(50)).Times(n) 161 | 162 | for _, result := range results { 163 | if !result { 164 | c.Metric.On("Name").Return("Metric").Once() 165 | c.Aggregator.On("Name").Return("Aggregator").Once() 166 | c.Comparator.On("Name").Return("?").Twice() 167 | } 168 | 169 | c.Comparator.On("Compare", float64(50), float64(100)).Return(result).Once() 170 | } 171 | 172 | return c 173 | } 174 | 175 | func (m *MockedConstraint) Constraint() *tester.Constraint { 176 | return &tester.Constraint{ 177 | Metric: m.Metric, 178 | Aggregator: m.Aggregator, 179 | Comparator: m.Comparator, 180 | Threshold: m.Threshold, 181 | } 182 | } 183 | 184 | func (m *MockedConstraint) AssertExpectations(t *testing.T) { 185 | t.Helper() 186 | m.Metric.AssertExpectations(t) 187 | m.Aggregator.AssertExpectations(t) 188 | m.Comparator.AssertExpectations(t) 189 | } 190 | -------------------------------------------------------------------------------- /tester/run/concurrency.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package run 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/tester" 15 | "github.com/pinterest/bender" 16 | ) 17 | 18 | // LoadTestConcurrencyFixed runs predefined set of throughput tests. 19 | func LoadTestConcurrencyFixed(r tester.ConcurrencyRunner, o interface{}, ws ...tester.Workers) error { 20 | t := r.Tester() 21 | if err := t.Before(o); err != nil { 22 | //nolint:wrapcheck 23 | return err 24 | } 25 | 26 | defer t.After(o) 27 | 28 | for _, workers := range ws { 29 | if err := loadTestConcurrency(r, t, o, workers); err != nil { 30 | return err 31 | } 32 | } 33 | 34 | return nil 35 | } 36 | 37 | // LoadTestConcurrencyConstraints automatically tries to find a breakpoint based on provided constraints checks. 38 | func LoadTestConcurrencyConstraints(r tester.ConcurrencyRunner, o interface{}, start tester.Workers, g tester.Growth, 39 | cs ...*tester.Constraint) error { 40 | t := r.Tester() 41 | if err := t.Before(o); err != nil { 42 | //nolint:wrapcheck 43 | return err 44 | } 45 | 46 | defer t.After(o) 47 | 48 | workers := start 49 | for workers > 0 { 50 | startTime := time.Now() 51 | 52 | if err := loadTestConcurrency(r, t, o, workers); err != nil { 53 | return err 54 | } 55 | 56 | duration := time.Since(startTime) 57 | 58 | if checkConstraints(startTime, duration, cs...) { 59 | workers = g.OnSuccess(workers) 60 | } else { 61 | workers = g.OnFail(workers) 62 | } 63 | } 64 | 65 | return nil 66 | } 67 | 68 | // loadTestConcurrency runs a single test for a desired QPS. 69 | func loadTestConcurrency(r tester.ConcurrencyRunner, t tester.Tester, o interface{}, workers tester.Workers) error { 70 | if err := t.BeforeEach(o); err != nil { 71 | //nolint:wrapcheck 72 | return err 73 | } 74 | defer t.AfterEach(o) 75 | 76 | if err := r.Before(workers, o); err != nil { 77 | //nolint:wrapcheck 78 | return err 79 | } 80 | defer r.After(workers, o) 81 | 82 | executor, err := t.RequestExecutor(o) 83 | if err != nil { 84 | //nolint:wrapcheck 85 | return err 86 | } 87 | 88 | bender.LoadTestConcurrency(r.WorkerSemaphore(), r.Requests(), executor, r.Recorder()) 89 | bender.Record(r.Recorder(), r.Recorders()...) 90 | 91 | return nil 92 | } 93 | -------------------------------------------------------------------------------- /tester/run/throughput.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package run 10 | 11 | import ( 12 | "time" 13 | 14 | "github.com/facebookincubator/fbender/tester" 15 | "github.com/pinterest/bender" 16 | ) 17 | 18 | // LoadTestThroughputFixed runs predefined set of throughput tests. 19 | func LoadTestThroughputFixed(r tester.ThroughputRunner, o interface{}, qs ...tester.QPS) error { 20 | t := r.Tester() 21 | if err := t.Before(o); err != nil { 22 | //nolint:wrapcheck 23 | return err 24 | } 25 | 26 | defer t.After(o) 27 | 28 | for _, qps := range qs { 29 | if err := loadTestThroughput(r, t, o, qps); err != nil { 30 | return err 31 | } 32 | } 33 | 34 | return nil 35 | } 36 | 37 | // LoadTestThroughputConstraints automatically tries to find a breakpoint based on provided constraints checks. 38 | func LoadTestThroughputConstraints(r tester.ThroughputRunner, o interface{}, start tester.QPS, g tester.Growth, 39 | cs ...*tester.Constraint) error { 40 | t := r.Tester() 41 | if err := t.Before(o); err != nil { 42 | //nolint:wrapcheck 43 | return err 44 | } 45 | 46 | defer t.After(o) 47 | 48 | qps := start 49 | for qps > 0 { 50 | startTime := time.Now() 51 | 52 | if err := loadTestThroughput(r, t, o, qps); err != nil { 53 | return err 54 | } 55 | 56 | duration := time.Since(startTime) 57 | 58 | if checkConstraints(startTime, duration, cs...) { 59 | qps = g.OnSuccess(qps) 60 | } else { 61 | qps = g.OnFail(qps) 62 | } 63 | } 64 | 65 | return nil 66 | } 67 | 68 | // loadTestThroughput runs a single test for a desired QPS. 69 | func loadTestThroughput(r tester.ThroughputRunner, t tester.Tester, o interface{}, qps tester.QPS) error { 70 | if err := t.BeforeEach(o); err != nil { 71 | //nolint:wrapcheck 72 | return err 73 | } 74 | 75 | defer t.AfterEach(o) 76 | 77 | if err := r.Before(qps, o); err != nil { 78 | //nolint:wrapcheck 79 | return err 80 | } 81 | 82 | defer r.After(qps, o) 83 | 84 | executor, err := t.RequestExecutor(o) 85 | if err != nil { 86 | //nolint:wrapcheck 87 | return err 88 | } 89 | 90 | bender.LoadTestThroughput(r.Intervals(), r.Requests(), executor, r.Recorder()) 91 | bender.Record(r.Recorder(), r.Recorders()...) 92 | 93 | return nil 94 | } 95 | -------------------------------------------------------------------------------- /tester/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester 10 | 11 | import ( 12 | "errors" 13 | 14 | "github.com/pinterest/bender" 15 | ) 16 | 17 | // ErrInvalidOptions is thrown when options don't implement the required interface. 18 | var ErrInvalidOptions = errors.New("invalid options") 19 | 20 | // Tester is used to setup the test for a specific endpoint. 21 | type Tester interface { 22 | // Before is called once, before any tests. 23 | Before(options interface{}) error 24 | // After is called once, after all tests (or after some of them if a test fails). 25 | // This should be used to cleanup everything that was set up in the Before. 26 | After(options interface{}) 27 | // BeforeEach is called before every test. 28 | BeforeEach(options interface{}) error 29 | // AfterEach is called after every test, even if the test fails. This should 30 | // be used to cleanup everything that was set up in the BeforeEach. 31 | AfterEach(options interface{}) 32 | // RequestExecutor is called every time a test is to be ran to get an executor. 33 | RequestExecutor(options interface{}) (bender.RequestExecutor, error) 34 | } 35 | 36 | // QPS is the test desired queries per second. 37 | type QPS = int 38 | 39 | // ThroughputRunner is used to setup the test execution. 40 | type ThroughputRunner interface { 41 | // Before is called before running a test. 42 | Before(qps QPS, options interface{}) error 43 | // After is called after test finishes. This should be used to clean up 44 | // everything that was ser up in the Before. 45 | After(qps QPS, options interface{}) 46 | 47 | // Protocol tester. 48 | Tester() Tester 49 | 50 | // Params used by LoadTestThroughput function. 51 | Intervals() bender.IntervalGenerator 52 | Requests() chan interface{} 53 | Recorder() chan interface{} 54 | Recorders() []bender.Recorder 55 | } 56 | 57 | // Workers is the test desired concurrent workers. 58 | type Workers = int 59 | 60 | // ConcurrencyRunner is used to setup concurrency test execution. 61 | type ConcurrencyRunner interface { 62 | // Before is called before running a test. 63 | Before(workers Workers, options interface{}) error 64 | // After is called after test finishes. This should be used to clean up 65 | // everything that was ser up in the Before. 66 | After(workers Workers, options interface{}) 67 | 68 | // Protocol tester. 69 | Tester() Tester 70 | 71 | // Params used by LoadTestConcurrency function. 72 | WorkerSemaphore() *bender.WorkerSemaphore 73 | Requests() chan interface{} 74 | Recorder() chan interface{} 75 | Recorders() []bender.Recorder 76 | } 77 | -------------------------------------------------------------------------------- /tester/tester_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tester_test 10 | 11 | import ( 12 | "reflect" 13 | "testing" 14 | 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | // assertPointerEqual checks whether two pointers are equal. 19 | func assertPointerEqual(t *testing.T, expected, value interface{}, args ...interface{}) { 20 | t.Helper() 21 | 22 | expectedPointer := reflect.ValueOf(expected).Pointer() 23 | valuePointer := reflect.ValueOf(value).Pointer() 24 | assert.Equal(t, expectedPointer, valuePointer, args...) 25 | } 26 | -------------------------------------------------------------------------------- /tester/tftp/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package tftp 10 | 11 | import ( 12 | "fmt" 13 | "time" 14 | 15 | "github.com/pin/tftp" 16 | "github.com/pinterest/bender" 17 | protocol "github.com/pinterest/bender/tftp" 18 | ) 19 | 20 | // Tester is a load tester for TFTP. 21 | type Tester struct { 22 | Target string 23 | Timeout time.Duration 24 | BlockSize int 25 | 26 | client *tftp.Client 27 | } 28 | 29 | // Before is called before the first test. 30 | func (t *Tester) Before(options interface{}) error { 31 | var err error 32 | 33 | t.client, err = tftp.NewClient(t.Target) 34 | if err != nil { 35 | return fmt.Errorf("unable to set up the tester: %w", err) 36 | } 37 | 38 | t.client.SetTimeout(t.Timeout) 39 | t.client.SetBlockSize(t.BlockSize) 40 | 41 | return nil 42 | } 43 | 44 | // After is called after all tests are finished. 45 | func (t *Tester) After(_ interface{}) {} 46 | 47 | // BeforeEach is called before every test. 48 | func (t *Tester) BeforeEach(_ interface{}) error { 49 | return nil 50 | } 51 | 52 | // AfterEach is called after every test. 53 | func (t *Tester) AfterEach(_ interface{}) {} 54 | 55 | // RequestExecutor returns a request executor. 56 | func (t *Tester) RequestExecutor(_ interface{}) (bender.RequestExecutor, error) { 57 | return protocol.CreateExecutor(t.client, protocol.DiscardingValidator), nil 58 | } 59 | -------------------------------------------------------------------------------- /tester/udp/tester.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package udp 10 | 11 | import ( 12 | "time" 13 | 14 | protocol "github.com/facebookincubator/fbender/protocols/udp" 15 | "github.com/pinterest/bender" 16 | ) 17 | 18 | // Tester is a load tester for UDP. 19 | type Tester struct { 20 | Target string 21 | Timeout time.Duration 22 | Validator protocol.ResponseValidator 23 | } 24 | 25 | // Before is called before the first test. 26 | func (t *Tester) Before(_ interface{}) error { 27 | return nil 28 | } 29 | 30 | // After is called after all tests are finished. 31 | func (t *Tester) After(_ interface{}) {} 32 | 33 | // BeforeEach is called before every test. 34 | func (t *Tester) BeforeEach(_ interface{}) error { 35 | return nil 36 | } 37 | 38 | // AfterEach is called after every test. 39 | func (t *Tester) AfterEach(_ interface{}) {} 40 | 41 | func validator(_ *protocol.Datagram, _ []byte) error { 42 | return nil 43 | } 44 | 45 | // RequestExecutor returns a request executor. 46 | func (t *Tester) RequestExecutor(options interface{}) (bender.RequestExecutor, error) { 47 | if t.Validator == nil { 48 | return protocol.CreateExecutor(t.Timeout, validator, t.Target), nil 49 | } 50 | 51 | return protocol.CreateExecutor(t.Timeout, t.Validator, t.Target), nil 52 | } 53 | -------------------------------------------------------------------------------- /utils/completion.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils 10 | 11 | import ( 12 | "fmt" 13 | "strings" 14 | 15 | "github.com/spf13/cobra" 16 | "github.com/spf13/pflag" 17 | ) 18 | 19 | const template = ` 20 | %s() { 21 | %s 22 | }` 23 | 24 | // BashCompletion annotates the flag with completion function and registers the 25 | // completion function in the root command if it hasn't been added already. 26 | func BashCompletion(cmd *cobra.Command, flags *pflag.FlagSet, flag string, fname, fbody string) error { 27 | if !strings.Contains(cmd.Root().BashCompletionFunction, fname) { 28 | cmd.Root().BashCompletionFunction += fmt.Sprintf(template, fname, fbody) 29 | } 30 | 31 | return cobra.MarkFlagCustom(flags, flag, fname) 32 | } 33 | -------------------------------------------------------------------------------- /utils/completion_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils_test 10 | 11 | import ( 12 | "strings" 13 | "testing" 14 | 15 | "github.com/facebookincubator/fbender/utils" 16 | "github.com/spf13/cobra" 17 | "github.com/stretchr/testify/assert" 18 | "github.com/stretchr/testify/require" 19 | ) 20 | 21 | func TestBashCompletion(t *testing.T) { 22 | const ( 23 | fname = "__fbender_test_handle_flag" 24 | fbody = `COMPREPLY=($(compgen -W "42 24" -- "${cur}"))` 25 | ) 26 | 27 | c := &cobra.Command{} 28 | c.Flags().Int("myint", 0, "set myint") 29 | 30 | // Check if the completion function is appended 31 | err := utils.BashCompletion(c, c.Flags(), "myint", fname, fbody) 32 | require.NoError(t, err) 33 | assert.Contains(t, c.BashCompletionFunction, fname) 34 | assert.Equal(t, ` 35 | __fbender_test_handle_flag() { 36 | COMPREPLY=($(compgen -W "42 24" -- "${cur}")) 37 | }`, c.BashCompletionFunction) 38 | 39 | // Check if the flag annotation has been added 40 | f := c.Flags().Lookup("myint") 41 | require.NotNil(t, f) 42 | require.Contains(t, f.Annotations, "cobra_annotation_bash_completion_custom") 43 | assert.Equal(t, []string{fname}, f.Annotations["cobra_annotation_bash_completion_custom"]) 44 | 45 | // Check if the completion function is appended only once 46 | err = utils.BashCompletion(c, c.Flags(), "myint", fname, fbody) 47 | require.NoError(t, err) 48 | assert.Contains(t, c.BashCompletionFunction, fname) 49 | count := strings.Count(c.BashCompletionFunction, fname) 50 | assert.Equal(t, 1, count, "Completion function should be added only once") 51 | } 52 | -------------------------------------------------------------------------------- /utils/net.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils 10 | 11 | import ( 12 | "net" 13 | "strconv" 14 | ) 15 | 16 | // WithDefaultPort adds a default port if no port is present. 17 | func WithDefaultPort(hostport string, port int) string { 18 | if _, _, err := net.SplitHostPort(hostport); err != nil { 19 | return net.JoinHostPort(hostport, strconv.Itoa(port)) 20 | } 21 | 22 | return hostport 23 | } 24 | -------------------------------------------------------------------------------- /utils/net_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/utils" 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | func TestWithDefaultPort(t *testing.T) { 19 | // Adds default port 20 | assert.Equal(t, "[::1]:53", utils.WithDefaultPort("::1", 53)) 21 | assert.Equal(t, "127.0.0.1:53", utils.WithDefaultPort("127.0.0.1", 53)) 22 | assert.Equal(t, "example.com:53", utils.WithDefaultPort("example.com", 53)) 23 | // Does not change port when present 24 | assert.Equal(t, "[::1]:5353", utils.WithDefaultPort("[::1]:5353", 53)) 25 | assert.Equal(t, "127.0.0.1:5353", utils.WithDefaultPort("127.0.0.1:5353", 53)) 26 | assert.Equal(t, "example.com:5353", utils.WithDefaultPort("example.com:5353", 53)) 27 | } 28 | -------------------------------------------------------------------------------- /utils/random.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils 10 | 11 | import ( 12 | "crypto/rand" 13 | "encoding/hex" 14 | ) 15 | 16 | // RandomHex generates a random hex string of specified length. 17 | func RandomHex(n int) (string, error) { 18 | b := make([]byte, (n+1)/2) 19 | if _, err := rand.Read(b); err != nil { 20 | //nolint:wrapcheck 21 | return "", err 22 | } 23 | 24 | return hex.EncodeToString(b)[:n], nil 25 | } 26 | -------------------------------------------------------------------------------- /utils/random_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/utils" 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | func TestRandomHex(t *testing.T) { 19 | // Even length 20 | hex, err := utils.RandomHex(16) 21 | assert.NoError(t, err) 22 | assert.Len(t, hex, 16) 23 | // Odd length 24 | hex, err = utils.RandomHex(9) 25 | assert.NoError(t, err) 26 | assert.Len(t, hex, 9) 27 | // Corner case 28 | hex, err = utils.RandomHex(0) 29 | assert.NoError(t, err) 30 | assert.Len(t, hex, 0) 31 | } 32 | -------------------------------------------------------------------------------- /utils/regexp.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils 10 | 11 | import ( 12 | "regexp" 13 | ) 14 | 15 | // NamedRegex is a regex which supports named capture groups. 16 | type NamedRegex struct { 17 | *regexp.Regexp 18 | } 19 | 20 | // FindStringSubmatchMap returns a map of named capture groups. 21 | func (r *NamedRegex) FindStringSubmatchMap(s string) map[string]string { 22 | captures := make(map[string]string) 23 | match := r.FindStringSubmatch(s) 24 | 25 | if match == nil { 26 | return captures 27 | } 28 | 29 | for i, name := range r.SubexpNames() { 30 | // Ignore the whole regexp match and unnamed groups 31 | if i == 0 || name == "" { 32 | continue 33 | } 34 | 35 | captures[name] = match[i] 36 | } 37 | 38 | return captures 39 | } 40 | 41 | // MustCompile compiles a string to a named regexp. 42 | func MustCompile(s string) NamedRegex { 43 | return NamedRegex{Regexp: regexp.MustCompile(s)} 44 | } 45 | -------------------------------------------------------------------------------- /utils/regexp_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils_test 10 | 11 | import ( 12 | "testing" 13 | 14 | "github.com/facebookincubator/fbender/utils" 15 | "github.com/stretchr/testify/assert" 16 | ) 17 | 18 | func TestFindStringSubmatchMap(t *testing.T) { 19 | r := utils.MustCompile(`^(?P[A-Z][a-z]*) (?P[A-Z][a-z]*)$`) 20 | // this should match 21 | match := r.FindStringSubmatchMap("Mikolaj Walczak") 22 | expected := map[string]string{"name": "Mikolaj", "surname": "Walczak"} 23 | assert.Equal(t, expected, match) 24 | // this should fail 25 | match = r.FindStringSubmatchMap("12345") 26 | expected = map[string]string{} 27 | assert.Equal(t, expected, match) 28 | 29 | // With optional fields 30 | r = utils.MustCompile(`^(?P[A-Z][a-z]*)( (?P[A-Z][a-z]*))?$`) 31 | // this should still match 32 | match = r.FindStringSubmatchMap("Mikolaj Walczak") 33 | expected = map[string]string{"name": "Mikolaj", "surname": "Walczak"} 34 | assert.Equal(t, expected, match) 35 | // this should also match and have an empty surname 36 | match = r.FindStringSubmatchMap("Mikolaj") 37 | expected = map[string]string{"name": "Mikolaj", "surname": ""} 38 | assert.Equal(t, expected, match) 39 | // this should fail 40 | match = r.FindStringSubmatchMap("12345") 41 | expected = map[string]string{} 42 | assert.Equal(t, expected, match) 43 | } 44 | -------------------------------------------------------------------------------- /utils/spinner.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils 10 | 11 | import ( 12 | "context" 13 | "time" 14 | 15 | "github.com/facebookincubator/fbender/log" 16 | spin "github.com/tj/go-spin" 17 | ) 18 | 19 | // Default refresh rate. 20 | const spinnerRefresh = 100 * time.Millisecond 21 | 22 | // NewBackgroundSpinner creates a new spinner which runs in background refreshing 23 | // its output on a constant rate. The spinner is prefixed with the provided 24 | // description. Returns a function which cancels the spinner. The default value 25 | // for refresh is used if provided refresh is less than or equal to zero. 26 | func NewBackgroundSpinner(description string, refresh time.Duration) context.CancelFunc { 27 | if refresh <= 0 { 28 | refresh = spinnerRefresh 29 | } 30 | 31 | ctx, cancel := context.WithCancel(context.Background()) 32 | 33 | spinner := spin.New() 34 | spinner.Set(spin.Spin1) 35 | 36 | sync := make(chan bool) 37 | 38 | go func() { 39 | for { 40 | log.Errorf("\r%s... ", description) 41 | select { 42 | case <-ctx.Done(): 43 | log.Errorf("Done.\n") 44 | sync <- false 45 | 46 | return 47 | default: 48 | log.Errorf("%s", spinner.Next()) 49 | } 50 | time.Sleep(refresh) 51 | } 52 | }() 53 | 54 | return func() { 55 | // cancel context and wait for goroutine to clean the output 56 | cancel() 57 | <-sync 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /utils/spinner_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) Facebook, Inc. and its affiliates. 3 | All rights reserved. 4 | 5 | This source code is licensed under the BSD-style license found in the 6 | LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | package utils_test 10 | 11 | import ( 12 | "bytes" 13 | "testing" 14 | "time" 15 | 16 | "github.com/facebookincubator/fbender/log" 17 | "github.com/facebookincubator/fbender/utils" 18 | "github.com/stretchr/testify/assert" 19 | ) 20 | 21 | func TestNewBackgroundSpinner(t *testing.T) { 22 | w := new(bytes.Buffer) 23 | log.Stderr = w 24 | 25 | cancel := utils.NewBackgroundSpinner("Testing", 100*time.Millisecond) 26 | 27 | time.Sleep(1 * time.Second) 28 | cancel() 29 | 30 | v := w.String() 31 | 32 | assert.Contains(t, v, "\rTesting... |") 33 | assert.Contains(t, v, "\rTesting... /") 34 | assert.Contains(t, v, "\rTesting... -") 35 | assert.Contains(t, v, "\rTesting... \\") 36 | assert.Contains(t, v, "\rTesting... Done.") 37 | } 38 | --------------------------------------------------------------------------------