[ { "path": ".dockerignore", "content": "*\n\n!pgo\n!src\n!Cargo.*\n" }, { "path": ".github/FUNDING.yml", "content": "# These are supported funding model platforms\n\ngithub: hatoo # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]\npatreon: # Replace with a single Patreon username\nopen_collective: # Replace with a single Open Collective username\nko_fi: hatoo # Replace with a single Ko-fi username\ntidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel\ncommunity_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry\nliberapay: # Replace with a single Liberapay username\nissuehunt: # Replace with a single IssueHunt username\notechie: # Replace with a single Otechie username\nlfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry\ncustom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']\n" }, { "path": ".github/dependabot.yml", "content": "# To get started with Dependabot version updates, you'll need to specify which\n# package ecosystems to update and where the package manifests are located.\n# Please see the documentation for all configuration options:\n# https://docs.github.com/en/code-security/dependabot/working-with-dependabot/dependabot-options-reference\n\nversion: 2\nupdates:\n - package-ecosystem: \"cargo\" # See documentation for possible values\n directories: # Locations of package manifests\n - \"/\"\n - \"pgo/server\"\n schedule:\n interval: \"weekly\"\n - package-ecosystem: \"docker\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n" }, { "path": ".github/workflows/CI.yml", "content": "on:\n push:\n branches:\n - master\n pull_request:\n\nname: CI\n\njobs:\n test:\n name: Test Suite\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n os: [ubuntu-latest, windows-latest, macos-latest, macos-14]\n additional_args: [\"\", \"--no-default-features --features native-tls\", \"--features http3\"]\n # vsock feature is only on linux\n include:\n - os: [\"ubuntu-latest\"]\n additional_args: \"--features vsock\"\n steps:\n # We need nasm to build aws-lc on windows\n - uses: ilammy/setup-nasm@v1\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n - uses: Swatinem/rust-cache@v2\n - run: cargo test ${{ matrix.additional_args }}\n\n fmt:\n name: Rustfmt\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n with:\n components: rustfmt\n - name: Enforce formatting\n run: cargo fmt --all --check\n\n clippy:\n name: Clippy\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n with:\n components: clippy\n - uses: Swatinem/rust-cache@v2\n - name: Linting\n run: cargo clippy --all-targets --all-features -- -D warnings\n" }, { "path": ".github/workflows/release-pgo.yml", "content": "name: Publish PGO\n\non:\n push:\n branches:\n - master\n tags:\n - \"v*.*.*\"\n pull_request:\n\njobs:\n publish:\n name: Publish for ${{ matrix.target }} PGO\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n include:\n - os: ubuntu-latest\n artifact_name: oha\n release_name: oha-linux-amd64-pgo\n target: x86_64-unknown-linux-musl\n additional_args: \"--features vsock\"\n # tikv-jemalloc-sys@0.6.0+5.3.0-1: background_threads_runtime_support not supported for aarch64-unknown-linux-musl\n# - os: ubuntu-24.04-arm\n# artifact_name: oha\n# release_name: oha-linux-arm64-pgo\n# target: aarch64-unknown-linux-musl\n# additional_args: \"--features vsock\"\n - os: windows-latest\n artifact_name: oha.exe\n release_name: oha-windows-amd64-pgo.exe\n target: x86_64-pc-windows-msvc\n additional_args: \"\"\n# - os: macos-14-large\n# artifact_name: oha\n# release_name: oha-macos-amd64-pgo\n# target: x86_64-apple-darwin\n# additional_args: \"\"\n# - os: macos-14\n# artifact_name: oha\n# release_name: oha-macos-arm64-pgo\n# target: aarch64-apple-darwin\n# additional_args: \"\"\n\n steps:\n - uses: ilammy/setup-nasm@v1\n - uses: actions/checkout@v6\n - name: Install musl-tools on Linux\n run: sudo apt-get update --yes && sudo apt-get install --yes musl-tools\n if: contains(matrix.target, 'musl')\n - uses: dtolnay/rust-toolchain@stable\n with:\n targets: ${{ matrix.target }}\n components: llvm-tools-preview\n - uses: Swatinem/rust-cache@v2\n - run: cargo install cargo-pgo --version 0.2.9\n - uses: oven-sh/setup-bun@v2\n - run: bun run pgo.js --target ${{ matrix.target }} ${{ matrix.additional_args }}\n - uses: actions/upload-artifact@v7\n with:\n name: ${{ matrix.release_name }}\n path: target/${{ matrix.target }}/pgo/${{ matrix.artifact_name }}\n - name: Upload binaries to release\n uses: svenstaro/upload-release-action@v2\n if: startsWith(github.ref, 'refs/tags/v')\n with:\n repo_token: ${{ secrets.GITHUB_TOKEN }}\n file: target/${{ matrix.target }}/pgo/${{ matrix.artifact_name }}\n asset_name: ${{ matrix.release_name }}\n" }, { "path": ".github/workflows/release.yml", "content": "name: Publish\n\non:\n push:\n branches:\n - master\n tags:\n - \"v*.*.*\"\n pull_request:\n\njobs:\n publish:\n name: Publish for ${{ matrix.target }}\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n include:\n - os: ubuntu-latest\n artifact_name: oha\n release_name: oha-linux-amd64\n target: x86_64-unknown-linux-musl\n additional_args: \"--features vsock\"\n use_cross: false\n - os: windows-latest\n artifact_name: oha.exe\n release_name: oha-windows-amd64.exe\n target: x86_64-pc-windows-msvc\n additional_args: \"\"\n use_cross: false\n - os: macos-latest\n artifact_name: oha\n release_name: oha-macos-amd64\n target: x86_64-apple-darwin\n additional_args: \"\"\n use_cross: false\n - os: ubuntu-latest\n artifact_name: oha\n release_name: oha-linux-arm64\n target: aarch64-unknown-linux-musl\n additional_args: \"--features vsock\"\n use_cross: true\n - os: macos-14\n artifact_name: oha\n release_name: oha-macos-arm64\n target: aarch64-apple-darwin\n additional_args: \"\"\n use_cross: false\n env:\n BUILD_CMD: cargo\n steps:\n - uses: ilammy/setup-nasm@v1\n - uses: actions/checkout@v6\n - name: Install musl-tools on Linux\n run: sudo apt-get update --yes && sudo apt-get install --yes musl-tools\n if: contains(matrix.target, 'musl')\n - name: Install Rust toolchain\n uses: dtolnay/rust-toolchain@stable\n with:\n targets: ${{ matrix.target }}\n - name: Install cross\n if: matrix.use_cross\n uses: taiki-e/install-action@v2\n with:\n tool: cross\n - uses: Swatinem/rust-cache@v2\n - name: Overwrite build command env variable\n if: matrix.use_cross\n shell: bash\n run: echo \"BUILD_CMD=cross\" >> $GITHUB_ENV\n - name: Build\n shell: bash\n run: $BUILD_CMD build --profile release-ci --target ${{ matrix.target }} --locked --no-default-features --features rustls ${{ matrix.additional_args }}\n - name: Upload\n uses: actions/upload-artifact@v7\n with:\n name: ${{ matrix.release_name }}\n path: target/${{ matrix.target }}/release-ci/${{ matrix.artifact_name }}\n - name: Upload binaries to release\n if: startsWith(github.ref, 'refs/tags/v')\n uses: svenstaro/upload-release-action@v2\n with:\n repo_token: ${{ secrets.GITHUB_TOKEN }}\n file: target/${{ matrix.target }}/release-ci/${{ matrix.artifact_name }}\n asset_name: ${{ matrix.release_name }}\n tag: ${{ github.ref }}\n\n docker:\n name: Build and Push Docker Image\n runs-on: ubuntu-latest\n needs: publish\n if: startsWith(github.ref, 'refs/tags/v')\n permissions:\n contents: read\n packages: write\n steps:\n - uses: actions/checkout@v6\n\n - name: Set up QEMU\n uses: docker/setup-qemu-action@v3\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n\n - name: Login to GitHub Container Registry\n if: startsWith(github.ref, 'refs/tags/v')\n uses: docker/login-action@v3\n with:\n registry: ghcr.io\n username: ${{ github.actor }}\n password: ${{ secrets.GITHUB_TOKEN }}\n\n - name: Extract metadata (tags, labels)\n id: meta\n uses: docker/metadata-action@v5\n with:\n images: ghcr.io/${{ github.repository }}\n tags: |\n type=semver,pattern={{version}}\n type=semver,pattern={{major}}.{{minor}}\n type=semver,pattern={{major}}\n type=raw,value=latest,enable=${{ startsWith(github.ref, 'refs/tags/v') }}\n\n - name: Build and push\n uses: docker/build-push-action@v6\n with:\n context: .\n platforms: linux/amd64,linux/arm64\n push: ${{ startsWith(github.ref, 'refs/tags/v') }}\n tags: ${{ steps.meta.outputs.tags }}\n labels: ${{ steps.meta.outputs.labels }}\n cache-from: type=gha\n cache-to: type=gha,mode=max\n" }, { "path": ".gitignore", "content": "**/target\n/.idea\n/.vscode\n" }, { "path": "CHANGELOG.md", "content": "# Unreleased\n\n# 1.14.0\n\n- fix Possible bug with latency_correction #857\n\n# 1.13.0 (2025-02-07)\n\n- Add first byte stats to JSON output#844\n- On http3, bind ipv6 for ipv6#839\n- Prefer ipv6 for localhost on macos#837\n\n# 1.12.1 (2025-11-29)\n\n- feat: add official docker img with better caching #830 \n\n# 1.12.0 (2025-11-29)\n\n- Fix dns and connection time stats #816\n- Add colors to --help/-h #822 \n- rename --disable-color to --no-color #824 \n- fixed --urls-from-file reuses host with path from different url #825\n\n# 1.11.0 (2025-11-01)\n\n- feat:support decimal when using n_requests,such as 2.6k or 0.001m #808 \n- connect-to now doing tcp connect instead of dns caching #807 \n- feat: make connect timeout configurable #805 \n- feat: support k/m suffixes for n-requests #800 \n- feat: Add option to read request body line-by-line from file. #729 #789\n\n# 1.10.0 (2025-09-06)\n\n- fix dns lookup on http2/http3 #771 \n- feat: add curl-compatible multipart form data support (-F option) #755\n\n# 1.9.0 (2025-06-21)\n\n- Implement experimental HTTP3 support #746 \n- Allow appending to database if oha table has already been created #742 \n- Add -u/--time-unit option #741 \n- Add RequestResult.first_byte field for measuring first body byte latency #727 \n- Add support for results in csv format #725 \n- Add support for fractional QPS values #724 \n\n# 1.8.0 (2025-02-15)\n\n- Support mtls #687 \n- Support Proxy headers #688 \n- Randomize --connect-to if multiple matching options #695 \n\n# 1.7.0 (2025-02-01)\n\n- Impl support for calling AWS APIs with sigv4 #666 \n- support -o #669 \n\n# 1.6.0 (2025-01-11)\n\n- Feature: Reading Urls from file #639\n- Add some optimized workers to `--no-tui` mode #646\n\n# 1.5.0 (2024-12-07)\n\n- Add `--debug` option to check actual request/response\n- Switch colors to justified latency thresholds (fixes #609) #610 \n- Fix Running with -q hangs #603 \n- Support HTTP proxy #614 \n\n# 1.4.7 (2024-10-26)\n\n- [rustls] Cache HTTPS certs\n\n# 1.4.6 (2024-08-17)\n\n- Add `--wait-ongoing-requests-after-deadline` option\n- Add `--db-url` option to save results to SQLite database\n- Add `--dump-urls` option to debug random URL generation\n\n# 1.4.5 (2024-05-29)\n\n- Some performance improvements\n\n# 1.4.4 (2024-04-20)\n\n- support Termux #464\n\n# 1.4.3 (2024-04-06)\n\n- fix rustls error #452\n\n# 1.4.2 (2024-04-06)\n\n- Fix printing of Size/request #447 \n\n# 1.4.1 (2024-03-16)\n\n- Enable: Profile-Guided Optimization (PGO) #268\n\n# 1.4.0 (2024-03-09)\n\n- No DNS lookup when unix socket or vsock #418\n- Add HTTP over VSOCK support #416\n\n# 1.3.0 (2024-02-04)\n\n- Optimize timeout #403 \n- Compact error #402 \n- fix tui layout #401 \n\n# 1.2.0 (2024-02-03)\n\n- Print help message when no argument is given #378\n- Lookup DNS at beginning and cache it #391\n- Report deadlined requests #392 \n- Fix MacOS Crash issues #384\n\n# 1.1.0 (2024-01-16)\n\n- [HTTP2] Reconnect TCP connection when it fails #369 \n\n# 1.0.0 (2023-11-16)\n\n- Update hyper dependency to 1.0.0\n\n# 1.0.0-rc.4.a8dcd7ca5df49c0701893c4d9d81ec8c1342f141 (2023-10-14)\n\nThis is a RC release for 1.0.0. Please test it and report any issues.\nThe version is named as same as `hyper`'s version and it's commit hash.\n\nSince this version depends on unreleased `hyper`'s version, we can't release on crates.io. Only on binary releases.\n\n- Support HTTP/2 #224 #201\n- Make `rustls` as a default TLS backend #331\n- Added `-p` option to set number of HTTP/2 parallel requests\n\n# 0.6.5 (2023-10-09)\n\n- Fix Apple Silicon's binary release #323\n\n# 0.6.4 (2023-09-24)\n\n- Fix -H option to overwrite default value #309\n- feat: display 99.90- and 99.99-percentile latency #315 \n\n# 0.6.3 (2023-09-05)\n\n- Add style and colors to the summary view #64\n- Added a stats-success-breakdown flag for more detailed status code specific response statistics #212\n\n# 0.6.2 (2023-08-12)\n\n- Support Burst feature #276\n\n# 0.6.1 (2023-07-12)\n\n- Fix sending HTTP uri #255\n- Add default user agent header #257\n\n# 0.6.0 (2023-06-24)\n\n- Support IDNA #236\n- Support randomly generated URL using rand_regex crate\n\n# 0.5.9 (2023-06-12)\n\n- Fix -H Header parser\n- Update printer #229\n - Use percentage for Success rate summary value #228 \n - Latency distribution -> Response time distribution\n\n# 0.5.8 (2023-03-25)\n\n- Add `--unix-socket` on `unix` profiles for HTTP. #220\n- Fix tui to not requiring True Color. #209\n\n# 0.5.7 (2023-02-25)\n\n- Fix `--latency-correction` adds the time of DNS. #211\n- Fix `-z` behaviour to cancel workers at the dead line. #211\n- Fix align of histogram #210\n\n# 0.5.6 (2023-02-02)\n\n- Update `clap` to version 4\n- Release `musl` binaries #206\n- Support [Ipv6] format requested_host in --connect-to #197\n\n# 0.5.5 (2022-09-19)\n\n- Add colors to the tui view #64\n\n# 0.5.4 (2022-08-27)\n\n- Support Ipv6 host #181\n- Print min, max, average and pXX for Requests per second in JSON output like bombardier #177\n- Add JSON Output #169\n- Fix QPS control to send with correct rate for first 1 sec\n- Make histogram compatible to hey\n - closes #161\n\n# 0.5.3 (2022-07-16)\n\n- Add support for bracketed IPv6 syntax in connect-to\n\n# 0.5.2 (2022-04-28)\n\n- Add `rustls` feature flag to build against `rustls` instead of `native-tls`.\n\n# 0.5.1 (2022-03-29)\n\n- Fix histogram to show correct response time\n - closes #157\n\n# 0.5.0 (2022-01-01)\n\n- Use clap 3.0.0 instead of structopt\n - closes #131\n\n# 0.4.6 (2021-07-05)\n\n- Add `--latency-correction` to avoid Coordinated Omission Problem.\n\n# 0.4.5 (2021-05-04)\n\n- Set '--no-tui' automatically when stdout isn't TTY\n\n# 0.4.4 (2020-11-18)\n\n- Bump `resolv-conf` to support `options edns0 trust-ad` on `/etc/resolv.conf`\n\n# 0.4.3 (2020-11-12)\n\n- Add --connect-to option to override DNS for a given host+port, similar to curl\n\n# 0.4.2 (2020-10-06)\n\n- Speed up on WSL Ubuntu 20.4\n\n# 0.4.1 (2020-07-28)\n\n- Support -q 0 option for unlimited qps\n- Fix performance on limiting query/second\n" }, { "path": "Cargo.toml", "content": "[package]\nauthors = [\"hatoo \"]\ncategories = [\n \"command-line-utilities\",\n \"network-programming\",\n \"web-programming::http-client\",\n \"development-tools::profiling\",\n]\ndescription = \"Ohayou(おはよう), HTTP load generator, inspired by rakyll/hey with tui animation.\"\nedition = \"2024\"\nkeywords = [\"cli\", \"load-testing\", \"performance\", \"http\"]\nlicense = \"MIT\"\nname = \"oha\"\nreadme = \"README.md\"\nrepository = \"https://github.com/hatoo/oha\"\nversion = \"1.14.0\"\nrust-version = \"1.87\"\n\n# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html\n\n[features]\ndefault = [\"rustls\"]\nnative-tls = [\"dep:native-tls\", \"dep:tokio-native-tls\"]\nrustls = [\n \"dep:rustls\",\n \"dep:tokio-rustls\",\n \"dep:rustls-native-certs\",\n \"dep:rustls-pki-types\",\n]\nvsock = [\"dep:tokio-vsock\"]\nhttp3 = [\"dep:h3\", \"dep:h3-quinn\", \"dep:quinn-proto\", \"dep:quinn\", \"dep:http\"]\n\n[dependencies]\nanyhow = \"1.0.86\"\naverage = \"0.16.0\"\nbyte-unit = \"5.1.4\"\nclap = { version = \"4.5.9\", features = [\"derive\", \"env\"] }\nclap_complete = \"4\"\nfloat-ord = \"0.3.2\"\nkanal = \"0.1.1\"\nhumantime = \"2.1.0\"\nlibc = \"0.2.155\"\nserde = { version = \"1.0.204\", features = [\"derive\"] }\nserde_json = \"1.0\"\nthiserror = \"2.0.12\"\ntokio = { version = \"1.38.1\", features = [\"full\"] }\nratatui = { version = \"0.30.0\", default-features = false, features = [\n \"crossterm\",\n] }\naws-sign-v4 = \"0.3\"\nchrono = \"0.4\"\nbytes = \"1\"\n\nhyper = { version = \"1.4\", features = [\"client\", \"http1\", \"http2\"] }\n\n# native-tls\nnative-tls = { version = \"0.2.12\", features = [\"alpn\"], optional = true }\ntokio-native-tls = { version = \"0.3.1\", optional = true }\n\nrustls = { version = \"0.23.18\", optional = true }\nrustls-native-certs = { version = \"0.8.0\", optional = true }\ntokio-rustls = { version = \"0.26.0\", optional = true }\nrustls-pki-types = { version = \"1.7.0\", optional = true }\n\nh3 = { version = \"0.0.8\", optional = true }\nh3-quinn = { version = \"0.0.10\", optional = true }\nquinn-proto = { version = \"0.11.10\", optional = true, features = [\"aws-lc-rs\"] }\nhttp = { version = \"1.4.0\", optional = true }\nquinn = { version = \"0.11.7\", optional = true, features = [\n \"aws-lc-rs\",\n \"runtime-tokio\",\n] }\n\nbase64 = \"0.22.1\"\nrand = \"0.10.0\"\nhickory-resolver = { version = \"0.25.2\", features = [\"tokio\"] }\nrand_regex = \"0.19.0\"\nregex-syntax = \"0.8.5\"\nurl = \"2.5.2\"\nhttp-body-util = \"0.1.2\"\nhyper-util = { version = \"0.1.6\", features = [\"tokio\"] }\ntokio-vsock = { version = \"0.7.2\", optional = true }\nrusqlite = { version = \"0.38.0\", features = [\"bundled\"] }\nnum_cpus = \"1.16.0\"\ntokio-util = \"0.7.13\"\nclap-cargo = \"0.18.3\"\n\n[target.'cfg(not(target_env = \"msvc\"))'.dependencies]\ntikv-jemallocator = \"0.6\"\n\n[target.'cfg(unix)'.dependencies]\nrlimit = \"0.11.0\"\n\n[dev-dependencies]\naxum = { version = \"0.8.1\", features = [\"http2\"] }\naxum-server = { version = \"0.8.0\", features = [\"tls-rustls\"] }\nbytes = \"1.6\"\nfloat-cmp = \"0.10.0\"\nhttp-mitm-proxy = { version = \"0.18.0\", default-features = false }\njsonschema = \"0.42.2\"\nlazy_static = \"1.5.0\"\npredicates = \"3.1.0\"\n# features = [\"aws_lc_rs\"] is a workaround for mac & native-tls\n# https://github.com/sfackler/rust-native-tls/issues/225\nrcgen = { version = \"0.14.3\", features = [\"aws_lc_rs\"] }\nregex = \"1.10.5\"\ntempfile = \"3.10.1\"\nrustls = \"0.23.18\"\nrstest = \"0.26.0\"\nrstest_reuse = \"0.7.0\"\nctor = \"0.6.1\"\n\n[target.'cfg(unix)'.dev-dependencies]\nactix-web = \"4\"\n\n[profile.pgo]\ninherits = \"release\"\n# https://github.com/TechEmpower/FrameworkBenchmarks/blob/master/frameworks/Rust/faf/Cargo.toml + lto=true\nopt-level = 3\npanic = 'abort'\ncodegen-units = 1\nlto = true\ndebug = false\nincremental = false\noverflow-checks = false\n\n[profile.release-ci]\ninherits = \"pgo\"\n" }, { "path": "Cross.toml", "content": "# For Asahi linux\n[target.aarch64-unknown-linux-gnu.env]\npassthrough = [\"JEMALLOC_SYS_WITH_LG_PAGE=16\"]\n\n[target.aarch64-unknown-linux-musl.env]\npassthrough = [\"JEMALLOC_SYS_WITH_LG_PAGE=16\"]\n" }, { "path": "Dockerfile", "content": "ARG RUST_VERSION=slim\nFROM docker.io/library/rust:${RUST_VERSION} AS chef\n\nRUN cargo install cargo-chef --locked\nRUN apt-get update && apt-get install -y \\\n cmake \\\n && apt-get clean \\\n && rm -rf /var/lib/apt/lists/*\n\nWORKDIR /app\n\nFROM chef AS planner\nCOPY . .\nRUN cargo chef prepare --recipe-path recipe.json\n\nFROM chef AS builder\nCOPY --from=planner /app/recipe.json recipe.json\n\nRUN cargo chef cook --release --no-default-features --features rustls --recipe-path recipe.json\n\nCOPY . .\nRUN cargo build --release --no-default-features --features rustls --bin oha\nRUN strip /app/target/release/oha\n\nFROM registry.fedoraproject.org/fedora-minimal AS runtime\nUSER 65535\nCOPY --chown=65535:65535 --from=builder /app/target/release/oha /bin/\nENTRYPOINT [\"/bin/oha\"]\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2020 hatoo\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# oha (おはよう)\n\n[![GitHub Actions](https://github.com/hatoo/oha/workflows/CI/badge.svg)](https://github.com/hatoo/oha/actions?query=workflow%3ACI)\n[![Crates.io](https://img.shields.io/crates/v/oha.svg)](https://crates.io/crates/oha)\n[![Arch Linux](https://img.shields.io/archlinux/v/extra/x86_64/oha)](https://archlinux.org/packages/extra/x86_64/oha/)\n[![Homebrew](https://img.shields.io/homebrew/v/oha)](https://formulae.brew.sh/formula/oha)\n\n[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/hatoo)\n\noha is a tiny program that sends some load to a web application and show realtime tui inspired by [rakyll/hey](https://github.com/rakyll/hey).\n\nThis program is written in Rust and powered by [tokio](https://github.com/tokio-rs/tokio) and beautiful tui by [ratatui](https://github.com/ratatui-org/ratatui).\n\n![demo](demo.gif)\n\n# Installation\n\nThis program is built on stable Rust, with both `make` and `cmake` prerequisites to install via cargo.\n\n cargo install oha\n\nYou can optionally build oha against [native-tls](https://github.com/sfackler/rust-native-tls) instead of [rustls](https://github.com/rustls/rustls).\n\n cargo install --no-default-features --features native-tls oha\n\nYou can enable VSOCK support by enabling `vsock` feature.\n\n cargo install --features vsock oha\n\nYou can enable experimental HTTP3 support by enabling the `http3` feature. This uses the [H3](https://github.com/hyperium/h3) library by the developers of Hyper.\nIt will remain experimental as long as H3 is experimental. It currently depends on using `rustls` for TLS.\n\n## Download pre-built binary\n\nYou can download pre-built binary from [Release page](https://github.com/hatoo/oha/releases) for each version and from [Publish workflow](https://github.com/hatoo/oha/actions/workflows/release.yml) and [Publish PGO workflow](https://github.com/hatoo/oha/actions/workflows/release-pgo.yml) for each commit.\n\n## On Arch Linux\n\n pacman -S oha\n\n## On macOS (Homebrew)\n\n brew install oha\n\n## On Windows (winget)\n\n winget install hatoo.oha\n\n## On Debian ([Azlux's repository](http://packages.azlux.fr/))\n\n echo \"deb [signed-by=/usr/share/keyrings/azlux-archive-keyring.gpg] http://packages.azlux.fr/debian/ stable main\" | sudo tee /etc/apt/sources.list.d/azlux.list\n sudo wget -O /usr/share/keyrings/azlux-archive-keyring.gpg https://azlux.fr/repo.gpg\n apt update\n apt install oha\n\n## X-CMD (Linux, macOS, Windows WSL/GitBash)\n\nYou can install with [x-cmd](https://www.x-cmd.com).\n\n```sh\nx env use oha\n```\n\n## Containerized\n\n### Official Docker image\n\n[ghcr.io/hatoo/oha](https://github.com/hatoo/oha/pkgs/container/oha)\n\n### Build image locally\n\nYou can also build and create a container image including oha\n\n```sh\ndocker build -t hatoo/oha:latest .\n```\n\nThen you can use oha directly through the container\n\n```sh\ndocker run --rm -it --network=host hatoo/oha:latest https://example.com:3000\n```\n\n## Profile-Guided Optimization (PGO)\n\nYou can build `oha` with PGO by using the following commands:\n\n```sh\nbun run pgo.js\n```\n\nAnd the binary will be available at `target/[target-triple]/pgo/oha`.\n\n**Note**: Please keep in mind that in order to run the aforementioned command,\nyou need to have installed `cargo-pgo` cargo package.\n\nYou can install it via `cargo install cargo-pgo`.\n\n# Platform\n\n- Linux - Tested on Ubuntu 18.04 gnome-terminal\n- Windows 10 - Tested on Windows Powershell\n- MacOS - Tested on iTerm2\n\n# Usage\n\n`-q` option works different from [rakyll/hey](https://github.com/rakyll/hey). It's set overall query per second instead of for each workers.\n\n```sh\nOhayou(おはよう), HTTP load generator, inspired by rakyll/hey with tui animation.\n\nUsage: oha [OPTIONS] \n\nArguments:\n Target URL or file with multiple URLs.\n\nOptions:\n -n \n Number of requests to run. Accepts plain numbers or suffixes: k = 1,000, m = 1,000,000 (e.g. 10k, 1m). [default: 200]\n -c \n Number of connections to run concurrently. You may should increase limit to number of open files for larger `-c`. [default: 50]\n -p \n Number of parallel requests to send on HTTP/2. `oha` will run c * p concurrent workers in total. [default: 1]\n -z \n Duration of application to send requests.\n On HTTP/1, When the duration is reached, ongoing requests are aborted and counted as \"aborted due to deadline\"\n You can change this behavior with `-w` option.\n Currently, on HTTP/2, When the duration is reached, ongoing requests are waited. `-w` option is ignored.\n Examples: -z 10s -z 3m.\n -w, --wait-ongoing-requests-after-deadline\n When the duration is reached, ongoing requests are waited\n -q \n Rate limit for all, in queries per second (QPS)\n --burst-delay \n Introduce delay between a predefined number of requests.\n Note: If qps is specified, burst will be ignored\n --burst-rate \n Rates of requests for burst. Default is 1\n Note: If qps is specified, burst will be ignored\n --rand-regex-url\n Generate URL by rand_regex crate but dot is disabled for each query e.g. http://127.0.0.1/[a-z][a-z][0-9]. Currently dynamic scheme, host and port with keep-alive do not work well. See https://docs.rs/rand_regex/latest/rand_regex/struct.Regex.html for details of syntax.\n --urls-from-file\n Read the URLs to query from a file\n --max-repeat \n A parameter for the '--rand-regex-url'. The max_repeat parameter gives the maximum extra repeat counts the x*, x+ and x{n,} operators will become. [default: 4]\n --dump-urls \n Dump target Urls times to debug --rand-regex-url\n --latency-correction\n Correct latency to avoid coordinated omission problem. It's ignored if -q is not set.\n --no-tui\n No realtime tui\n --fps \n Frame per second for tui. [default: 16]\n -m, --method \n HTTP method [default: GET]\n -H \n Custom HTTP header. Examples: -H \"foo: bar\"\n --proxy-header \n Custom Proxy HTTP header. Examples: --proxy-header \"foo: bar\"\n -t \n Timeout for each request. Default to infinite.\n --connect-timeout \n Timeout for establishing a new connection. Default to 5s. [default: 5s]\n -A \n HTTP Accept Header.\n -d \n HTTP request body.\n -D \n HTTP request body from file.\n -Z \n HTTP request body from file line by line.\n -F, --form
\n Specify HTTP multipart POST data (curl compatible). Examples: -F 'name=value' -F 'file=@path/to/file'\n -T \n Content-Type.\n -a \n Basic authentication (username:password), or AWS credentials (access_key:secret_key)\n --aws-session \n AWS session token\n --aws-sigv4 \n AWS SigV4 signing params (format: aws:amz:region:service)\n -x \n HTTP proxy\n --proxy-http-version \n HTTP version to connect to proxy. Available values 0.9, 1.0, 1.1, 2.\n --proxy-http2\n Use HTTP/2 to connect to proxy. Shorthand for --proxy-http-version=2\n --http-version \n HTTP version. Available values 0.9, 1.0, 1.1, 2, 3\n --http2\n Use HTTP/2. Shorthand for --http-version=2\n --host \n HTTP Host header\n --disable-compression\n Disable compression.\n -r, --redirect \n Limit for number of Redirect. Set 0 for no redirection. Redirection isn't supported for HTTP/2. [default: 10]\n --disable-keepalive\n Disable keep-alive, prevents re-use of TCP connections between different HTTP requests. This isn't supported for HTTP/2.\n --no-pre-lookup\n *Not* perform a DNS lookup at beginning to cache it\n --ipv6\n Lookup only ipv6.\n --ipv4\n Lookup only ipv4.\n --cacert \n (TLS) Use the specified certificate file to verify the peer. Native certificate store is used even if this argument is specified.\n --cert \n (TLS) Use the specified client certificate file. --key must be also specified\n --key \n (TLS) Use the specified client key file. --cert must be also specified\n --insecure\n Accept invalid certs.\n --connect-to \n Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'\n Note: if used several times for the same host:port:target_host:target_port, a random choice is made\n --no-color\n Disable the color scheme. [env: NO_COLOR=]\n --unix-socket \n Connect to a unix socket instead of the domain in the URL. Only for non-HTTPS URLs.\n --stats-success-breakdown\n Include a response status code successful or not successful breakdown for the time histogram and distribution statistics\n --db-url \n Write succeeded requests to sqlite database url E.G test.db\n --debug\n Perform a single request and dump the request and response\n -o, --output \n Output file to write the results to. If not specified, results are written to stdout.\n --output-format \n Output format [default: text] [possible values: text, json, csv, quiet]\n -u, --time-unit \n Time unit to be used. If not specified, the time unit is determined automatically. This option affects only text format. [possible values: ns, us, ms, s, m, h]\n -h, --help\n Print help\n -V, --version\n Print version\n```\n\n# Performance\n\n`oha` uses faster implementation when `--no-tui` option is set and both `-q` and `--burst-delay` are not set because it can avoid overhead to gather data realtime.\n\n# Output\n\nBy default `oha` outputs a text summary of the results.\n\n`oha` prints JSON summary output when `--output-format json` option is set.\nThe schema of JSON output is defined in [schema.json](./schema.json).\n\nWhen `--output-format csv` is used result of each request is printed as a line of comma separated values.\n\n# Tips\n\n## Stress test in more realistic condition\n\n`oha` uses default options inherited from [rakyll/hey](https://github.com/rakyll/hey) but you may need to change options to stress test in more realistic condition.\n\nI suggest to run `oha` with following options.\n\n```sh\noha <-z or -n> -c -q --latency-correction --disable-keepalive \n```\n\n- --disable-keepalive\n\n In real, user doesn't query same URL using [Keep-Alive](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Keep-Alive). You may want to run without `Keep-Alive`.\n- --latency-correction\n\n You can avoid `Coordinated Omission Problem` by using `--latency-correction`.\n\n## Burst feature\n\nYou can use `--burst-delay` along with `--burst-rate` option to introduce delay between a defined number of requests.\n\n```sh\noha -n 10 --burst-delay 2s --burst-rate 4\n```\n\nIn this particular scenario, every 2 seconds, 4 requests will be processed, and after 6s the total of 10 requests will be processed.\n*NOTE: If you don't set `--burst-rate` option, the amount is default to 1*\n\n## Dynamic url feature\n\nYou can use `--rand-regex-url` option to generate random url for each connection.\n\n```sh\noha --rand-regex-url http://127.0.0.1/[a-z][a-z][0-9]\n```\n\nEach Urls are generated by [rand_regex](https://github.com/kennytm/rand_regex) crate but regex's dot is disabled since it's not useful for this purpose and it's very inconvenient if url's dots are interpreted as regex's dot.\n\nOptionally you can set `--max-repeat` option to limit max repeat count for each regex. e.g http://127.0.0.1/[a-z]* with `--max-repeat 4` will generate url like http://127.0.0.1/[a-z]{0,4}\n\nCurrently dynamic scheme, host and port with keep-alive are not works well.\n\n## URLs from file feature\n\nYou can use `--urls-from-file` to read the target URLs from a file. Each line of this file needs to contain one valid URL as in the example below.\n\n```\nhttp://domain.tld/foo/bar\nhttp://domain.tld/assets/vendors-node_modules_highlight_js_lib_index_js-node_modules_tanstack_react-query_build_modern-3fdf40-591fb51c8a6e.js\nhttp://domain.tld/images/test.png\nhttp://domain.tld/foo/bar?q=test\nhttp://domain.tld/foo\n```\n\nSuch a file can for example be created from an access log to generate a more realistic load distribution over the different pages of a server. \n\nWhen this type of URL specification is used, every request goes to a random URL given in the file.\n\n# Contribution\n\nFeel free to help us!\n\nHere are some areas which need improving.\n\n- Write tests\n- Improve tui design.\n - Show more information?\n- Improve speed\n - I'm new to tokio. I think there are some space to optimize query scheduling.\n" }, { "path": "pgo/server/Cargo.toml", "content": "[package]\nname = \"server\"\nversion = \"0.1.0\"\nedition = \"2021\"\n\n# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html\n\n[dependencies]\naxum = \"0.8.1\"\ntokio = { version = \"1\", features = [\"full\"] }\n" }, { "path": "pgo/server/src/main.rs", "content": "use std::net::SocketAddr;\nuse tokio::net::TcpListener;\n\nuse axum::{routing::get, Router};\n\n#[tokio::main]\nasync fn main() {\n // build our application with a route\n let app = Router::new()\n // `GET /` goes to `root`\n .route(\"/\", get(root));\n\n // run our app with hyper\n // `axum::Server` is a re-export of `hyper::Server`\n let addr = SocketAddr::from(([127, 0, 0, 1], 8888));\n let listener = TcpListener::bind(&addr).await.unwrap();\n axum::serve(listener, app).await.unwrap();\n}\n\nasync fn root() -> &'static str {\n \"Hello, World!\"\n}\n" }, { "path": "pgo.js", "content": "import { $ } from \"bun\";\n\nlet additional = [];\n\nif (Bun.argv.length >= 3) {\n additional = Bun.argv.slice(2);\n}\n\nlet server = null;\n\ntry {\n server = Bun.spawn(['cargo', 'run', '--release', '--manifest-path', 'pgo/server/Cargo.toml']);\n await $`cargo pgo run -- --profile pgo ${additional} -- -z 3m -c 900 --no-tui http://localhost:8888`;\n await $`cargo pgo optimize build -- --profile pgo ${additional}`\n} finally {\n if (server !== null) {\n server.kill();\n }\n}\n" }, { "path": "schema.json", "content": "{\n \"$schema\": \"https://json-schema.org/draft/2020-12/schema\",\n \"description\": \"JSON schema for the output of the `oha --output-format json`\",\n \"type\": \"object\",\n \"properties\": {\n \"summary\": {\n \"description\": \"Important statistics\",\n \"type\": \"object\",\n \"properties\": {\n \"successRate\": {\n \"description\": \"The number of success requests / All requests which isn't includes deadline\",\n \"type\": \"number\"\n },\n \"total\": {\n \"description\": \"Total duration in seconds\",\n \"type\": \"number\"\n },\n \"slowest\": {\n \"description\": \"The slowest request duration in seconds\",\n \"type\": \"number\"\n },\n \"fastest\": {\n \"description\": \"The fastest request duration in seconds\",\n \"type\": \"number\"\n },\n \"average\": {\n \"description\": \"The average request duration in seconds\",\n \"type\": \"number\"\n },\n \"requestsPerSec\": {\n \"description\": \"The number of requests per second\",\n \"type\": \"number\"\n },\n \"totalData\": {\n \"description\": \"Total data of HTTP bodies in bytes\",\n \"type\": \"integer\"\n },\n \"sizePerRequest\": {\n \"description\": \"The average size of HTTP bodies in bytes\",\n \"type\": \"integer\"\n },\n \"sizePerSec\": {\n \"description\": \"The average size of HTTP bodies per second in bytes\",\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"successRate\",\n \"total\",\n \"slowest\",\n \"fastest\",\n \"average\",\n \"requestsPerSec\",\n \"totalData\",\n \"sizePerRequest\",\n \"sizePerSec\"\n ]\n },\n \"responseTimeHistogram\": {\n \"description\": \"The histogram of response time in seconds. The key is the response time in seconds and the value is the number of requests\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n },\n \"latencyPercentiles\": {\n \"description\": \"The latency percentiles in seconds\",\n \"type\": \"object\",\n \"properties\": {\n \"p10\": {\n \"type\": \"number\"\n },\n \"p25\": {\n \"type\": \"number\"\n },\n \"p50\": {\n \"type\": \"number\"\n },\n \"p75\": {\n \"type\": \"number\"\n },\n \"p90\": {\n \"type\": \"number\"\n },\n \"p95\": {\n \"type\": \"number\"\n },\n \"p99\": {\n \"type\": \"number\"\n },\n \"p99.9\": {\n \"type\": \"number\"\n },\n \"p99.99\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"p10\",\n \"p25\",\n \"p50\",\n \"p75\",\n \"p90\",\n \"p95\",\n \"p99\",\n \"p99.9\",\n \"p99.99\"\n ]\n },\n \"firstByteHistogram\": {\n \"description\": \"The histogram of first byte time in seconds. The key is the first byte time in seconds and the value is the number of requests\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n },\n \"firstBytePercentiles\": {\n \"description\": \"The first byte percentiles in seconds\",\n \"type\": \"object\",\n \"properties\": {\n \"p10\": {\n \"type\": \"number\"\n },\n \"p25\": {\n \"type\": \"number\"\n },\n \"p50\": {\n \"type\": \"number\"\n },\n \"p75\": {\n \"type\": \"number\"\n },\n \"p90\": {\n \"type\": \"number\"\n },\n \"p95\": {\n \"type\": \"number\"\n },\n \"p99\": {\n \"type\": \"number\"\n },\n \"p99.9\": {\n \"type\": \"number\"\n },\n \"p99.99\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"p10\",\n \"p25\",\n \"p50\",\n \"p75\",\n \"p90\",\n \"p95\",\n \"p99\",\n \"p99.9\",\n \"p99.99\"\n ]\n },\n \"responseTimeHistogramSuccessful\": {\n \"description\": \"Only present if `--stats-success-breakdown` argument is passed. The histogram of response time in seconds for successful requests. The key is the response time in seconds and the value is the number of requests\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n },\n \"latencyPercentileSuccessful\": {\n \"description\": \"Only present if `--stats-success-breakdown` argument is passed. The latency percentiles in seconds for successful requests\",\n \"type\": \"object\",\n \"properties\": {\n \"p10\": {\n \"type\": \"number\"\n },\n \"p25\": {\n \"type\": \"number\"\n },\n \"p50\": {\n \"type\": \"number\"\n },\n \"p75\": {\n \"type\": \"number\"\n },\n \"p90\": {\n \"type\": \"number\"\n },\n \"p95\": {\n \"type\": \"number\"\n },\n \"p99\": {\n \"type\": \"number\"\n },\n \"p99.9\": {\n \"type\": \"number\"\n },\n \"p99.99\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"p10\",\n \"p25\",\n \"p50\",\n \"p75\",\n \"p90\",\n \"p95\",\n \"p99\",\n \"p99.9\",\n \"p99.99\"\n ]\n },\n \"responseTimeHistogramNotSuccessful\": {\n \"description\": \"Only present if `--stats-success-breakdown` argument is passed. The histogram of response time in seconds for not successful requests. The key is the response time in seconds and the value is the number of requests\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n },\n \"latencyPercentileNotSuccessful\": {\n \"description\": \"Only present if `--stats-success-breakdown` argument is passed. The latency percentiles in seconds for not successful requests\",\n \"type\": \"object\",\n \"properties\": {\n \"p10\": {\n \"type\": \"number\"\n },\n \"p25\": {\n \"type\": \"number\"\n },\n \"p50\": {\n \"type\": \"number\"\n },\n \"p75\": {\n \"type\": \"number\"\n },\n \"p90\": {\n \"type\": \"number\"\n },\n \"p95\": {\n \"type\": \"number\"\n },\n \"p99\": {\n \"type\": \"number\"\n },\n \"p99.9\": {\n \"type\": \"number\"\n },\n \"p99.99\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"p10\",\n \"p25\",\n \"p50\",\n \"p75\",\n \"p90\",\n \"p95\",\n \"p99\",\n \"p99.9\",\n \"p99.99\"\n ]\n },\n \"rps\": {\n \"description\": \"The statistics for requests per second. Note: the way of calculating rps over time isn't obvious, see source code for details.\",\n \"type\": \"object\",\n \"properties\": {\n \"mean\": {\n \"type\": \"number\"\n },\n \"stddev\": {\n \"type\": [\n \"number\",\n \"null\"\n ]\n },\n \"max\": {\n \"type\": \"number\"\n },\n \"min\": {\n \"type\": \"number\"\n },\n \"percentiles\": {\n \"type\": \"object\",\n \"properties\": {\n \"p10\": {\n \"type\": \"number\"\n },\n \"p25\": {\n \"type\": \"number\"\n },\n \"p50\": {\n \"type\": \"number\"\n },\n \"p75\": {\n \"type\": \"number\"\n },\n \"p90\": {\n \"type\": \"number\"\n },\n \"p95\": {\n \"type\": \"number\"\n },\n \"p99\": {\n \"type\": \"number\"\n },\n \"p99.9\": {\n \"type\": \"number\"\n },\n \"p99.99\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"p10\",\n \"p25\",\n \"p50\",\n \"p75\",\n \"p90\",\n \"p95\",\n \"p99\",\n \"p99.9\",\n \"p99.99\"\n ]\n }\n },\n \"required\": [\n \"mean\",\n \"stddev\",\n \"max\",\n \"min\",\n \"percentiles\"\n ]\n },\n \"details\": {\n \"description\": \"The details of connection time. Note: `oha` uses keep-alive connections in default. So, the connection time may added only for the first request.\",\n \"type\": \"object\",\n \"properties\": {\n \"DNSDialup\": {\n \"description\": \"The time of DNS resolution + TCP handshake in seconds\",\n \"type\": \"object\",\n \"properties\": {\n \"average\": {\n \"type\": \"number\"\n },\n \"fastest\": {\n \"type\": \"number\"\n },\n \"slowest\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"average\",\n \"fastest\",\n \"slowest\"\n ]\n },\n \"DNSLookup\": {\n \"description\": \"The time of DNS resolution in seconds\",\n \"type\": \"object\",\n \"properties\": {\n \"average\": {\n \"type\": \"number\"\n },\n \"fastest\": {\n \"type\": \"number\"\n },\n \"slowest\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"average\",\n \"fastest\",\n \"slowest\"\n ]\n },\n \"firstByte\": {\n \"description\": \"The time to first byte in seconds\",\n \"type\": \"object\",\n \"properties\": {\n \"average\": {\n \"type\": \"number\"\n },\n \"fastest\": {\n \"type\": \"number\"\n },\n \"slowest\": {\n \"type\": \"number\"\n }\n },\n \"required\": [\n \"average\",\n \"fastest\",\n \"slowest\"\n ]\n }\n },\n \"required\": [\n \"DNSDialup\",\n \"DNSLookup\",\n \"firstByte\"\n ]\n },\n \"statusCodeDistribution\": {\n \"description\": \"The distribution of status codes. The key is the status code and the value is the number of requests\",\n \"type\": \"object\",\n \"propertyNames\": {\n \"type\": \"string\",\n \"pattern\": \"^[0-9]+$\"\n },\n \"additionalProperties\": {\n \"type\": \"integer\"\n },\n \"minProperties\": 1\n },\n \"errorDistribution\": {\n \"description\": \"The distribution of errors. The key is the error message and the value is the number of errors. Note: the error message is from internal libraries so the detail may change in future.\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n }\n },\n \"required\": [\n \"summary\",\n \"responseTimeHistogram\",\n \"latencyPercentiles\",\n \"firstByteHistogram\",\n \"firstBytePercentiles\",\n \"rps\",\n \"details\",\n \"statusCodeDistribution\",\n \"errorDistribution\"\n ]\n}" }, { "path": "src/aws_auth.rs", "content": "use anyhow::Result;\n\nuse bytes::Bytes;\nuse hyper::{\n HeaderMap,\n header::{self, HeaderName},\n};\nuse thiserror::Error;\nuse url::Url;\n\npub struct AwsSignatureConfig {\n pub access_key: String,\n pub secret_key: String,\n pub session_token: Option,\n pub service: String,\n pub region: String,\n}\n\n#[derive(Error, Debug)]\npub enum AwsSignatureError {\n #[error(\"URL must contain a host {0}\")]\n NoHost(Url),\n #[error(\"Invalid host header name {0}\")]\n InvalidHost(String),\n #[error(\"Invalid authorization header name {0}\")]\n InvalidAuthorization(String),\n}\n\n// Initialize unsignable headers as a static constant\nstatic UNSIGNABLE_HEADERS: [HeaderName; 8] = [\n header::ACCEPT,\n header::ACCEPT_ENCODING,\n header::USER_AGENT,\n header::EXPECT,\n header::RANGE,\n header::CONNECTION,\n HeaderName::from_static(\"presigned-expires\"),\n HeaderName::from_static(\"x-amzn-trace-id\"),\n];\n\nimpl AwsSignatureConfig {\n pub fn sign_request(\n &self,\n method: &str,\n headers: &mut HeaderMap,\n url: &Url,\n body: &Bytes,\n ) -> Result<(), AwsSignatureError> {\n let datetime = chrono::Utc::now();\n\n let header_amz_date = datetime\n .format(\"%Y%m%dT%H%M%SZ\")\n .to_string()\n .parse()\n .unwrap();\n\n if !headers.contains_key(header::HOST) {\n let host = url\n .host_str()\n .ok_or_else(|| AwsSignatureError::NoHost(url.clone()))?;\n headers.insert(\n header::HOST,\n host.parse()\n .map_err(|_| AwsSignatureError::InvalidHost(host.to_string()))?,\n );\n }\n headers.insert(\"x-amz-date\", header_amz_date);\n\n if let Some(session_token) = &self.session_token {\n headers.insert(\"x-amz-security-token\", session_token.parse().unwrap());\n }\n\n headers.remove(header::AUTHORIZATION);\n\n //remove and store headers in a vec from unsignable_headers\n let removed_headers: Vec<(header::HeaderName, header::HeaderValue)> = UNSIGNABLE_HEADERS\n .iter()\n .filter_map(|k| headers.remove(k).map(|v| (k.clone(), v)))\n .collect();\n\n headers.insert(\n header::CONTENT_LENGTH,\n body.len().to_string().parse().unwrap(),\n );\n\n let aws_sign = aws_sign_v4::AwsSign::new(\n method,\n url.as_str(),\n &datetime,\n headers,\n &self.region,\n &self.access_key,\n &self.secret_key,\n &self.service,\n body,\n );\n\n let signature = aws_sign.sign();\n\n //insert headers\n for (key, value) in removed_headers {\n headers.insert(key, value);\n }\n\n headers.insert(\n header::AUTHORIZATION,\n signature\n .parse()\n .map_err(|_| AwsSignatureError::InvalidAuthorization(signature.to_string()))?,\n );\n\n Ok(())\n }\n\n pub fn new(\n access_key: &str,\n secret_key: &str,\n signing_params: &str,\n session_token: Option,\n ) -> Result {\n let parts: Vec<&str> = signing_params\n .strip_prefix(\"aws:amz:\")\n .unwrap_or_default()\n .split(':')\n .collect();\n if parts.len() != 2 {\n anyhow::bail!(\"Invalid AWS signing params format. Expected aws:amz:region:service\");\n }\n\n Ok(Self {\n access_key: access_key.into(),\n secret_key: secret_key.into(),\n session_token,\n region: parts[0].to_string(),\n service: parts[1].to_string(),\n })\n }\n}\n" }, { "path": "src/cli.rs", "content": "use hyper::http::header::{HeaderName, HeaderValue};\nuse std::str::FromStr;\n\npub fn parse_header(s: &str) -> Result<(HeaderName, HeaderValue), anyhow::Error> {\n let header = s.splitn(2, ':').collect::>();\n anyhow::ensure!(header.len() == 2, anyhow::anyhow!(\"Parse header\"));\n let name = HeaderName::from_str(header[0])?;\n let value = HeaderValue::from_str(header[1].trim_start_matches(' '))?;\n Ok::<(HeaderName, HeaderValue), anyhow::Error>((name, value))\n}\n\npub fn parse_n_requests(s: &str) -> Result {\n let s = s.trim().to_lowercase();\n if let Some(num) = s.strip_suffix('k') {\n num.parse::()\n .map(|n| (n * 1000_f64) as usize)\n .map_err(|e| e.to_string())\n } else if let Some(num) = s.strip_suffix('m') {\n num.parse::()\n .map(|n| (n * 1_000_000_f64) as usize)\n .map_err(|e| e.to_string())\n } else {\n s.parse::().map_err(|e| e.to_string())\n }\n}\n\n/// An entry specified by `connect-to` to override DNS resolution and default\n/// port numbers. For example, `example.org:80:localhost:5000` will connect to\n/// `localhost:5000` whenever `http://example.org` is requested.\n#[derive(Clone, Debug)]\npub struct ConnectToEntry {\n pub requested_host: String,\n pub requested_port: u16,\n pub target_host: String,\n pub target_port: u16,\n}\n\nimpl FromStr for ConnectToEntry {\n type Err = String;\n\n fn from_str(s: &str) -> Result {\n let expected_syntax: &str = \"syntax for --connect-to is host:port:target_host:target_port\";\n\n let (s, target_port) = s.rsplit_once(':').ok_or(expected_syntax)?;\n let (s, target_host) = if s.ends_with(']') {\n // ipv6\n let i = s.rfind(\":[\").ok_or(expected_syntax)?;\n (&s[..i], &s[i + 1..])\n } else {\n s.rsplit_once(':').ok_or(expected_syntax)?\n };\n let (requested_host, requested_port) = s.rsplit_once(':').ok_or(expected_syntax)?;\n\n Ok(ConnectToEntry {\n requested_host: requested_host.into(),\n requested_port: requested_port.parse().map_err(|err| {\n format!(\"requested port must be an u16, but got {requested_port}: {err}\")\n })?,\n target_host: target_host.into(),\n target_port: target_port.parse().map_err(|err| {\n format!(\"target port must be an u16, but got {target_port}: {err}\")\n })?,\n })\n }\n}\n\n#[cfg(feature = \"vsock\")]\npub fn parse_vsock_addr(s: &str) -> Result {\n let (cid, port) = s\n .split_once(':')\n .ok_or(\"syntax for --vsock-addr is cid:port\")?;\n Ok(tokio_vsock::VsockAddr::new(\n cid.parse()\n .map_err(|err| format!(\"cid must be a u32, but got {cid}: {err}\"))?,\n port.parse()\n .map_err(|err| format!(\"port must be a u32, but got {port}: {err}\"))?,\n ))\n}\n" }, { "path": "src/client.rs", "content": "use bytes::Bytes;\n#[cfg(test)]\nuse hickory_resolver::config::{ResolverConfig, ResolverOpts};\nuse http_body_util::{BodyExt, Full};\nuse hyper::{Method, Request, http};\nuse hyper_util::rt::{TokioExecutor, TokioIo};\nuse rand::{prelude::*, rng};\nuse std::{\n borrow::Cow,\n io::Write,\n sync::{\n Arc,\n atomic::{AtomicBool, Ordering::Relaxed},\n },\n time::Instant,\n};\nuse thiserror::Error;\nuse tokio::{\n io::{AsyncRead, AsyncWrite},\n net::TcpStream,\n};\nuse url::{ParseError, Url};\n\nuse crate::{\n ConnectToEntry,\n pcg64si::Pcg64Si,\n request_generator::{RequestGenerationError, RequestGenerator},\n url_generator::UrlGeneratorError,\n};\n\n#[cfg(feature = \"http3\")]\nuse crate::client_h3::send_debug_request_http3;\n\ntype SendRequestHttp1 = hyper::client::conn::http1::SendRequest>;\ntype SendRequestHttp2 = hyper::client::conn::http2::SendRequest>;\n\nfn format_host_port(host: &str, port: u16) -> String {\n if host.contains(':') && !(host.starts_with('[') && host.ends_with(']')) {\n format!(\"[{host}]:{port}\")\n } else {\n format!(\"{host}:{port}\")\n }\n}\n\n#[derive(Debug, Clone, Copy)]\npub struct ConnectionTime {\n pub dns_lookup: std::time::Duration,\n pub dialup: std::time::Duration,\n}\n\n#[derive(Debug, Clone)]\n/// a result for a request\npub struct RequestResult {\n pub rng: Pcg64Si,\n // When the query should started\n pub start_latency_correction: Option,\n /// When the query started\n pub start: std::time::Instant,\n /// DNS + dialup\n /// None when reuse connection\n pub connection_time: Option,\n /// First body byte received\n pub first_byte: Option,\n /// When the query ends\n pub end: std::time::Instant,\n /// HTTP status\n pub status: http::StatusCode,\n /// Length of body\n pub len_bytes: usize,\n}\n\nimpl RequestResult {\n /// Duration the request takes.\n pub fn duration(&self) -> std::time::Duration {\n self.end - self.start_latency_correction.unwrap_or(self.start)\n }\n}\n\n// encapsulates the HTTP generation of the work type. Used internally only for conditional logic.\n#[derive(Debug, Clone, Copy, PartialEq)]\nenum HttpWorkType {\n H1,\n H2,\n #[cfg(feature = \"http3\")]\n H3,\n}\n\npub struct Dns {\n pub connect_to: Vec,\n pub resolver:\n hickory_resolver::Resolver,\n}\n\nimpl Dns {\n fn select_connect_to<'a, R: Rng>(\n &'a self,\n host: &str,\n port: u16,\n rng: &mut R,\n ) -> Option<&'a ConnectToEntry> {\n self.connect_to\n .iter()\n .filter(|entry| entry.requested_port == port && entry.requested_host == host)\n .collect::>()\n .choose(rng)\n .copied()\n }\n\n /// Perform a DNS lookup for a given url and returns (ip_addr, port)\n async fn lookup(\n &self,\n url: &Url,\n rng: &mut R,\n ) -> Result<(std::net::IpAddr, u16), ClientError> {\n let host = url.host_str().ok_or(ClientError::HostNotFound)?;\n let port = url\n .port_or_known_default()\n .ok_or(ClientError::PortNotFound)?;\n\n // Try to find an override (passed via `--connect-to`) that applies to this (host, port),\n // choosing one randomly if several match.\n let (host, port) = if let Some(entry) = self.select_connect_to(host, port, rng) {\n (entry.target_host.as_str(), entry.target_port)\n } else {\n (host, port)\n };\n\n let host = if host.starts_with('[') && host.ends_with(']') {\n // host is [ipv6] format\n // remove first [ and last ]\n &host[1..host.len() - 1]\n } else {\n host\n };\n\n // Perform actual DNS lookup, either on the original (host, port), or\n // on the (host, port) specified with `--connect-to`.\n let addrs = self\n .resolver\n .lookup_ip(host)\n .await\n .map_err(Box::new)?\n .iter()\n .collect::>();\n\n let addr = *addrs.choose(rng).ok_or(ClientError::DNSNoRecord)?;\n\n Ok((addr, port))\n }\n}\n\n#[derive(Error, Debug)]\npub enum ClientError {\n #[error(\"failed to get port from URL\")]\n PortNotFound,\n #[error(\"failed to get host from URL\")]\n HostNotFound,\n #[error(\"No record returned from DNS\")]\n DNSNoRecord,\n #[error(\"Redirection limit has reached\")]\n TooManyRedirect,\n #[error(transparent)]\n // Use Box here because ResolveError is big.\n Resolve(#[from] Box),\n\n #[cfg(feature = \"native-tls\")]\n #[error(transparent)]\n NativeTls(#[from] native_tls::Error),\n\n #[cfg(feature = \"rustls\")]\n #[error(transparent)]\n Rustls(#[from] rustls::Error),\n\n #[cfg(feature = \"rustls\")]\n #[error(transparent)]\n InvalidDnsName(#[from] rustls_pki_types::InvalidDnsNameError),\n\n #[error(transparent)]\n Io(#[from] std::io::Error),\n #[error(transparent)]\n Http(#[from] http::Error),\n #[error(transparent)]\n Hyper(#[from] hyper::Error),\n #[error(transparent)]\n InvalidUriParts(#[from] http::uri::InvalidUriParts),\n #[error(transparent)]\n InvalidHeaderValue(#[from] http::header::InvalidHeaderValue),\n #[error(\"Failed to get header from builder\")]\n GetHeaderFromBuilder,\n #[error(transparent)]\n HeaderToStr(#[from] http::header::ToStrError),\n #[error(transparent)]\n InvalidUri(#[from] http::uri::InvalidUri),\n #[error(\"timeout\")]\n Timeout,\n #[error(\"aborted due to deadline\")]\n Deadline,\n #[error(transparent)]\n UrlGenerator(#[from] UrlGeneratorError),\n #[error(transparent)]\n UrlParse(#[from] ParseError),\n #[error(\"Request generation error: {0}\")]\n RequestGeneration(#[from] RequestGenerationError),\n #[cfg(feature = \"http3\")]\n #[error(transparent)]\n Http3(#[from] crate::client_h3::Http3Error),\n}\n\npub struct Client {\n pub request_generator: RequestGenerator,\n pub proxy_http_version: http::Version,\n pub proxy_headers: http::header::HeaderMap,\n pub dns: Dns,\n pub timeout: Option,\n pub connect_timeout: std::time::Duration,\n pub redirect_limit: usize,\n pub disable_keepalive: bool,\n pub proxy_url: Option,\n #[cfg(unix)]\n pub unix_socket: Option,\n #[cfg(feature = \"vsock\")]\n pub vsock_addr: Option,\n #[cfg(feature = \"rustls\")]\n pub rustls_configs: crate::tls_config::RuslsConfigs,\n #[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\n pub native_tls_connectors: crate::tls_config::NativeTlsConnectors,\n}\n\n#[cfg(test)]\nimpl Default for Client {\n fn default() -> Self {\n use crate::request_generator::BodyGenerator;\n\n let (resolver_config, resolver_opts) = crate::system_resolv_conf()\n .unwrap_or_else(|_| (ResolverConfig::default(), ResolverOpts::default()));\n let resolver = hickory_resolver::Resolver::builder_with_config(\n resolver_config,\n hickory_resolver::name_server::TokioConnectionProvider::default(),\n )\n .with_options(resolver_opts)\n .build();\n\n Self {\n request_generator: RequestGenerator {\n url_generator: crate::url_generator::UrlGenerator::new_static(\n \"http://example.com\".parse().unwrap(),\n ),\n https: false,\n http_proxy: None,\n method: http::Method::GET,\n version: http::Version::HTTP_11,\n headers: http::header::HeaderMap::new(),\n body_generator: BodyGenerator::Static(Bytes::new()),\n aws_config: None,\n },\n proxy_http_version: http::Version::HTTP_11,\n proxy_headers: http::header::HeaderMap::new(),\n dns: Dns {\n resolver,\n connect_to: Vec::new(),\n },\n timeout: None,\n connect_timeout: std::time::Duration::from_secs(5),\n redirect_limit: 0,\n disable_keepalive: false,\n proxy_url: None,\n #[cfg(unix)]\n unix_socket: None,\n #[cfg(feature = \"vsock\")]\n vsock_addr: None,\n #[cfg(feature = \"rustls\")]\n rustls_configs: crate::tls_config::RuslsConfigs::new(false, None, None),\n #[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\n native_tls_connectors: crate::tls_config::NativeTlsConnectors::new(false, None, None),\n }\n }\n}\n\nstruct ClientStateHttp1 {\n rng: Pcg64Si,\n send_request: Option,\n}\n\nimpl Default for ClientStateHttp1 {\n fn default() -> Self {\n Self {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: None,\n }\n }\n}\n\nstruct ClientStateHttp2 {\n rng: Pcg64Si,\n send_request: SendRequestHttp2,\n}\n\npub enum QueryLimit {\n Qps(f64),\n Burst(std::time::Duration, usize),\n}\n\n// To avoid dynamic dispatch\n// I'm not sure how much this is effective\npub(crate) enum Stream {\n Tcp(TcpStream),\n #[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\n Tls(tokio_native_tls::TlsStream),\n #[cfg(feature = \"rustls\")]\n // Box for large variant\n Tls(Box>),\n #[cfg(unix)]\n Unix(tokio::net::UnixStream),\n #[cfg(feature = \"vsock\")]\n Vsock(tokio_vsock::VsockStream),\n #[cfg(feature = \"http3\")]\n Quic(quinn::Connection),\n}\n\nimpl Stream {\n async fn handshake_http1(self, with_upgrade: bool) -> Result {\n match self {\n Stream::Tcp(stream) => {\n let (send_request, conn) =\n hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;\n if with_upgrade {\n tokio::spawn(conn.with_upgrades());\n } else {\n tokio::spawn(conn);\n }\n Ok(send_request)\n }\n Stream::Tls(stream) => {\n let (send_request, conn) =\n hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;\n if with_upgrade {\n tokio::spawn(conn.with_upgrades());\n } else {\n tokio::spawn(conn);\n }\n Ok(send_request)\n }\n #[cfg(unix)]\n Stream::Unix(stream) => {\n let (send_request, conn) =\n hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;\n if with_upgrade {\n tokio::spawn(conn.with_upgrades());\n } else {\n tokio::spawn(conn);\n }\n Ok(send_request)\n }\n #[cfg(feature = \"vsock\")]\n Stream::Vsock(stream) => {\n let (send_request, conn) =\n hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;\n if with_upgrade {\n tokio::spawn(conn.with_upgrades());\n } else {\n tokio::spawn(conn);\n }\n Ok(send_request)\n }\n #[cfg(feature = \"http3\")]\n Stream::Quic(_) => {\n panic!(\"quic is not supported in http1\")\n }\n }\n }\n async fn handshake_http2(self) -> Result {\n let mut builder = hyper::client::conn::http2::Builder::new(TokioExecutor::new());\n builder\n // from nghttp2's default\n .initial_stream_window_size((1 << 30) - 1)\n .initial_connection_window_size((1 << 30) - 1);\n\n match self {\n Stream::Tcp(stream) => {\n let (send_request, conn) = builder.handshake(TokioIo::new(stream)).await?;\n tokio::spawn(conn);\n Ok(send_request)\n }\n Stream::Tls(stream) => {\n let (send_request, conn) = builder.handshake(TokioIo::new(stream)).await?;\n tokio::spawn(conn);\n Ok(send_request)\n }\n #[cfg(unix)]\n Stream::Unix(stream) => {\n let (send_request, conn) = builder.handshake(TokioIo::new(stream)).await?;\n tokio::spawn(conn);\n Ok(send_request)\n }\n #[cfg(feature = \"vsock\")]\n Stream::Vsock(stream) => {\n let (send_request, conn) = builder.handshake(TokioIo::new(stream)).await?;\n tokio::spawn(conn);\n Ok(send_request)\n }\n #[cfg(feature = \"http3\")]\n Stream::Quic(_) => {\n panic!(\"quic is not supported in http2\")\n }\n }\n }\n}\n\nimpl Client {\n #[inline]\n fn is_http2(&self) -> bool {\n self.request_generator.version == http::Version::HTTP_2\n }\n\n #[inline]\n fn is_proxy_http2(&self) -> bool {\n self.proxy_http_version == http::Version::HTTP_2\n }\n\n fn is_work_http2(&self) -> bool {\n if self.proxy_url.is_some() {\n if self.request_generator.https {\n self.is_http2()\n } else {\n self.is_proxy_http2()\n }\n } else {\n self.is_http2()\n }\n }\n\n fn work_type(&self) -> HttpWorkType {\n #[cfg(feature = \"http3\")]\n if self.request_generator.version == http::Version::HTTP_3 {\n return HttpWorkType::H3;\n }\n if self.is_work_http2() {\n HttpWorkType::H2\n } else {\n HttpWorkType::H1\n }\n }\n\n /// Perform a DNS lookup to cache it\n /// This is useful to avoid DNS lookup latency at the first concurrent requests\n pub async fn pre_lookup(&self) -> Result<(), ClientError> {\n // If the client is using a unix socket, we don't need to do a DNS lookup\n #[cfg(unix)]\n if self.unix_socket.is_some() {\n return Ok(());\n }\n // If the client is using a vsock address, we don't need to do a DNS lookup\n #[cfg(feature = \"vsock\")]\n if self.vsock_addr.is_some() {\n return Ok(());\n }\n\n let mut rng = rand::rng();\n let url = self.request_generator.url_generator.generate(&mut rng)?;\n // It automatically caches the result\n self.dns.lookup(&url, &mut rng).await?;\n Ok(())\n }\n\n #[allow(clippy::type_complexity)]\n pub fn generate_request(\n &self,\n rng: &mut R,\n ) -> Result<(Cow<'_, Url>, Request>, R), ClientError> {\n let snapshot = *rng;\n let (url, mut req) = self.request_generator.generate(rng)?;\n if self.proxy_url.is_some() && req.uri().scheme_str() == Some(\"http\") {\n if let Some(authority) = req.uri().authority() {\n let requested_host = authority.host();\n let requested_port = authority.port_u16().unwrap_or(80);\n if let Some(entry) = self\n .dns\n .select_connect_to(requested_host, requested_port, rng)\n {\n let new_authority: http::uri::Authority =\n format_host_port(entry.target_host.as_str(), entry.target_port).parse()?;\n let mut parts = req.uri().clone().into_parts();\n parts.authority = Some(new_authority);\n let new_uri = http::Uri::from_parts(parts)?;\n *req.uri_mut() = new_uri;\n }\n }\n }\n Ok((url, req, snapshot))\n }\n\n /**\n * Returns a stream of the underlying transport. NOT a HTTP client\n */\n pub(crate) async fn client(\n &self,\n url: &Url,\n rng: &mut R,\n http_version: http::Version,\n ) -> Result<(Instant, Stream), ClientError> {\n let timeout_duration = self.connect_timeout;\n\n #[cfg(feature = \"http3\")]\n if http_version == http::Version::HTTP_3 {\n let addr = self.dns.lookup(url, rng).await?;\n let dns_lookup = Instant::now();\n let stream = tokio::time::timeout(timeout_duration, self.quic_client(addr, url)).await;\n return match stream {\n Ok(Ok(stream)) => Ok((dns_lookup, stream)),\n Ok(Err(err)) => Err(err),\n Err(_) => Err(ClientError::Timeout),\n };\n }\n if url.scheme() == \"https\" {\n let addr = self.dns.lookup(url, rng).await?;\n let dns_lookup = Instant::now();\n // If we do not put a timeout here then the connections attempts will\n // linger long past the configured timeout\n let stream =\n tokio::time::timeout(timeout_duration, self.tls_client(addr, url, http_version))\n .await;\n return match stream {\n Ok(Ok(stream)) => Ok((dns_lookup, stream)),\n Ok(Err(err)) => Err(err),\n Err(_) => Err(ClientError::Timeout),\n };\n }\n #[cfg(unix)]\n if let Some(socket_path) = &self.unix_socket {\n let dns_lookup = Instant::now();\n let stream = tokio::time::timeout(\n timeout_duration,\n tokio::net::UnixStream::connect(socket_path),\n )\n .await;\n return match stream {\n Ok(Ok(stream)) => Ok((dns_lookup, Stream::Unix(stream))),\n Ok(Err(err)) => Err(ClientError::Io(err)),\n Err(_) => Err(ClientError::Timeout),\n };\n }\n #[cfg(feature = \"vsock\")]\n if let Some(addr) = self.vsock_addr {\n let dns_lookup = Instant::now();\n let stream =\n tokio::time::timeout(timeout_duration, tokio_vsock::VsockStream::connect(addr))\n .await;\n return match stream {\n Ok(Ok(stream)) => Ok((dns_lookup, Stream::Vsock(stream))),\n Ok(Err(err)) => Err(ClientError::Io(err)),\n Err(_) => Err(ClientError::Timeout),\n };\n }\n // HTTP\n let addr = self.dns.lookup(url, rng).await?;\n let dns_lookup = Instant::now();\n let stream =\n tokio::time::timeout(timeout_duration, tokio::net::TcpStream::connect(addr)).await;\n match stream {\n Ok(Ok(stream)) => {\n stream.set_nodelay(true)?;\n Ok((dns_lookup, Stream::Tcp(stream)))\n }\n Ok(Err(err)) => Err(ClientError::Io(err)),\n Err(_) => Err(ClientError::Timeout),\n }\n }\n\n async fn tls_client(\n &self,\n addr: (std::net::IpAddr, u16),\n url: &Url,\n http_version: http::Version,\n ) -> Result {\n let stream = tokio::net::TcpStream::connect(addr).await?;\n stream.set_nodelay(true)?;\n\n let stream = self.connect_tls(stream, url, http_version).await?;\n\n Ok(Stream::Tls(stream))\n }\n\n #[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\n async fn connect_tls(\n &self,\n stream: S,\n url: &Url,\n http_version: http::Version,\n ) -> Result, ClientError>\n where\n S: AsyncRead + AsyncWrite + Unpin,\n {\n let connector = self\n .native_tls_connectors\n .connector(http_version >= http::Version::HTTP_2);\n let stream = connector\n .connect(url.host_str().ok_or(ClientError::HostNotFound)?, stream)\n .await?;\n\n Ok(stream)\n }\n\n #[cfg(feature = \"rustls\")]\n async fn connect_tls(\n &self,\n stream: S,\n url: &Url,\n http_version: http::Version,\n ) -> Result>, ClientError>\n where\n S: AsyncRead + AsyncWrite + Unpin,\n {\n let connector =\n tokio_rustls::TlsConnector::from(self.rustls_configs.config(http_version).clone());\n let domain = rustls_pki_types::ServerName::try_from(\n url.host_str().ok_or(ClientError::HostNotFound)?,\n )?;\n let stream = connector.connect(domain.to_owned(), stream).await?;\n\n Ok(Box::new(stream))\n }\n\n async fn client_http1(\n &self,\n url: &Url,\n rng: &mut R,\n ) -> Result<(Instant, SendRequestHttp1), ClientError> {\n if let Some(proxy_url) = &self.proxy_url {\n let http_proxy_version = if self.is_proxy_http2() {\n http::Version::HTTP_2\n } else {\n http::Version::HTTP_11\n };\n let (dns_lookup, stream) = self.client(proxy_url, rng, http_proxy_version).await?;\n if url.scheme() == \"https\" {\n let requested_host = url.host_str().ok_or(ClientError::HostNotFound)?;\n let requested_port = url\n .port_or_known_default()\n .ok_or(ClientError::PortNotFound)?;\n let (connect_host, connect_port) = if let Some(entry) =\n self.dns\n .select_connect_to(requested_host, requested_port, rng)\n {\n (entry.target_host.as_str(), entry.target_port)\n } else {\n (requested_host, requested_port)\n };\n let connect_authority = format_host_port(connect_host, connect_port);\n // Do CONNECT request to proxy\n let req = {\n let mut builder = http::Request::builder()\n .method(Method::CONNECT)\n .uri(connect_authority);\n *builder\n .headers_mut()\n .ok_or(ClientError::GetHeaderFromBuilder)? = self.proxy_headers.clone();\n builder.body(http_body_util::Full::default())?\n };\n let res = if self.proxy_http_version == http::Version::HTTP_2 {\n let mut send_request = stream.handshake_http2().await?;\n send_request.send_request(req).await?\n } else {\n let mut send_request = stream.handshake_http1(true).await?;\n send_request.send_request(req).await?\n };\n let stream = hyper::upgrade::on(res).await?;\n let stream = self\n .connect_tls(TokioIo::new(stream), url, self.request_generator.version)\n .await?;\n let (send_request, conn) =\n hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;\n tokio::spawn(conn);\n Ok((dns_lookup, send_request))\n } else {\n // Send full URL in request() for HTTP proxy\n Ok((dns_lookup, stream.handshake_http1(false).await?))\n }\n } else {\n let (dns_lookup, stream) = self.client(url, rng, http::Version::HTTP_11).await?;\n Ok((dns_lookup, stream.handshake_http1(false).await?))\n }\n }\n\n async fn work_http1(\n &self,\n client_state: &mut ClientStateHttp1,\n ) -> Result {\n let do_req = async {\n let (url, request, rng) = self.generate_request(&mut client_state.rng)?;\n let mut start = std::time::Instant::now();\n let mut first_byte: Option = None;\n let mut connection_time: Option = None;\n\n let mut send_request = if let Some(send_request) = client_state.send_request.take() {\n send_request\n } else {\n let (dns_lookup, send_request) =\n self.client_http1(&url, &mut client_state.rng).await?;\n let dialup = std::time::Instant::now();\n\n connection_time = Some(ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n });\n send_request\n };\n while send_request.ready().await.is_err() {\n // This gets hit when the connection for HTTP/1.1 faults\n // This re-connects\n start = std::time::Instant::now();\n let (dns_lookup, send_request_) =\n self.client_http1(&url, &mut client_state.rng).await?;\n send_request = send_request_;\n let dialup = std::time::Instant::now();\n connection_time = Some(ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n });\n }\n match send_request.send_request(request).await {\n Ok(res) => {\n let (parts, mut stream) = res.into_parts();\n let mut status = parts.status;\n\n let mut len_bytes = 0;\n while let Some(chunk) = stream.frame().await {\n if first_byte.is_none() {\n first_byte = Some(std::time::Instant::now())\n }\n len_bytes += chunk?.data_ref().map(|d| d.len()).unwrap_or_default();\n }\n\n if self.redirect_limit != 0 {\n if let Some(location) = parts.headers.get(\"Location\") {\n let (send_request_redirect, new_status, len) = self\n .redirect(\n url,\n rng,\n send_request,\n location,\n self.redirect_limit,\n &mut client_state.rng,\n )\n .await?;\n\n send_request = send_request_redirect;\n status = new_status;\n len_bytes = len;\n }\n }\n\n let end = std::time::Instant::now();\n\n let result = RequestResult {\n rng,\n start_latency_correction: None,\n start,\n first_byte,\n end,\n status,\n len_bytes,\n connection_time,\n };\n\n if !self.disable_keepalive {\n client_state.send_request = Some(send_request);\n }\n\n Ok::<_, ClientError>(result)\n }\n Err(e) => {\n client_state.send_request = Some(send_request);\n Err(e.into())\n }\n }\n };\n\n if let Some(timeout) = self.timeout {\n tokio::select! {\n res = do_req => {\n res\n }\n _ = tokio::time::sleep(timeout) => {\n Err(ClientError::Timeout)\n }\n }\n } else {\n do_req.await\n }\n }\n\n async fn connect_http2(\n &self,\n url: &Url,\n rng: &mut R,\n ) -> Result<(ConnectionTime, SendRequestHttp2), ClientError> {\n let start = std::time::Instant::now();\n if let Some(proxy_url) = &self.proxy_url {\n let http_proxy_version = if self.is_proxy_http2() {\n http::Version::HTTP_2\n } else {\n http::Version::HTTP_11\n };\n let (dns_lookup, stream) = self.client(proxy_url, rng, http_proxy_version).await?;\n if url.scheme() == \"https\" {\n let requested_host = url.host_str().ok_or(ClientError::HostNotFound)?;\n let requested_port = url\n .port_or_known_default()\n .ok_or(ClientError::PortNotFound)?;\n let (connect_host, connect_port) = if let Some(entry) =\n self.dns\n .select_connect_to(requested_host, requested_port, rng)\n {\n (entry.target_host.as_str(), entry.target_port)\n } else {\n (requested_host, requested_port)\n };\n let connect_authority = format_host_port(connect_host, connect_port);\n let req = {\n let mut builder = http::Request::builder()\n .method(Method::CONNECT)\n .uri(connect_authority);\n *builder\n .headers_mut()\n .ok_or(ClientError::GetHeaderFromBuilder)? = self.proxy_headers.clone();\n builder.body(http_body_util::Full::default())?\n };\n let res = if self.proxy_http_version == http::Version::HTTP_2 {\n let mut send_request = stream.handshake_http2().await?;\n send_request.send_request(req).await?\n } else {\n let mut send_request = stream.handshake_http1(true).await?;\n send_request.send_request(req).await?\n };\n let stream = hyper::upgrade::on(res).await?;\n let stream = self\n .connect_tls(TokioIo::new(stream), url, http::Version::HTTP_2)\n .await?;\n let (send_request, conn) =\n hyper::client::conn::http2::Builder::new(TokioExecutor::new())\n // from nghttp2's default\n .initial_stream_window_size((1 << 30) - 1)\n .initial_connection_window_size((1 << 30) - 1)\n .handshake(TokioIo::new(stream))\n .await?;\n tokio::spawn(conn);\n let dialup = std::time::Instant::now();\n\n Ok((\n ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n },\n send_request,\n ))\n } else {\n let send_request = stream.handshake_http2().await?;\n let dialup = std::time::Instant::now();\n Ok((\n ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n },\n send_request,\n ))\n }\n } else {\n let (dns_lookup, stream) = self\n .client(url, rng, self.request_generator.version)\n .await?;\n let send_request = stream.handshake_http2().await?;\n let dialup = std::time::Instant::now();\n Ok((\n ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n },\n send_request,\n ))\n }\n }\n\n async fn work_http2(\n &self,\n client_state: &mut ClientStateHttp2,\n ) -> Result {\n let do_req = async {\n let (_url, request, rng) = self.generate_request(&mut client_state.rng)?;\n let start = std::time::Instant::now();\n let mut first_byte: Option = None;\n let connection_time: Option = None;\n\n match client_state.send_request.send_request(request).await {\n Ok(res) => {\n let (parts, mut stream) = res.into_parts();\n let status = parts.status;\n\n let mut len_bytes = 0;\n while let Some(chunk) = stream.frame().await {\n if first_byte.is_none() {\n first_byte = Some(std::time::Instant::now())\n }\n len_bytes += chunk?.data_ref().map(|d| d.len()).unwrap_or_default();\n }\n\n let end = std::time::Instant::now();\n\n let result = RequestResult {\n rng,\n start_latency_correction: None,\n start,\n first_byte,\n end,\n status,\n len_bytes,\n connection_time,\n };\n\n Ok::<_, ClientError>(result)\n }\n Err(e) => Err(e.into()),\n }\n };\n\n if let Some(timeout) = self.timeout {\n tokio::select! {\n res = do_req => {\n res\n }\n _ = tokio::time::sleep(timeout) => {\n Err(ClientError::Timeout)\n }\n }\n } else {\n do_req.await\n }\n }\n\n #[allow(clippy::type_complexity)]\n async fn redirect(\n &self,\n base_url: Cow<'_, Url>,\n seed: R,\n send_request: SendRequestHttp1,\n location: &http::header::HeaderValue,\n limit: usize,\n rng: &mut R,\n ) -> Result<(SendRequestHttp1, http::StatusCode, usize), ClientError> {\n if limit == 0 {\n return Err(ClientError::TooManyRedirect);\n }\n let url = match Url::parse(location.to_str()?) {\n Ok(url) => url,\n Err(ParseError::RelativeUrlWithoutBase) => Url::options()\n .base_url(Some(&base_url))\n .parse(location.to_str()?)?,\n Err(err) => Err(err)?,\n };\n\n let (mut send_request, send_request_base) =\n if base_url.authority() == url.authority() && !self.disable_keepalive {\n // reuse connection\n (send_request, None)\n } else {\n let (_dns_lookup, stream) = self.client_http1(&url, rng).await?;\n (stream, Some(send_request))\n };\n\n while send_request.ready().await.is_err() {\n let (_dns_lookup, stream) = self.client_http1(&url, rng).await?;\n send_request = stream;\n }\n\n let mut request = self.generate_request(&mut seed.clone())?.1;\n if url.authority() != base_url.authority() {\n request.headers_mut().insert(\n http::header::HOST,\n http::HeaderValue::from_str(url.authority())?,\n );\n }\n *request.uri_mut() = if self.proxy_url.is_some() && url.scheme() == \"http\" {\n // Full URL in request() for HTTP proxy\n url.as_str().parse()?\n } else {\n url[url::Position::BeforePath..].parse()?\n };\n if self.proxy_url.is_some() && request.uri().scheme_str() == Some(\"http\") {\n if let Some(authority) = request.uri().authority() {\n let requested_host = authority.host();\n let requested_port = authority.port_u16().unwrap_or(80);\n if let Some(entry) = self\n .dns\n .select_connect_to(requested_host, requested_port, rng)\n {\n let new_authority: http::uri::Authority =\n format_host_port(entry.target_host.as_str(), entry.target_port).parse()?;\n let mut parts = request.uri().clone().into_parts();\n parts.authority = Some(new_authority);\n let new_uri = http::Uri::from_parts(parts)?;\n *request.uri_mut() = new_uri;\n }\n }\n }\n let res = send_request.send_request(request).await?;\n let (parts, mut stream) = res.into_parts();\n let mut status = parts.status;\n\n let mut len_bytes = 0;\n while let Some(chunk) = stream.frame().await {\n len_bytes += chunk?.data_ref().map(|d| d.len()).unwrap_or_default();\n }\n\n if let Some(location) = parts.headers.get(\"Location\") {\n let (send_request_redirect, new_status, len) =\n Box::pin(self.redirect(base_url, seed, send_request, location, limit - 1, rng))\n .await?;\n send_request = send_request_redirect;\n status = new_status;\n len_bytes = len;\n }\n\n if let Some(send_request_base) = send_request_base {\n Ok((send_request_base, status, len_bytes))\n } else {\n Ok((send_request, status, len_bytes))\n }\n }\n}\n\n/// Check error and decide whether to cancel the connection\npub(crate) fn is_cancel_error(res: &Result) -> bool {\n matches!(res, Err(ClientError::Deadline)) || is_too_many_open_files(res)\n}\n\n/// Check error was \"Too many open file\"\nfn is_too_many_open_files(res: &Result) -> bool {\n res.as_ref()\n .err()\n .map(|err| match err {\n ClientError::Io(io_error) => io_error.raw_os_error() == Some(libc::EMFILE),\n _ => false,\n })\n .unwrap_or(false)\n}\n\n/// Check error was any Hyper error (primarily for HTTP2 connection errors)\nfn is_hyper_error(res: &Result) -> bool {\n res.as_ref()\n .err()\n .map(|err| match err {\n // REVIEW: IoErrors, if indicating the underlying connection has failed,\n // should also cause a stop of HTTP2 requests\n ClientError::Io(_) => true,\n ClientError::Hyper(_) => true,\n _ => false,\n })\n .unwrap_or(false)\n}\n\nasync fn setup_http2(\n client: &Client,\n rng: &mut R,\n) -> Result<(ConnectionTime, SendRequestHttp2), ClientError> {\n let url = client.request_generator.url_generator.generate(rng)?;\n let (connection_time, send_request) = client.connect_http2(&url, rng).await?;\n\n Ok((connection_time, send_request))\n}\n\nasync fn work_http2_once(\n client: &Client,\n client_state: &mut ClientStateHttp2,\n report_tx: &kanal::Sender>,\n connection_time: ConnectionTime,\n start_latency_correction: Option,\n) -> (bool, bool) {\n let mut res = client.work_http2(client_state).await;\n let is_cancel = is_cancel_error(&res);\n let is_reconnect = is_hyper_error(&res);\n set_connection_time(&mut res, connection_time);\n if let Some(start_latency_correction) = start_latency_correction {\n set_start_latency_correction(&mut res, start_latency_correction);\n }\n report_tx.send(res).unwrap();\n (is_cancel, is_reconnect)\n}\n\npub(crate) fn set_connection_time(\n res: &mut Result,\n connection_time: ConnectionTime,\n) {\n if let Ok(res) = res {\n res.connection_time = Some(connection_time);\n }\n}\n\npub(crate) fn set_start_latency_correction(\n res: &mut Result,\n start_latency_correction: std::time::Instant,\n) {\n if let Ok(res) = res {\n res.start_latency_correction = Some(start_latency_correction);\n }\n}\n\npub async fn work_debug(w: &mut W, client: Arc) -> Result<(), ClientError> {\n let mut rng = Pcg64Si::from_rng(&mut rng());\n let (url, request, _) = client.generate_request(&mut rng)?;\n\n writeln!(w, \"{request:#?}\")?;\n\n let response = match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => {\n let (_, (h3_connection, client_state)) = client.connect_http3(&url, &mut rng).await?;\n\n send_debug_request_http3(h3_connection, client_state, request).await?\n }\n HttpWorkType::H2 => {\n let (_, mut client_state) = client.connect_http2(&url, &mut rng).await?;\n let response = client_state.send_request(request).await?;\n let (parts, body) = response.into_parts();\n let body = body.collect().await.unwrap().to_bytes();\n\n http::Response::from_parts(parts, body)\n }\n HttpWorkType::H1 => {\n let (_dns_lookup, mut send_request) = client.client_http1(&url, &mut rng).await?;\n\n let response = send_request.send_request(request).await?;\n let (parts, body) = response.into_parts();\n let body = body.collect().await.unwrap().to_bytes();\n\n http::Response::from_parts(parts, body)\n }\n };\n\n writeln!(w, \"{response:#?}\")?;\n\n Ok(())\n}\n\n/// Run n tasks by m workers\npub async fn work(\n client: Arc,\n report_tx: kanal::Sender>,\n n_tasks: usize,\n n_connections: usize,\n n_http2_parallel: usize,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work(client, report_tx, n_tasks, n_connections, n_http2_parallel).await;\n return;\n }\n\n use std::sync::atomic::{AtomicUsize, Ordering};\n let counter = Arc::new(AtomicUsize::new(0));\n\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n tokio::spawn(async move {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http2_parallel)\n .map(|_| {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n tokio::spawn(async move {\n while counter.fetch_add(1, Ordering::Relaxed)\n < n_tasks\n {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n None,\n )\n .await;\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n\n true\n })\n })\n .collect::>();\n\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n\n if connection_gone {\n return;\n }\n }\n Err(err) => {\n if counter.fetch_add(1, Ordering::Relaxed) < n_tasks {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n }\n }\n }\n })\n })\n .collect::>();\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n tokio::spawn(async move {\n let mut client_state = ClientStateHttp1::default();\n while counter.fetch_add(1, Ordering::Relaxed) < n_tasks {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel {\n break;\n }\n }\n })\n })\n .collect::>();\n for f in futures {\n let _ = f.await;\n }\n }\n };\n}\n\n/// n tasks by m workers limit to qps works in a second\npub async fn work_with_qps(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n n_tasks: usize,\n n_connections: usize,\n n_http_parallel: usize,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work_with_qps(\n client,\n report_tx,\n query_limit,\n n_tasks,\n n_connections,\n n_http_parallel,\n )\n .await;\n return;\n }\n\n let (tx, rx) = kanal::unbounded::<()>();\n\n let work_queue = async move {\n match query_limit {\n QueryLimit::Qps(qps) => {\n let start = std::time::Instant::now();\n for i in 0..n_tasks {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n tx.send(())?;\n }\n }\n QueryLimit::Burst(duration, rate) => {\n let mut n = 0;\n // Handle via rate till n_tasks out of bound\n while n + rate < n_tasks {\n tokio::time::sleep(duration).await;\n for _ in 0..rate {\n tx.send(())?;\n }\n n += rate;\n }\n // Handle the remaining tasks\n if n_tasks > n {\n tokio::time::sleep(duration).await;\n for _ in 0..n_tasks - n {\n tx.send(())?;\n }\n }\n }\n }\n // tx gone\n drop(tx);\n Ok::<(), kanal::SendError>(())\n };\n\n let rx = rx.to_async();\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n tokio::spawn(async move {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http_parallel)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n tokio::spawn(async move {\n while let Ok(()) = rx.recv().await {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n None,\n )\n .await;\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n })\n })\n .collect::>();\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n if connection_gone {\n return;\n }\n }\n Err(err) => {\n // Consume a task\n if let Ok(()) = rx.recv().await {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n }\n }\n }\n })\n })\n .collect::>();\n\n work_queue.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n tokio::spawn(async move {\n let mut client_state = ClientStateHttp1::default();\n while let Ok(()) = rx.recv().await {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel {\n break;\n }\n }\n })\n })\n .collect::>();\n\n work_queue.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n }\n };\n}\n\n/// n tasks by m workers limit to qps works in a second with latency correction\npub async fn work_with_qps_latency_correction(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n n_tasks: usize,\n n_connections: usize,\n n_http2_parallel: usize,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work_with_qps_latency_correction(\n client,\n report_tx,\n query_limit,\n n_tasks,\n n_connections,\n n_http2_parallel,\n )\n .await;\n return;\n }\n\n let (tx, rx) = kanal::unbounded();\n\n let work_queue = async move {\n match query_limit {\n QueryLimit::Qps(qps) => {\n let start = std::time::Instant::now();\n for i in 0..n_tasks {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let now = std::time::Instant::now();\n tx.send(now)?;\n }\n }\n QueryLimit::Burst(duration, rate) => {\n let mut n = 0;\n // Handle via rate till n_tasks out of bound\n while n + rate < n_tasks {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n for _ in 0..rate {\n tx.send(now)?;\n }\n n += rate;\n }\n // Handle the remaining tasks\n if n_tasks > n {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n for _ in 0..n_tasks - n {\n tx.send(now)?;\n }\n }\n }\n }\n\n // tx gone\n drop(tx);\n Ok::<(), kanal::SendError>(())\n };\n\n let rx = rx.to_async();\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n tokio::spawn(async move {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http2_parallel)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n tokio::spawn(async move {\n while let Ok(start) = rx.recv().await {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n Some(start),\n )\n .await;\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n })\n })\n .collect::>();\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n if connection_gone {\n return;\n }\n }\n Err(err) => {\n // Consume a task\n if rx.recv().await.is_ok() {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n }\n }\n }\n })\n })\n .collect::>();\n\n work_queue.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let mut client_state = ClientStateHttp1::default();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n tokio::spawn(async move {\n while let Ok(start) = rx.recv().await {\n let mut res = client.work_http1(&mut client_state).await;\n set_start_latency_correction(&mut res, start);\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel {\n break;\n }\n }\n })\n })\n .collect::>();\n\n work_queue.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n }\n }\n}\n\n/// Run until dead_line by n workers\npub async fn work_until(\n client: Arc,\n report_tx: kanal::Sender>,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http_parallel: usize,\n wait_ongoing_requests_after_deadline: bool,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work_until(\n client,\n report_tx,\n dead_line,\n n_connections,\n n_http_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n return;\n }\n\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n // Using semaphore to control the deadline\n // Maybe there is a better concurrent primitive to do this\n let s = Arc::new(tokio::sync::Semaphore::new(0));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let s = s.clone();\n tokio::spawn(async move {\n let s = s.clone();\n // Keep trying to establish or re-establish connections up to the deadline\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n // Setup the parallel workers for each HTTP2 connection\n let futures = (0..n_http_parallel)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n let s = s.clone();\n tokio::spawn(async move {\n // This is where HTTP2 loops to make all the requests for a given client and worker\n loop {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n None,\n )\n .await;\n\n let is_cancel = is_cancel || s.is_closed();\n if is_cancel || is_reconnect {\n break is_cancel;\n }\n }\n })\n })\n .collect::>();\n\n let mut connection_gone = false;\n for f in futures {\n tokio::select! {\n r = f => {\n match r {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n _ = s.acquire() => {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n connection_gone = true;\n }\n }\n }\n if connection_gone {\n return;\n }\n }\n\n Err(err) => {\n report_tx.send(Err(err)).unwrap();\n if s.is_closed() {\n break;\n }\n }\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n s.close();\n\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let is_end = Arc::new(AtomicBool::new(false));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let mut client_state = ClientStateHttp1::default();\n let is_end = is_end.clone();\n tokio::spawn(async move {\n loop {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel || is_end.load(Relaxed) {\n break;\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n is_end.store(true, Relaxed);\n\n if wait_ongoing_requests_after_deadline {\n for f in futures {\n let _ = f.await;\n }\n } else {\n for f in futures {\n f.abort();\n if let Err(e) = f.await {\n if e.is_cancelled() {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n }\n }\n }\n }\n }\n };\n}\n\n/// Run until dead_line by n workers limit to qps works in a second\n#[allow(clippy::too_many_arguments)]\npub async fn work_until_with_qps(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n start: std::time::Instant,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http2_parallel: usize,\n wait_ongoing_requests_after_deadline: bool,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work_until_with_qps(\n client,\n report_tx,\n query_limit,\n start,\n dead_line,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n return;\n }\n\n let rx = match query_limit {\n QueryLimit::Qps(qps) => {\n let (tx, rx) = kanal::unbounded::<()>();\n tokio::spawn(async move {\n for i in 0.. {\n if std::time::Instant::now() > dead_line {\n break;\n }\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let _ = tx.send(());\n }\n // tx gone\n });\n rx\n }\n QueryLimit::Burst(duration, rate) => {\n let (tx, rx) = kanal::unbounded();\n tokio::spawn(async move {\n // Handle via rate till deadline is reached\n for _ in 0.. {\n if std::time::Instant::now() > dead_line {\n break;\n }\n\n tokio::time::sleep(duration).await;\n for _ in 0..rate {\n let _ = tx.send(());\n }\n }\n // tx gone\n });\n rx\n }\n };\n\n let rx = rx.to_async();\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n let s = Arc::new(tokio::sync::Semaphore::new(0));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let s = s.clone();\n tokio::spawn(async move {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http2_parallel)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n let s = s.clone();\n tokio::spawn(async move {\n while let Ok(()) = rx.recv().await {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n None,\n )\n .await;\n\n let is_cancel = is_cancel || s.is_closed();\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n })\n })\n .collect::>();\n let mut connection_gone = false;\n for f in futures {\n tokio::select! {\n r = f => {\n match r {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n _ = s.acquire() => {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n connection_gone = true;\n }\n }\n }\n if connection_gone {\n return;\n }\n }\n Err(err) => {\n // Consume a task\n if rx.recv().await.is_ok() {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n\n if s.is_closed() {\n return;\n }\n }\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n s.close();\n\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let is_end = Arc::new(AtomicBool::new(false));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let mut client_state = ClientStateHttp1::default();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let is_end = is_end.clone();\n tokio::spawn(async move {\n while let Ok(()) = rx.recv().await {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel || is_end.load(Relaxed) {\n break;\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n is_end.store(true, Relaxed);\n\n if wait_ongoing_requests_after_deadline {\n for f in futures {\n let _ = f.await;\n }\n } else {\n for f in futures {\n f.abort();\n if let Err(e) = f.await {\n if e.is_cancelled() {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n }\n }\n }\n }\n }\n }\n}\n\n/// Run until dead_line by n workers limit to qps works in a second with latency correction\n#[allow(clippy::too_many_arguments)]\npub async fn work_until_with_qps_latency_correction(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n start: std::time::Instant,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http2_parallel: usize,\n wait_ongoing_requests_after_deadline: bool,\n) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::work_until_with_qps_latency_correction(\n client,\n report_tx,\n query_limit,\n start,\n dead_line,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n return;\n }\n\n let (tx, rx) = kanal::unbounded();\n match query_limit {\n QueryLimit::Qps(qps) => {\n tokio::spawn(async move {\n for i in 0.. {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let now = std::time::Instant::now();\n if now > dead_line {\n break;\n }\n let _ = tx.send(now);\n }\n // tx gone\n });\n }\n QueryLimit::Burst(duration, rate) => {\n tokio::spawn(async move {\n // Handle via rate till deadline is reached\n loop {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n if now > dead_line {\n break;\n }\n\n for _ in 0..rate {\n let _ = tx.send(now);\n }\n }\n // tx gone\n });\n }\n };\n\n let rx = rx.to_async();\n match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n let s = Arc::new(tokio::sync::Semaphore::new(0));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let s = s.clone();\n tokio::spawn(async move {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http2_parallel)\n .map(|_| {\n let client = client.clone();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n let s = s.clone();\n tokio::spawn(async move {\n while let Ok(start) = rx.recv().await {\n let (is_cancel, is_reconnect) =\n work_http2_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n Some(start),\n )\n .await;\n let is_cancel = is_cancel || s.is_closed();\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n })\n })\n .collect::>();\n let mut connection_gone = false;\n for f in futures {\n tokio::select! {\n r = f => {\n match r {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n _ = s.acquire() => {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n connection_gone = true;\n }\n }\n }\n if connection_gone {\n return;\n }\n }\n\n Err(err) => {\n if rx.recv().await.is_ok() {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n\n if s.is_closed() {\n return;\n }\n }\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n s.close();\n\n for f in futures {\n let _ = f.await;\n }\n }\n HttpWorkType::H1 => {\n let is_end = Arc::new(AtomicBool::new(false));\n\n let futures = (0..n_connections)\n .map(|_| {\n let client = client.clone();\n let mut client_state = ClientStateHttp1::default();\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let is_end = is_end.clone();\n tokio::spawn(async move {\n while let Ok(start) = rx.recv().await {\n let mut res = client.work_http1(&mut client_state).await;\n set_start_latency_correction(&mut res, start);\n let is_cancel = is_cancel_error(&res);\n report_tx.send(res).unwrap();\n if is_cancel || is_end.load(Relaxed) {\n break;\n }\n }\n })\n })\n .collect::>();\n\n tokio::time::sleep_until(dead_line.into()).await;\n is_end.store(true, Relaxed);\n\n if wait_ongoing_requests_after_deadline {\n for f in futures {\n let _ = f.await;\n }\n } else {\n for f in futures {\n f.abort();\n if let Err(e) = f.await {\n if e.is_cancelled() {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n }\n }\n }\n }\n }\n }\n}\n\n/// Optimized workers for `--no-tui` mode\npub mod fast {\n use std::sync::{\n Arc,\n atomic::{AtomicBool, AtomicIsize, Ordering},\n };\n\n use rand::SeedableRng;\n\n use crate::{\n client::{\n ClientError, ClientStateHttp1, ClientStateHttp2, HttpWorkType, is_cancel_error,\n is_hyper_error, set_connection_time, setup_http2,\n },\n pcg64si::Pcg64Si,\n result_data::ResultData,\n };\n\n use super::Client;\n\n /// Run n tasks by m workers\n pub async fn work(\n client: Arc,\n report_tx: kanal::Sender,\n n_tasks: usize,\n n_connections: usize,\n n_http_parallel: usize,\n ) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::fast::work(\n client,\n report_tx,\n n_tasks,\n n_connections,\n n_http_parallel,\n )\n .await;\n return;\n }\n\n let counter = Arc::new(AtomicIsize::new(n_tasks as isize));\n let num_threads = num_cpus::get_physical();\n let connections = (0..num_threads).filter_map(|i| {\n let num_connection = n_connections / num_threads\n + (if (n_connections % num_threads) > i {\n 1\n } else {\n 0\n });\n if num_connection > 0 {\n Some(num_connection)\n } else {\n None\n }\n });\n let token = tokio_util::sync::CancellationToken::new();\n\n let handles = match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n connections\n .map(|num_connections| {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n let token = token.clone();\n\n std::thread::spawn(move || {\n let client = client.clone();\n let local = tokio::task::LocalSet::new();\n for _ in 0..num_connections {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n let token = token.clone();\n local.spawn_local(Box::pin(async move {\n let mut has_err = false;\n let mut result_data_err = ResultData::default();\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n let client = client.clone();\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http_parallel)\n .map(|_| {\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(\n &mut rand::rng(),\n ),\n send_request: send_request.clone(),\n };\n let counter = counter.clone();\n let client = client.clone();\n let report_tx = report_tx.clone();\n let token = token.clone();\n tokio::task::spawn_local(async move {\n let mut result_data =\n ResultData::default();\n\n let work = async {\n while counter\n .fetch_sub(1, Ordering::Relaxed)\n > 0\n {\n let mut res = client\n .work_http2(\n &mut client_state,\n )\n .await;\n let is_cancel =\n is_cancel_error(&res);\n let is_reconnect =\n is_hyper_error(&res);\n set_connection_time(\n &mut res,\n connection_time,\n );\n\n result_data.push(res);\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n };\n\n let is_cancel = tokio::select! {\n is_cancel = work => {\n is_cancel\n }\n _ = token.cancelled() => {\n true\n }\n };\n\n report_tx.send(result_data).unwrap();\n is_cancel\n })\n })\n .collect::>();\n\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n\n if connection_gone {\n break;\n }\n }\n Err(err) => {\n if counter.fetch_sub(1, Ordering::Relaxed) > 0 {\n has_err = true;\n result_data_err.push(Err(err));\n } else {\n break;\n }\n }\n }\n }\n if has_err {\n report_tx.send(result_data_err).unwrap();\n }\n }));\n }\n\n rt.block_on(local);\n })\n })\n .collect::>()\n }\n HttpWorkType::H1 => connections\n .map(|num_connection| {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n\n let token = token.clone();\n std::thread::spawn(move || {\n let local = tokio::task::LocalSet::new();\n\n for _ in 0..num_connection {\n let report_tx = report_tx.clone();\n let counter = counter.clone();\n let client = client.clone();\n let token = token.clone();\n local.spawn_local(Box::pin(async move {\n let mut result_data = ResultData::default();\n\n tokio::select! {\n _ = token.cancelled() => {}\n _ = async {\n let mut client_state = ClientStateHttp1::default();\n while counter.fetch_sub(1, Ordering::Relaxed) > 0 {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n result_data.push(res);\n if is_cancel {\n break;\n }\n }\n } => {}\n }\n report_tx.send(result_data).unwrap();\n }));\n }\n rt.block_on(local);\n })\n })\n .collect::>(),\n };\n\n tokio::spawn(async move {\n tokio::signal::ctrl_c().await.unwrap();\n token.cancel();\n });\n\n tokio::task::block_in_place(|| {\n for handle in handles {\n let _ = handle.join();\n }\n });\n }\n\n /// Run until dead_line by n workers\n pub async fn work_until(\n client: Arc,\n report_tx: kanal::Sender,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http_parallel: usize,\n wait_ongoing_requests_after_deadline: bool,\n ) {\n #[cfg(feature = \"http3\")]\n if matches!(client.work_type(), HttpWorkType::H3) {\n crate::client_h3::fast::work_until(\n client,\n report_tx,\n dead_line,\n n_connections,\n n_http_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n return;\n }\n\n let num_threads = num_cpus::get_physical();\n\n let is_end = Arc::new(AtomicBool::new(false));\n let connections = (0..num_threads).filter_map(|i| {\n let num_connection = n_connections / num_threads\n + (if (n_connections % num_threads) > i {\n 1\n } else {\n 0\n });\n if num_connection > 0 {\n Some(num_connection)\n } else {\n None\n }\n });\n let token = tokio_util::sync::CancellationToken::new();\n let handles = match client.work_type() {\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => unreachable!(),\n HttpWorkType::H2 => {\n connections\n .map(|num_connections| {\n let report_tx = report_tx.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n let token = token.clone();\n let is_end = is_end.clone();\n\n std::thread::spawn(move || {\n let client = client.clone();\n let local = tokio::task::LocalSet::new();\n for _ in 0..num_connections {\n let report_tx = report_tx.clone();\n let client = client.clone();\n let token = token.clone();\n let is_end = is_end.clone();\n local.spawn_local(Box::pin(async move {\n let mut has_err = false;\n let mut result_data_err = ResultData::default();\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n let client = client.clone();\n match setup_http2(&client, &mut rng).await {\n Ok((connection_time, send_request)) => {\n let futures = (0..n_http_parallel)\n .map(|_| {\n let mut client_state = ClientStateHttp2 {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request: send_request.clone(),\n };\n let client = client.clone();\n let report_tx = report_tx.clone();\n let token = token.clone();\n let is_end = is_end.clone();\n tokio::task::spawn_local(async move {\n let mut result_data = ResultData::default();\n\n let work = async {\n loop {\n let mut res = client\n .work_http2(&mut client_state)\n .await;\n let is_cancel = is_cancel_error(&res) || is_end.load(Ordering::Relaxed);\n let is_reconnect = is_hyper_error(&res);\n set_connection_time(\n &mut res,\n connection_time,\n );\n\n result_data.push(res);\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n };\n\n let is_cancel = tokio::select! {\n is_cancel = work => {\n is_cancel\n }\n _ = token.cancelled() => {\n result_data.push(Err(ClientError::Deadline));\n true\n }\n };\n\n report_tx.send(result_data).unwrap();\n is_cancel\n })\n })\n .collect::>();\n\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n\n if connection_gone {\n break;\n }\n }\n Err(err) => {\n has_err = true;\n result_data_err.push(Err(err));\n if is_end.load(Ordering::Relaxed) {\n break;\n }\n }\n }\n }\n if has_err {\n report_tx.send(result_data_err).unwrap();\n }\n }));\n }\n\n rt.block_on(local);\n })\n })\n .collect::>()\n }\n HttpWorkType::H1 => connections\n .map(|num_connection| {\n let report_tx = report_tx.clone();\n let is_end = is_end.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n\n let token = token.clone();\n std::thread::spawn(move || {\n let local = tokio::task::LocalSet::new();\n\n for _ in 0..num_connection {\n let report_tx = report_tx.clone();\n let is_end = is_end.clone();\n let client = client.clone();\n let token = token.clone();\n local.spawn_local(Box::pin(async move {\n let mut result_data = ResultData::default();\n\n let work = async {\n let mut client_state = ClientStateHttp1::default();\n loop {\n let res = client.work_http1(&mut client_state).await;\n let is_cancel = is_cancel_error(&res);\n result_data.push(res);\n if is_cancel || is_end.load(Ordering::Relaxed) {\n break;\n }\n }\n };\n\n tokio::select! {\n _ = work => {\n }\n _ = token.cancelled() => {\n result_data.push(Err(ClientError::Deadline));\n }\n }\n report_tx.send(result_data).unwrap();\n }));\n }\n rt.block_on(local);\n })\n })\n .collect::>(),\n };\n tokio::select! {\n _ = tokio::time::sleep_until(dead_line.into()) => {\n }\n _ = tokio::signal::ctrl_c() => {\n }\n }\n\n is_end.store(true, Ordering::Relaxed);\n\n if !wait_ongoing_requests_after_deadline {\n token.cancel();\n }\n tokio::task::block_in_place(|| {\n for handle in handles {\n let _ = handle.join();\n }\n });\n }\n}\n" }, { "path": "src/client_h3.rs", "content": "use bytes::Buf;\nuse bytes::Bytes;\nuse core::sync::atomic::Ordering;\nuse http::Request;\nuse http_body_util::BodyExt;\nuse hyper::http;\nuse kanal::AsyncReceiver;\nuse quinn::default_runtime;\nuse std::net::SocketAddr;\nuse std::net::UdpSocket;\nuse std::sync::Arc;\nuse std::sync::atomic::AtomicBool;\nuse std::sync::atomic::AtomicIsize;\nuse std::time::Instant;\n\nuse tokio::sync::Semaphore;\nuse url::Url;\n\npub type SendRequestHttp3 = (\n h3::client::Connection,\n h3::client::SendRequest,\n);\n\n// HTTP3-specific error types\n#[derive(thiserror::Error, Debug)]\npub enum Http3Error {\n #[error(\"QUIC Client: {0}\")]\n QuicClientConfig(#[from] quinn::crypto::rustls::NoInitialCipherSuite),\n #[error(\"QUIC connect: {0}\")]\n QuicConnect(#[from] quinn::ConnectError),\n #[error(\"QUIC connection: {0}\")]\n QuicConnection(#[from] quinn::ConnectionError),\n #[error(\"Quic connection closed earlier than expected\")]\n QuicDriverClosedEarly(#[from] tokio::sync::oneshot::error::RecvError),\n #[error(\"HTTP3 connection: {0}\")]\n H3Connection(#[from] h3::error::ConnectionError),\n #[error(\"HTTP3 Stream: {0}\")]\n H3Stream(#[from] h3::error::StreamError),\n}\n\nuse crate::client::QueryLimit;\nuse crate::client::{\n Client, ClientError, ConnectionTime, RequestResult, Stream, is_cancel_error,\n set_connection_time, set_start_latency_correction,\n};\nuse crate::pcg64si::Pcg64Si;\nuse crate::result_data::ResultData;\nuse rand::SeedableRng;\nuse rand::prelude::Rng;\n\npub(crate) struct ClientStateHttp3 {\n pub(crate) rng: Pcg64Si,\n pub(crate) send_request: h3::client::SendRequest,\n}\n\nimpl ClientStateHttp3 {\n fn new(send_request: h3::client::SendRequest) -> Self {\n Self {\n rng: SeedableRng::from_rng(&mut rand::rng()),\n send_request,\n }\n }\n}\n\nimpl Client {\n pub(crate) async fn connect_http3(\n &self,\n url: &Url,\n rng: &mut R,\n ) -> Result<(ConnectionTime, SendRequestHttp3), ClientError> {\n let start = std::time::Instant::now();\n let (dns_lookup, stream) = self.client(url, rng, http::Version::HTTP_3).await?;\n let send_request = stream.handshake_http3().await?;\n let dialup = std::time::Instant::now();\n Ok((\n ConnectionTime {\n dns_lookup: dns_lookup - start,\n dialup: dialup - start,\n },\n send_request,\n ))\n }\n\n pub(crate) async fn quic_client(\n &self,\n addr: (std::net::IpAddr, u16),\n url: &Url,\n ) -> Result {\n let endpoint_config = h3_quinn::quinn::EndpointConfig::default();\n let local_socket = if addr.0.is_ipv6() {\n UdpSocket::bind(\"[::]:0\").expect(\"couldn't bind to address\")\n } else {\n UdpSocket::bind(\"0.0.0.0:0\").expect(\"couldn't bind to address\")\n };\n // If we can set the right build flags, we can use `h3_quinn::quinn::Endpoint::client` instead\n let mut client_endpoint = h3_quinn::quinn::Endpoint::new(\n endpoint_config,\n None,\n local_socket,\n default_runtime().unwrap(),\n )\n .unwrap();\n\n let tls_config = self.rustls_configs.config(http::Version::HTTP_3).clone();\n let client_config = quinn::ClientConfig::new(Arc::new(\n quinn::crypto::rustls::QuicClientConfig::try_from(tls_config)\n .map_err(Http3Error::from)?,\n ));\n client_endpoint.set_default_client_config(client_config);\n\n let remote_socket_address = SocketAddr::new(addr.0, addr.1);\n let server_name = url.host_str().ok_or(ClientError::HostNotFound)?;\n let conn = client_endpoint\n .connect(remote_socket_address, server_name)\n .map_err(Http3Error::from)?\n .await\n .map_err(Http3Error::from)?;\n Ok(Stream::Quic(conn))\n }\n\n pub(crate) async fn work_http3(\n &self,\n client_state: &mut ClientStateHttp3,\n ) -> Result {\n let do_req = async {\n let (_url, request, rng) = self.generate_request(&mut client_state.rng)?;\n let start = std::time::Instant::now();\n let connection_time: Option = None;\n let mut first_byte: Option = None;\n\n // if we implement http_body::Body on our H3 SendRequest, we can do some nice streaming stuff\n // with the response here. However as we don't really use the response we can get away\n // with not doing this for now\n let (head, mut req_body) = request.into_parts();\n let request = http::request::Request::from_parts(head, ());\n let mut stream = client_state\n .send_request\n .send_request(request)\n .await\n .map_err(Http3Error::from)?;\n // send the request body now\n if let Some(Ok(frame)) = req_body.frame().await {\n if let Ok(data) = frame.into_data() {\n stream.send_data(data).await.map_err(Http3Error::from)?;\n }\n }\n stream.finish().await.map_err(Http3Error::from)?;\n\n // now read the response headers\n let response = stream.recv_response().await.map_err(Http3Error::from)?;\n let (parts, _) = response.into_parts();\n let status = parts.status;\n // now read the response body\n let mut len_bytes = 0;\n while let Some(chunk) = stream.recv_data().await.map_err(Http3Error::from)? {\n if first_byte.is_none() {\n first_byte = Some(std::time::Instant::now())\n }\n len_bytes += chunk.remaining();\n }\n let end = std::time::Instant::now();\n\n let result = RequestResult {\n rng,\n start_latency_correction: None,\n start,\n first_byte,\n end,\n status,\n len_bytes,\n connection_time,\n };\n\n Ok::<_, ClientError>(result)\n };\n\n if let Some(timeout) = self.timeout {\n tokio::select! {\n res = do_req => {\n res\n }\n _ = tokio::time::sleep(timeout) => {\n Err(ClientError::Timeout)\n }\n }\n } else {\n do_req.await\n }\n }\n}\n\nimpl Stream {\n async fn handshake_http3(self) -> Result {\n let Stream::Quic(quic_conn) = self else {\n panic!(\"You cannot call http3 handshake on a non-quic stream\");\n };\n let h3_quinn_conn = h3_quinn::Connection::new(quic_conn);\n // TODO add configuration settings to allow 'send_grease' etc.\n\n Ok(h3::client::new(h3_quinn_conn).await?)\n }\n}\n\npub(crate) async fn send_debug_request_http3(\n h3_connection: h3::client::Connection,\n mut client_state: h3::client::SendRequest,\n request: Request>,\n) -> Result, Http3Error> {\n // Prepare a channel to stop the driver thread\n let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();\n // Run the driver\n let http3_driver = spawn_http3_driver(h3_connection, shutdown_rx);\n\n let (head, mut req_body) = request.into_parts();\n let request = http::request::Request::from_parts(head, ());\n\n let mut stream = client_state.send_request(request).await?;\n if let Some(Ok(frame)) = req_body.frame().await {\n if let Ok(data) = frame.into_data() {\n stream.send_data(data).await?;\n }\n }\n\n stream.finish().await?;\n\n let response = stream.recv_response().await.unwrap_or_else(|err| {\n panic!(\"{}\", err);\n });\n let mut body_bytes = bytes::BytesMut::new();\n\n while let Some(mut chunk) = stream.recv_data().await? {\n let bytes = chunk.copy_to_bytes(chunk.remaining());\n body_bytes.extend_from_slice(&bytes);\n }\n let body = body_bytes.freeze();\n let (parts, _) = response.into_parts();\n let _ = shutdown_tx.send(());\n let _ = http3_driver.await.unwrap();\n Ok(http::Response::from_parts(parts, body))\n}\n\n/**\n * Create `n_connections` parallel HTTP3 connections (on independent QUIC connections).\n * On each of those, run `n_http3_parallel` requests continuously until `deadline` is reached.\n */\npub(crate) async fn parallel_work_http3(\n n_connections: usize,\n n_http_parallel: usize,\n rx: AsyncReceiver>,\n report_tx: kanal::Sender>,\n client: Arc,\n deadline: Option,\n) -> Vec> {\n let s = Arc::new(tokio::sync::Semaphore::new(0));\n let has_deadline = deadline.is_some();\n\n let futures = (0..n_connections)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n let s = s.clone();\n tokio::spawn(create_and_load_up_single_connection_http3(\n n_http_parallel,\n rx,\n report_tx,\n client,\n s,\n ))\n })\n .collect::>();\n\n if has_deadline {\n tokio::time::sleep_until(deadline.unwrap().into()).await;\n s.close();\n }\n\n futures\n}\n\n/**\n * For use in the 'slow' functions - send a report of every response in real time for display to the end-user.\n * Semaphore is closed to shut down all the tasks.\n * Very similar to how http2 loops work, just that we explicitly spawn the HTTP3 connection driver.\n */\nasync fn create_and_load_up_single_connection_http3(\n n_http_parallel: usize,\n rx: AsyncReceiver>,\n report_tx: kanal::Sender>,\n client: Arc,\n s: Arc,\n) {\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n // create a HTTP3 connection\n match setup_http3(&client, &mut rng).await {\n Ok((connection_time, (h3_connection, send_request))) => {\n let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();\n let http3_driver = spawn_http3_driver(h3_connection, shutdown_rx);\n let futures = (0..n_http_parallel)\n .map(|_| {\n let report_tx = report_tx.clone();\n let rx = rx.clone();\n let client = client.clone();\n let mut client_state = ClientStateHttp3::new(send_request.clone());\n let s = s.clone();\n tokio::spawn(async move {\n // This is where HTTP3 loops to make all the requests for a given client and worker\n while let Ok(start_time_option) = rx.recv().await {\n let (is_cancel, is_reconnect) = work_http3_once(\n &client,\n &mut client_state,\n &report_tx,\n connection_time,\n start_time_option,\n )\n .await;\n\n let is_cancel = is_cancel || s.is_closed();\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n true\n })\n })\n .collect::>();\n drop(send_request);\n\n // collect all the requests we have spawned, and end the process if/when the semaphore says\n let mut connection_gone = false;\n for f in futures {\n tokio::select! {\n r = f => {\n match r {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n _ = s.acquire() => {\n report_tx.send(Err(ClientError::Deadline)).unwrap();\n connection_gone = true;\n }\n }\n }\n if connection_gone {\n // Try and politely shut down the HTTP3 connection\n let _ = shutdown_tx.send(());\n let _ = http3_driver.await;\n return;\n }\n }\n Err(err) => {\n if s.is_closed() {\n break;\n // Consume a task\n } else if rx.recv().await.is_ok() {\n report_tx.send(Err(err)).unwrap();\n } else {\n return;\n }\n }\n }\n }\n}\n\n/**\n * This is structured to work very similarly to the `setup_http2`\n * function in `client.rs`\n */\npub(crate) async fn setup_http3(\n client: &Client,\n rng: &mut R,\n) -> Result<(ConnectionTime, SendRequestHttp3), ClientError> {\n let url = client.request_generator.url_generator.generate(rng)?;\n // Whatever rng state, all urls should have the same authority\n let (connection_time, send_request) = client.connect_http3(&url, rng).await?;\n\n Ok((connection_time, send_request))\n}\n\npub(crate) fn spawn_http3_driver(\n mut h3_connection: h3::client::Connection,\n shutdown_rx: tokio::sync::oneshot::Receiver<()>,\n) -> tokio::task::JoinHandle> {\n tokio::spawn(async move {\n tokio::select! {\n // Drive the connection\n closed = std::future::poll_fn(|cx| h3_connection.poll_close(cx)) => {\n if closed.is_h3_no_error() {\n Ok(())\n } else {\n Err(Http3Error::H3Connection(closed))\n }\n },\n // Listen for shutdown condition\n _ = shutdown_rx => {\n // Initiate shutdown\n h3_connection.shutdown(0).await?;\n // Wait for ongoing work to complete\n let closed = std::future::poll_fn(|cx| h3_connection.poll_close(cx)).await;\n if closed.is_h3_no_error() {\n Ok(())\n } else {\n Err(Http3Error::H3Connection(closed))\n }\n }\n }\n })\n}\n\npub(crate) async fn work_http3_once(\n client: &Client,\n client_state: &mut ClientStateHttp3,\n report_tx: &kanal::Sender>,\n connection_time: ConnectionTime,\n start_latency_correction: Option,\n) -> (bool, bool) {\n let mut res = client.work_http3(client_state).await;\n let is_cancel = is_cancel_error(&res);\n let is_reconnect = is_h3_error(&res);\n set_connection_time(&mut res, connection_time);\n if let Some(start_latency_correction) = start_latency_correction {\n set_start_latency_correction(&mut res, start_latency_correction);\n }\n report_tx.send(res).unwrap();\n (is_cancel, is_reconnect)\n}\n\nfn is_h3_error(res: &Result) -> bool {\n res.as_ref()\n .err()\n .map(|err| matches!(err, ClientError::Http3(_) | ClientError::Io(_)))\n .unwrap_or(false)\n}\n\n/**\n * 'Fast' implementation of HTTP3 load generation.\n * If `n_tasks` is set, it will generate up to that many tasks.\n * Othrwise it will terminate when `is_end` becomes set to true.\n */\n#[allow(clippy::too_many_arguments)]\npub(crate) fn http3_connection_fast_work_until(\n num_connections: usize,\n n_http_parallel: usize,\n report_tx: kanal::Sender,\n client: Arc,\n token: tokio_util::sync::CancellationToken,\n counter: Option>,\n is_end: Arc,\n rt: tokio::runtime::Runtime,\n) {\n let is_counting_tasks = counter.is_some();\n let client = client.clone();\n let local = tokio::task::LocalSet::new();\n for _ in 0..num_connections {\n let report_tx = report_tx.clone();\n let client = client.clone();\n let token = token.clone();\n let is_end = is_end.clone();\n let counter = counter.clone();\n local.spawn_local(Box::pin(async move {\n let mut has_err = false;\n let mut result_data_err = ResultData::default();\n let mut rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n loop {\n let client = client.clone();\n match setup_http3(&client, &mut rng).await {\n Ok((connection_time, (h3_connection, send_request))) => {\n let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();\n let http3_driver = spawn_http3_driver(h3_connection, shutdown_rx);\n let futures = (0..n_http_parallel)\n .map(|_| {\n let mut client_state = ClientStateHttp3::new(send_request.clone());\n let client = client.clone();\n let report_tx = report_tx.clone();\n let token = token.clone();\n let is_end = is_end.clone();\n let counter = counter.clone();\n tokio::task::spawn_local(async move {\n let mut result_data = ResultData::default();\n\n let work = async {\n loop {\n if is_counting_tasks\n && counter\n .as_ref()\n .unwrap()\n .fetch_sub(1, Ordering::Relaxed)\n <= 0\n {\n return true;\n }\n let mut res =\n client.work_http3(&mut client_state).await;\n let is_cancel = is_cancel_error(&res)\n || is_end.load(Ordering::Relaxed);\n let is_reconnect = is_h3_error(&res);\n set_connection_time(&mut res, connection_time);\n\n result_data.push(res);\n\n if is_cancel || is_reconnect {\n return is_cancel;\n }\n }\n };\n\n let is_cancel = tokio::select! {\n is_cancel = work => {\n is_cancel\n }\n _ = token.cancelled() => {\n result_data.push(Err(ClientError::Deadline));\n true\n }\n };\n\n report_tx.send(result_data).unwrap();\n is_cancel\n })\n })\n .collect::>();\n\n let mut connection_gone = false;\n for f in futures {\n match f.await {\n Ok(true) => {\n // All works done\n connection_gone = true;\n }\n Err(_) => {\n // Unexpected\n connection_gone = true;\n }\n _ => {}\n }\n }\n\n if connection_gone {\n let _ = shutdown_tx.send(());\n let _ = http3_driver.await;\n break;\n }\n }\n Err(err) => {\n has_err = true;\n result_data_err.push(Err(err));\n if is_end.load(Ordering::Relaxed)\n || (is_counting_tasks\n && counter.as_ref().unwrap().fetch_sub(1, Ordering::Relaxed) <= 0)\n {\n break;\n }\n }\n }\n }\n if has_err {\n report_tx.send(result_data_err).unwrap();\n }\n }));\n }\n rt.block_on(local);\n}\n\n/// Work function for HTTP3 client that generates `n_tasks` tasks.\npub async fn work(\n client: Arc,\n report_tx: kanal::Sender>,\n n_tasks: usize,\n n_connections: usize,\n n_http2_parallel: usize,\n) {\n let (tx, rx) = kanal::unbounded::>();\n let rx = rx.to_async();\n\n let n_tasks_emitter = async move {\n for _ in 0..n_tasks {\n tx.send(None)?\n }\n drop(tx);\n Ok::<(), kanal::SendError>(())\n };\n let futures =\n parallel_work_http3(n_connections, n_http2_parallel, rx, report_tx, client, None).await;\n n_tasks_emitter.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n}\n\n/// n tasks by m workers limit to qps works in a second\npub async fn work_with_qps(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n n_tasks: usize,\n n_connections: usize,\n n_http_parallel: usize,\n) {\n let (tx, rx) = kanal::unbounded::>();\n\n let work_queue = async move {\n match query_limit {\n QueryLimit::Qps(qps) => {\n let start = std::time::Instant::now();\n for i in 0..n_tasks {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n tx.send(None)?;\n }\n }\n QueryLimit::Burst(duration, rate) => {\n let mut n = 0;\n // Handle via rate till n_tasks out of bound\n while n + rate < n_tasks {\n tokio::time::sleep(duration).await;\n for _ in 0..rate {\n tx.send(None)?;\n }\n n += rate;\n }\n // Handle the remaining tasks\n if n_tasks > n {\n tokio::time::sleep(duration).await;\n for _ in 0..n_tasks - n {\n tx.send(None)?;\n }\n }\n }\n }\n // tx gone\n drop(tx);\n Ok::<(), kanal::SendError>(())\n };\n\n let rx = rx.to_async();\n let futures =\n parallel_work_http3(n_connections, n_http_parallel, rx, report_tx, client, None).await;\n work_queue.await.unwrap();\n for f in futures {\n let _ = f.await;\n }\n}\n\n/// n tasks by m workers limit to qps works in a second with latency correction\npub async fn work_with_qps_latency_correction(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n n_tasks: usize,\n n_connections: usize,\n n_http2_parallel: usize,\n) {\n let (tx, rx) = kanal::unbounded();\n\n let _work_queue = async move {\n match query_limit {\n QueryLimit::Qps(qps) => {\n let start = std::time::Instant::now();\n for i in 0..n_tasks {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let now = std::time::Instant::now();\n tx.send(Some(now))?;\n }\n }\n QueryLimit::Burst(duration, rate) => {\n let mut n = 0;\n // Handle via rate till n_tasks out of bound\n while n + rate < n_tasks {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n for _ in 0..rate {\n tx.send(Some(now))?;\n }\n n += rate;\n }\n // Handle the remaining tasks\n if n_tasks > n {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n for _ in 0..n_tasks - n {\n tx.send(Some(now))?;\n }\n }\n }\n }\n\n // tx gone\n drop(tx);\n Ok::<(), kanal::SendError>(())\n };\n\n let rx = rx.to_async();\n let futures =\n parallel_work_http3(n_connections, n_http2_parallel, rx, report_tx, client, None).await;\n for f in futures {\n let _ = f.await;\n }\n}\n\n/// Run until dead_line by n workers\npub async fn work_until(\n client: Arc,\n report_tx: kanal::Sender>,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http_parallel: usize,\n _wait_ongoing_requests_after_deadline: bool,\n) {\n let (tx, rx) = kanal::bounded_async::>(5000);\n // This emitter is used for H3 to give it unlimited tokens to emit work.\n let cancel_token = tokio_util::sync::CancellationToken::new();\n let emitter_handle = endless_emitter(cancel_token.clone(), tx).await;\n let futures = parallel_work_http3(\n n_connections,\n n_http_parallel,\n rx,\n report_tx.clone(),\n client.clone(),\n Some(dead_line),\n )\n .await;\n for f in futures {\n let _ = f.await;\n }\n // Cancel the emitter when we're done with the futures\n cancel_token.cancel();\n // Wait for the emitter to exit cleanly\n let _ = emitter_handle.await;\n}\n\n/// Run until dead_line by n workers limit to qps works in a second\n#[allow(clippy::too_many_arguments)]\npub async fn work_until_with_qps(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n start: std::time::Instant,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http2_parallel: usize,\n _wait_ongoing_requests_after_deadline: bool,\n) {\n let rx = match query_limit {\n QueryLimit::Qps(qps) => {\n let (tx, rx) = kanal::unbounded::>();\n tokio::spawn(async move {\n for i in 0.. {\n if std::time::Instant::now() > dead_line {\n break;\n }\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let _ = tx.send(None);\n }\n // tx gone\n });\n rx\n }\n QueryLimit::Burst(duration, rate) => {\n let (tx, rx) = kanal::unbounded();\n tokio::spawn(async move {\n // Handle via rate till deadline is reached\n for _ in 0.. {\n if std::time::Instant::now() > dead_line {\n break;\n }\n\n tokio::time::sleep(duration).await;\n for _ in 0..rate {\n let _ = tx.send(None);\n }\n }\n // tx gone\n });\n rx\n }\n };\n let rx = rx.to_async();\n let futures = parallel_work_http3(\n n_connections,\n n_http2_parallel,\n rx,\n report_tx,\n client,\n Some(dead_line),\n )\n .await;\n for f in futures {\n let _ = f.await;\n }\n}\n\n/// Run until dead_line by n workers limit to qps works in a second with latency correction\n#[allow(clippy::too_many_arguments)]\npub async fn work_until_with_qps_latency_correction(\n client: Arc,\n report_tx: kanal::Sender>,\n query_limit: QueryLimit,\n start: std::time::Instant,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http2_parallel: usize,\n _wait_ongoing_requests_after_deadline: bool,\n) {\n let (tx, rx) = kanal::unbounded();\n match query_limit {\n QueryLimit::Qps(qps) => {\n tokio::spawn(async move {\n for i in 0.. {\n tokio::time::sleep_until(\n (start + std::time::Duration::from_secs_f64(i as f64 * 1f64 / qps)).into(),\n )\n .await;\n let now = std::time::Instant::now();\n if now > dead_line {\n break;\n }\n let _ = tx.send(Some(now));\n }\n // tx gone\n });\n }\n QueryLimit::Burst(duration, rate) => {\n tokio::spawn(async move {\n // Handle via rate till deadline is reached\n loop {\n tokio::time::sleep(duration).await;\n let now = std::time::Instant::now();\n if now > dead_line {\n break;\n }\n\n for _ in 0..rate {\n let _ = tx.send(Some(now));\n }\n }\n // tx gone\n });\n }\n };\n\n let rx = rx.to_async();\n let futures = parallel_work_http3(\n n_connections,\n n_http2_parallel,\n rx,\n report_tx,\n client,\n Some(dead_line),\n )\n .await;\n for f in futures {\n let _ = f.await;\n }\n}\n\n#[cfg(feature = \"http3\")]\nasync fn endless_emitter(\n cancellation_token: tokio_util::sync::CancellationToken,\n tx: kanal::AsyncSender>,\n) -> tokio::task::JoinHandle<()> {\n tokio::spawn(async move {\n loop {\n tokio::select! {\n _ = cancellation_token.cancelled() => {\n break;\n }\n _ = async {\n // As we our `work_http2_once` function is limited by the number of `tx` we send, but we only\n // want to stop when our semaphore is closed, just dump unlimited `Nones` into the tx to un-constrain it\n let _ = tx.send(None).await;\n } => {}\n }\n }\n })\n}\n\npub mod fast {\n use std::sync::{\n Arc,\n atomic::{AtomicBool, AtomicIsize, Ordering},\n };\n\n use crate::{\n client::Client, client_h3::http3_connection_fast_work_until, result_data::ResultData,\n };\n\n /// Run n tasks by m workers\n pub async fn work(\n client: Arc,\n report_tx: kanal::Sender,\n n_tasks: usize,\n n_connections: usize,\n n_http_parallel: usize,\n ) {\n let counter = Arc::new(AtomicIsize::new(n_tasks as isize));\n let num_threads = num_cpus::get_physical();\n let connections = (0..num_threads).filter_map(|i| {\n let num_connection = n_connections / num_threads\n + (if (n_connections % num_threads) > i {\n 1\n } else {\n 0\n });\n if num_connection > 0 {\n Some(num_connection)\n } else {\n None\n }\n });\n let token = tokio_util::sync::CancellationToken::new();\n let handles = connections\n .map(|num_connections| {\n let report_tx = report_tx.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n let token = token.clone();\n let counter = counter.clone();\n // will let is_end just stay false permanently\n let is_end = Arc::new(AtomicBool::new(false));\n std::thread::spawn(move || {\n http3_connection_fast_work_until(\n num_connections,\n n_http_parallel,\n report_tx,\n client,\n token,\n Some(counter),\n is_end,\n rt,\n )\n })\n })\n .collect::>();\n tokio::spawn(async move {\n tokio::signal::ctrl_c().await.unwrap();\n token.cancel();\n });\n\n tokio::task::block_in_place(|| {\n for handle in handles {\n let _ = handle.join();\n }\n });\n }\n\n /// Run until dead_line by n workers\n pub async fn work_until(\n client: Arc,\n report_tx: kanal::Sender,\n dead_line: std::time::Instant,\n n_connections: usize,\n n_http_parallel: usize,\n wait_ongoing_requests_after_deadline: bool,\n ) {\n let num_threads = num_cpus::get_physical();\n\n let is_end = Arc::new(AtomicBool::new(false));\n let connections = (0..num_threads).filter_map(|i| {\n let num_connection = n_connections / num_threads\n + (if (n_connections % num_threads) > i {\n 1\n } else {\n 0\n });\n if num_connection > 0 {\n Some(num_connection)\n } else {\n None\n }\n });\n let token = tokio_util::sync::CancellationToken::new();\n let handles = connections\n .map(|num_connections| {\n let report_tx = report_tx.clone();\n let client = client.clone();\n let rt = tokio::runtime::Builder::new_current_thread()\n .enable_all()\n .build()\n .unwrap();\n let token = token.clone();\n let is_end = is_end.clone();\n std::thread::spawn(move || {\n http3_connection_fast_work_until(\n num_connections,\n n_http_parallel,\n report_tx,\n client,\n token,\n None,\n is_end,\n rt,\n )\n })\n })\n .collect::>();\n tokio::select! {\n _ = tokio::time::sleep_until(dead_line.into()) => {\n }\n _ = tokio::signal::ctrl_c() => {\n }\n }\n\n is_end.store(true, Ordering::Relaxed);\n\n if !wait_ongoing_requests_after_deadline {\n token.cancel();\n }\n tokio::task::block_in_place(|| {\n for handle in handles {\n let _ = handle.join();\n }\n });\n }\n}\n" }, { "path": "src/curl_compat.rs", "content": "//! Curl compatibility utilities\nuse std::str::FromStr;\n\npub struct Form {\n pub boundary: String,\n pub parts: Vec,\n}\n\npub struct FormPart {\n pub name: String,\n pub filename: Option,\n pub content_type: Option,\n pub data: Vec,\n}\n\nimpl Form {\n pub fn new() -> Self {\n Self {\n boundary: Self::generate_boundary(),\n parts: Vec::new(),\n }\n }\n\n pub fn add_part(&mut self, part: FormPart) {\n self.parts.push(part);\n }\n\n pub fn content_type(&self) -> String {\n format!(\"multipart/form-data; boundary={}\", self.boundary)\n }\n\n pub fn body(&self) -> Vec {\n let mut body = Vec::new();\n\n for part in &self.parts {\n // Add boundary separator\n body.extend_from_slice(b\"--\");\n body.extend_from_slice(self.boundary.as_bytes());\n body.extend_from_slice(b\"\\r\\n\");\n\n // Add Content-Disposition header\n body.extend_from_slice(b\"Content-Disposition: form-data; name=\\\"\");\n body.extend_from_slice(part.name.as_bytes());\n body.extend_from_slice(b\"\\\"\");\n\n // Add filename if present\n if let Some(filename) = &part.filename {\n body.extend_from_slice(b\"; filename=\\\"\");\n body.extend_from_slice(filename.as_bytes());\n body.extend_from_slice(b\"\\\"\");\n }\n body.extend_from_slice(b\"\\r\\n\");\n\n // Add Content-Type header if present\n if let Some(content_type) = &part.content_type {\n body.extend_from_slice(b\"Content-Type: \");\n body.extend_from_slice(content_type.as_bytes());\n body.extend_from_slice(b\"\\r\\n\");\n }\n\n // Empty line before data\n body.extend_from_slice(b\"\\r\\n\");\n\n // Add the actual data\n body.extend_from_slice(&part.data);\n body.extend_from_slice(b\"\\r\\n\");\n }\n\n // Add final boundary\n body.extend_from_slice(b\"--\");\n body.extend_from_slice(self.boundary.as_bytes());\n body.extend_from_slice(b\"--\\r\\n\");\n\n body\n }\n fn generate_boundary() -> String {\n let random_bytes: [u8; 16] = rand::random();\n\n // Convert to hex string manually to avoid external hex dependency\n let hex_string = random_bytes\n .iter()\n .map(|b| format!(\"{b:02x}\"))\n .collect::();\n\n format!(\"----formdata-oha-{hex_string}\")\n }\n}\n\nimpl FromStr for FormPart {\n type Err = anyhow::Error;\n\n /// Parse curl's -F format string\n /// Supports formats like:\n /// - `name=value`\n /// - `name=@filename` (file upload with filename)\n /// - `name= Result {\n // Split on first '=' to separate name from value/options\n let (name, rest) = s\n .split_once('=')\n .ok_or_else(|| anyhow::anyhow!(\"Invalid form format: missing '=' in '{}'\", s))?;\n\n let name = name.to_string();\n\n // Parse the value part which may contain semicolon-separated options\n let parts: Vec<&str> = rest.split(';').collect();\n let value_part = parts[0];\n\n let mut filename = None;\n let mut content_type = None;\n let data;\n\n // Check if this is a file upload (@filename or {\n filename = Some(value.trim().to_string());\n }\n \"type\" => {\n content_type = Some(value.trim().to_string());\n }\n _ => {\n // Ignore unknown options for compatibility\n }\n }\n }\n }\n\n Ok(FormPart {\n name,\n filename,\n content_type,\n data,\n })\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_parse_simple_field() {\n let part: FormPart = \"name=value\".parse().unwrap();\n assert_eq!(part.name, \"name\");\n assert_eq!(part.data, b\"value\");\n assert_eq!(part.filename, None);\n assert_eq!(part.content_type, None);\n }\n\n #[test]\n fn test_parse_field_with_filename() {\n let part: FormPart = \"upload=data;filename=test.txt\".parse().unwrap();\n assert_eq!(part.name, \"upload\");\n assert_eq!(part.data, b\"data\");\n assert_eq!(part.filename, Some(\"test.txt\".to_string()));\n assert_eq!(part.content_type, None);\n }\n\n #[test]\n fn test_parse_field_with_type() {\n let part: FormPart = \"data=content;type=text/plain\".parse().unwrap();\n assert_eq!(part.name, \"data\");\n assert_eq!(part.data, b\"content\");\n assert_eq!(part.filename, None);\n assert_eq!(part.content_type, Some(\"text/plain\".to_string()));\n }\n\n #[test]\n fn test_parse_field_with_filename_and_type() {\n let part: FormPart = \"file=content;filename=test.txt;type=text/plain\"\n .parse()\n .unwrap();\n assert_eq!(part.name, \"file\");\n assert_eq!(part.data, b\"content\");\n assert_eq!(part.filename, Some(\"test.txt\".to_string()));\n assert_eq!(part.content_type, Some(\"text/plain\".to_string()));\n }\n\n #[test]\n fn test_parse_invalid_format() {\n let result: Result = \"invalid\".parse();\n assert!(result.is_err());\n }\n\n #[test]\n fn test_parse_file_upload() {\n // Create a temporary file for testing\n let temp_dir = std::env::temp_dir();\n let test_file = temp_dir.join(\"test_form_upload.txt\");\n std::fs::write(&test_file, b\"test file content\").unwrap();\n\n let form_str = format!(\"upload=@{}\", test_file.display());\n let part: FormPart = form_str.parse().unwrap();\n\n assert_eq!(part.name, \"upload\");\n assert_eq!(part.data, b\"test file content\");\n assert_eq!(part.filename, Some(\"test_form_upload.txt\".to_string()));\n assert_eq!(part.content_type, None);\n\n // Clean up\n std::fs::remove_file(&test_file).ok();\n }\n\n #[test]\n fn test_parse_file_upload_without_filename() {\n // Create a temporary file for testing\n let temp_dir = std::env::temp_dir();\n let test_file = temp_dir.join(\"test_form_upload_no_filename.txt\");\n std::fs::write(&test_file, b\"test file content without filename\").unwrap();\n\n let form_str = format!(\"upload=<{}\", test_file.display());\n let part: FormPart = form_str.parse().unwrap();\n\n assert_eq!(part.name, \"upload\");\n assert_eq!(part.data, b\"test file content without filename\");\n assert_eq!(part.filename, None); // No filename set for '<' format\n assert_eq!(part.content_type, None);\n\n // Clean up\n std::fs::remove_file(&test_file).ok();\n }\n\n #[test]\n fn test_form_creation_and_body_generation() {\n let mut form = Form::new();\n\n // Add a simple text field\n let text_part: FormPart = \"name=John\".parse().unwrap();\n form.add_part(text_part);\n\n // Add a field with filename\n let file_part: FormPart = \"file=content;filename=test.txt;type=text/plain\"\n .parse()\n .unwrap();\n form.add_part(file_part);\n\n let body = form.body();\n let body_str = String::from_utf8_lossy(&body);\n\n // Check that boundary is present\n assert!(body_str.contains(&format!(\"--{}\", form.boundary)));\n\n // Check Content-Disposition headers\n assert!(body_str.contains(\"Content-Disposition: form-data; name=\\\"name\\\"\"));\n assert!(\n body_str\n .contains(\"Content-Disposition: form-data; name=\\\"file\\\"; filename=\\\"test.txt\\\"\")\n );\n\n // Check Content-Type header\n assert!(body_str.contains(\"Content-Type: text/plain\"));\n\n // Check data content\n assert!(body_str.contains(\"John\"));\n assert!(body_str.contains(\"content\"));\n\n // Check final boundary\n assert!(body_str.ends_with(&format!(\"--{}--\\r\\n\", form.boundary)));\n }\n\n #[test]\n fn test_form_content_type() {\n let form = Form::new();\n let content_type = form.content_type();\n\n assert!(content_type.starts_with(\"multipart/form-data; boundary=\"));\n assert!(content_type.contains(&form.boundary));\n }\n\n #[test]\n fn test_empty_form_body() {\n let form = Form::new();\n let body = form.body();\n let body_str = String::from_utf8_lossy(&body);\n\n // Should only contain final boundary for empty form\n assert_eq!(body_str, format!(\"--{}--\\r\\n\", form.boundary));\n }\n\n #[test]\n fn test_form_with_file_upload() {\n // Create a temporary file for testing\n let temp_dir = std::env::temp_dir();\n let test_file = temp_dir.join(\"form_test_upload.txt\");\n std::fs::write(&test_file, b\"file content for form\").unwrap();\n\n let mut form = Form::new();\n\n // Parse and add file upload part\n let form_str = format!(\"upload=@{}\", test_file.display());\n let file_part: FormPart = form_str.parse().unwrap();\n form.add_part(file_part);\n\n let body = form.body();\n let body_str = String::from_utf8_lossy(&body);\n\n // Check file upload formatting\n assert!(body_str.contains(\n \"Content-Disposition: form-data; name=\\\"upload\\\"; filename=\\\"form_test_upload.txt\\\"\"\n ));\n assert!(body_str.contains(\"file content for form\"));\n\n // Clean up\n std::fs::remove_file(&test_file).ok();\n }\n\n #[test]\n fn test_boundary_generation_is_random() {\n let form1 = Form::new();\n let form2 = Form::new();\n\n // Boundaries should be different for different forms\n assert_ne!(form1.boundary, form2.boundary);\n\n // Boundaries should follow the expected format\n assert!(form1.boundary.starts_with(\"----formdata-oha-\"));\n assert!(form2.boundary.starts_with(\"----formdata-oha-\"));\n\n // Boundaries should have the expected length (prefix + 32 hex chars)\n assert_eq!(form1.boundary.len(), \"----formdata-oha-\".len() + 32);\n assert_eq!(form2.boundary.len(), \"----formdata-oha-\".len() + 32);\n }\n}\n" }, { "path": "src/db.rs", "content": "use rusqlite::Connection;\n\nuse crate::client::{Client, RequestResult};\n\nfn create_db(conn: &Connection) -> Result {\n conn.execute(\n \"CREATE TABLE IF NOT EXISTS oha (\n url TEXT NOT NULL,\n start REAL NOT NULL,\n start_latency_correction REAL,\n end REAL NOT NULL,\n duration REAL NOT NULL,\n status INTEGER NOT NULL,\n len_bytes INTEGER NOT NULL,\n run INTEGER NOT NULL\n )\",\n (),\n )\n}\n\npub fn store(\n client: &Client,\n db_url: &str,\n start: std::time::Instant,\n request_records: &[RequestResult],\n run: i64,\n) -> Result {\n let mut conn = Connection::open(db_url)?;\n create_db(&conn)?;\n\n let t = conn.transaction()?;\n let mut affected_rows = 0;\n\n for request in request_records {\n let req = client.generate_request(&mut request.rng.clone()).unwrap().1;\n let url = req.uri();\n affected_rows += t.execute(\n \"INSERT INTO oha (url, start, start_latency_correction, end, duration, status, len_bytes, run) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)\",\n (\n url.to_string(),\n (request.start - start).as_secs_f64(),\n request.start_latency_correction.map(|d| (d - start).as_secs_f64()),\n (request.end - start).as_secs_f64(),\n request.duration().as_secs_f64(),\n request.status.as_u16() as i64,\n request.len_bytes as i64,\n run ,\n ),\n )?;\n }\n\n t.commit()?;\n\n Ok(affected_rows)\n}\n\n#[cfg(test)]\nmod test_db {\n use rand::SeedableRng;\n\n use super::*;\n\n #[test]\n fn test_store() {\n let run = std::time::SystemTime::now()\n .duration_since(std::time::UNIX_EPOCH)\n .unwrap()\n .as_secs() as i64;\n let start = std::time::Instant::now();\n let test_val = RequestResult {\n rng: SeedableRng::seed_from_u64(0),\n status: hyper::StatusCode::OK,\n len_bytes: 100,\n start_latency_correction: None,\n start: std::time::Instant::now(),\n connection_time: None,\n first_byte: None,\n end: std::time::Instant::now(),\n };\n let test_vec = vec![test_val.clone(), test_val.clone()];\n let client = Client::default();\n let result = store(&client, \":memory:\", start, &test_vec, run);\n assert_eq!(result.unwrap(), 2);\n }\n}\n" }, { "path": "src/histogram.rs", "content": "pub fn histogram(values: &[f64], bins: usize) -> Vec<(f64, usize)> {\n assert!(bins >= 2);\n let mut bucket: Vec = vec![0; bins];\n let min = values.iter().collect::().min();\n let max = values.iter().collect::().max();\n let step = (max - min) / (bins - 1) as f64;\n\n for &v in values {\n let i = std::cmp::min(((v - min) / step).ceil() as usize, bins - 1);\n bucket[i] += 1;\n }\n\n bucket\n .into_iter()\n .enumerate()\n .map(|(i, v)| (min + step * i as f64, v))\n .collect()\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_histogram() {\n let values1: [f64; 10] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];\n assert_eq!(\n histogram(&values1, 10),\n vec![\n (1.0, 1),\n (2.0, 1),\n (3.0, 1),\n (4.0, 1),\n (5.0, 1),\n (6.0, 1),\n (7.0, 1),\n (8.0, 1),\n (9.0, 1),\n (10.0, 1)\n ]\n );\n assert_eq!(\n histogram(&values1, 4),\n vec![(1.0, 1), (4.0, 3), (7.0, 3), (10.0, 3)]\n );\n assert_eq!(\n histogram(&values1, 17),\n vec![\n (1.0, 1),\n (1.5625, 0),\n (2.125, 1),\n (2.6875, 0),\n (3.25, 1),\n (3.8125, 0),\n (4.375, 1),\n (4.9375, 0),\n (5.5, 1),\n (6.0625, 1),\n (6.625, 0),\n (7.1875, 1),\n (7.75, 0),\n (8.3125, 1),\n (8.875, 0),\n (9.4375, 1),\n (10.0, 1)\n ]\n );\n\n let values2: [f64; 10] = [1.0, 1.0, 1.0, 1.0, 1.0, 10.0, 10.0, 10.0, 10.0, 10.0];\n assert_eq!(\n histogram(&values2, 10),\n vec![\n (1.0, 5),\n (2.0, 0),\n (3.0, 0),\n (4.0, 0),\n (5.0, 0),\n (6.0, 0),\n (7.0, 0),\n (8.0, 0),\n (9.0, 0),\n (10.0, 5)\n ]\n );\n assert_eq!(histogram(&values2, 2), vec![(1.0, 5), (10.0, 5)]);\n }\n}\n" }, { "path": "src/lib.rs", "content": "use anyhow::Context;\nuse aws_auth::AwsSignatureConfig;\nuse bytes::Bytes;\nuse clap::Parser;\nuse crossterm::tty::IsTty;\nuse hickory_resolver::config::{ResolverConfig, ResolverOpts};\nuse humantime::Duration;\nuse hyper::{\n HeaderMap,\n http::{self, header::HeaderName, header::HeaderValue},\n};\nuse printer::{PrintConfig, PrintMode};\nuse rand_regex::Regex;\nuse ratatui::crossterm;\nuse result_data::ResultData;\nuse std::{\n env,\n fs::File,\n io::{BufRead, BufReader, Read},\n path::{Path, PathBuf},\n pin::Pin,\n sync::Arc,\n};\nuse timescale::TimeScale;\nuse url::Url;\nuse url_generator::UrlGenerator;\n\nmod aws_auth;\nmod cli;\nmod client;\n#[cfg(feature = \"http3\")]\nmod client_h3;\nmod curl_compat;\nmod db;\nmod histogram;\nmod monitor;\nmod pcg64si;\nmod printer;\nmod request_generator;\nmod result_data;\nmod timescale;\nmod tls_config;\nmod url_generator;\n\n#[cfg(not(target_env = \"msvc\"))]\nuse tikv_jemallocator::Jemalloc;\n\nuse crate::{\n cli::{ConnectToEntry, parse_header},\n request_generator::{BodyGenerator, Proxy, RequestGenerator},\n};\n\n#[cfg(not(target_env = \"msvc\"))]\n#[global_allocator]\nstatic GLOBAL: Jemalloc = Jemalloc;\n\n#[derive(Parser)]\n#[command(version, about, long_about = None)]\n#[command(arg_required_else_help(true))]\n#[command(styles = clap_cargo::style::CLAP_STYLING)]\npub struct Opts {\n #[arg(help = \"Target URL or file with multiple URLs.\")]\n url: Option,\n\n #[arg(long = \"completions\", hide = true)]\n pub completions: Option,\n\n #[arg(\n help = \"Number of requests to run. Accepts plain numbers or suffixes: k = 1,000, m = 1,000,000 (e.g. 10k, 1m).\",\n short = 'n',\n default_value = \"200\",\n conflicts_with = \"duration\",\n value_parser = cli::parse_n_requests\n )]\n n_requests: usize,\n #[arg(\n help = \"Number of connections to run concurrently. You may should increase limit to number of open files for larger `-c`.\",\n short = 'c',\n default_value = \"50\"\n )]\n n_connections: usize,\n #[arg(\n help = \"Number of parallel requests to send on HTTP/2. `oha` will run c * p concurrent workers in total.\",\n short = 'p',\n default_value = \"1\"\n )]\n n_http2_parallel: usize,\n #[arg(\n help = \"Duration of application to send requests.\nOn HTTP/1, When the duration is reached, ongoing requests are aborted and counted as \\\"aborted due to deadline\\\"\nYou can change this behavior with `-w` option.\nCurrently, on HTTP/2, When the duration is reached, ongoing requests are waited. `-w` option is ignored.\nExamples: -z 10s -z 3m.\",\n short = 'z',\n conflicts_with = \"n_requests\"\n )]\n duration: Option,\n #[arg(\n help = \"When the duration is reached, ongoing requests are waited\",\n short,\n long,\n default_value = \"false\",\n requires = \"duration\"\n )]\n wait_ongoing_requests_after_deadline: bool,\n #[arg(help = \"Rate limit for all, in queries per second (QPS)\", short = 'q', conflicts_with_all = [\"burst_duration\", \"burst_requests\"])]\n query_per_second: Option,\n #[arg(\n help = \"Introduce delay between a predefined number of requests.\nNote: If qps is specified, burst will be ignored\",\n long = \"burst-delay\",\n requires = \"burst_requests\",\n conflicts_with = \"query_per_second\"\n )]\n burst_duration: Option,\n #[arg(\n help = \"Rates of requests for burst. Default is 1\nNote: If qps is specified, burst will be ignored\",\n long = \"burst-rate\",\n requires = \"burst_duration\",\n conflicts_with = \"query_per_second\"\n )]\n burst_requests: Option,\n\n #[arg(\n help = \"Generate URL by rand_regex crate but dot is disabled for each query e.g. http://127.0.0.1/[a-z][a-z][0-9]. Currently dynamic scheme, host and port with keep-alive do not work well. See https://docs.rs/rand_regex/latest/rand_regex/struct.Regex.html for details of syntax.\",\n default_value = \"false\",\n long\n )]\n rand_regex_url: bool,\n\n #[arg(\n help = \"Read the URLs to query from a file\",\n default_value = \"false\",\n long\n )]\n urls_from_file: bool,\n\n #[arg(\n help = \"A parameter for the '--rand-regex-url'. The max_repeat parameter gives the maximum extra repeat counts the x*, x+ and x{n,} operators will become.\",\n default_value = \"4\",\n long,\n requires = \"rand_regex_url\"\n )]\n max_repeat: u32,\n #[arg(\n help = \"Dump target Urls times to debug --rand-regex-url\",\n long\n )]\n dump_urls: Option,\n #[arg(\n help = \"Correct latency to avoid coordinated omission problem. It's ignored if -q is not set.\",\n long = \"latency-correction\"\n )]\n latency_correction: bool,\n #[arg(help = \"No realtime tui\", long = \"no-tui\")]\n no_tui: bool,\n #[arg(help = \"Frame per second for tui.\", default_value = \"16\", long = \"fps\")]\n fps: usize,\n #[arg(\n help = \"HTTP method\",\n short = 'm',\n long = \"method\",\n default_value = \"GET\"\n )]\n method: http::Method,\n #[arg(help = \"Custom HTTP header. Examples: -H \\\"foo: bar\\\"\", short = 'H', value_parser = parse_header)]\n headers: Vec<(HeaderName, HeaderValue)>,\n #[arg(\n help = \"Custom Proxy HTTP header. Examples: --proxy-header \\\"foo: bar\\\"\",\n long = \"proxy-header\",\n value_parser = parse_header\n )]\n proxy_headers: Vec<(HeaderName, HeaderValue)>,\n #[arg(help = \"Timeout for each request. Default to infinite.\", short = 't')]\n timeout: Option,\n #[arg(\n help = \"Timeout for establishing a new connection. Default to 5s.\",\n long = \"connect-timeout\",\n default_value = \"5s\"\n )]\n connect_timeout: humantime::Duration,\n #[arg(help = \"HTTP Accept Header.\", short = 'A')]\n accept_header: Option,\n #[arg(help = \"HTTP request body.\", short = 'd', conflicts_with_all = [\"body_path\", \"body_path_lines\", \"form\"])]\n body_string: Option,\n #[arg(help = \"HTTP request body from file.\", short = 'D', conflicts_with_all = [\"body_string\", \"body_path_lines\", \"form\"])]\n body_path: Option,\n #[arg(help = \"HTTP request body from file line by line.\", short = 'Z', conflicts_with_all = [\"body_string\", \"body_path\", \"form\"])]\n body_path_lines: Option,\n #[arg(\n help = \"Specify HTTP multipart POST data (curl compatible). Examples: -F 'name=value' -F 'file=@path/to/file'\",\n short = 'F',\n long = \"form\",\n conflicts_with_all = [\"body_string\", \"body_path\", \"body_path_lines\"]\n )]\n form: Vec,\n #[arg(help = \"Content-Type.\", short = 'T')]\n content_type: Option,\n #[arg(\n help = \"Basic authentication (username:password), or AWS credentials (access_key:secret_key)\",\n short = 'a'\n )]\n basic_auth: Option,\n #[arg(help = \"AWS session token\", long = \"aws-session\")]\n aws_session: Option,\n #[arg(\n help = \"AWS SigV4 signing params (format: aws:amz:region:service)\",\n long = \"aws-sigv4\"\n )]\n aws_sigv4: Option,\n #[arg(help = \"HTTP proxy\", short = 'x')]\n proxy: Option,\n #[arg(\n help = \"HTTP version to connect to proxy. Available values 0.9, 1.0, 1.1, 2.\",\n long = \"proxy-http-version\"\n )]\n proxy_http_version: Option,\n #[arg(\n help = \"Use HTTP/2 to connect to proxy. Shorthand for --proxy-http-version=2\",\n long = \"proxy-http2\"\n )]\n proxy_http2: bool,\n #[arg(\n help = \"HTTP version. Available values 0.9, 1.0, 1.1, 2, 3\",\n long = \"http-version\"\n )]\n http_version: Option,\n #[arg(help = \"Use HTTP/2. Shorthand for --http-version=2\", long = \"http2\")]\n http2: bool,\n #[arg(help = \"HTTP Host header\", long = \"host\")]\n host: Option,\n #[arg(help = \"Disable compression.\", long = \"disable-compression\")]\n disable_compression: bool,\n #[arg(\n help = \"Limit for number of Redirect. Set 0 for no redirection. Redirection isn't supported for HTTP/2.\",\n default_value = \"10\",\n short = 'r',\n long = \"redirect\"\n )]\n redirect: usize,\n #[arg(\n help = \"Disable keep-alive, prevents re-use of TCP connections between different HTTP requests. This isn't supported for HTTP/2.\",\n long = \"disable-keepalive\"\n )]\n disable_keepalive: bool,\n #[arg(\n help = \"*Not* perform a DNS lookup at beginning to cache it\",\n long = \"no-pre-lookup\",\n default_value = \"false\"\n )]\n no_pre_lookup: bool,\n #[arg(help = \"Lookup only ipv6.\", long = \"ipv6\")]\n ipv6: bool,\n #[arg(help = \"Lookup only ipv4.\", long = \"ipv4\")]\n ipv4: bool,\n #[arg(\n help = \"(TLS) Use the specified certificate file to verify the peer. Native certificate store is used even if this argument is specified.\",\n long\n )]\n cacert: Option,\n #[arg(\n help = \"(TLS) Use the specified client certificate file. --key must be also specified\",\n long,\n requires = \"key\"\n )]\n cert: Option,\n #[arg(\n help = \"(TLS) Use the specified client key file. --cert must be also specified\",\n long,\n requires = \"cert\"\n )]\n key: Option,\n #[arg(help = \"Accept invalid certs.\", long = \"insecure\")]\n insecure: bool,\n #[arg(\n help = \"Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'\nNote: if used several times for the same host:port:target_host:target_port, a random choice is made\",\n long = \"connect-to\"\n )]\n connect_to: Vec,\n #[arg(\n help = \"Disable the color scheme.\",\n alias = \"disable-color\",\n long = \"no-color\",\n env = \"NO_COLOR\"\n )]\n no_color: bool,\n #[cfg(unix)]\n #[arg(\n help = \"Connect to a unix socket instead of the domain in the URL. Only for non-HTTPS URLs.\",\n long = \"unix-socket\",\n group = \"socket-type\"\n )]\n unix_socket: Option,\n #[cfg(feature = \"vsock\")]\n #[arg(\n help = \"Connect to a VSOCK socket using 'cid:port' instead of the domain in the URL. Only for non-HTTPS URLs.\",\n long = \"vsock-addr\",\n value_parser = cli::parse_vsock_addr,\n group = \"socket-type\"\n )]\n vsock_addr: Option,\n #[arg(\n help = \"Include a response status code successful or not successful breakdown for the time histogram and distribution statistics\",\n long = \"stats-success-breakdown\"\n )]\n stats_success_breakdown: bool,\n #[arg(\n help = \"Write succeeded requests to sqlite database url E.G test.db\",\n long = \"db-url\"\n )]\n db_url: Option,\n #[arg(\n long,\n help = \"Perform a single request and dump the request and response\"\n )]\n debug: bool,\n #[arg(\n help = \"Output file to write the results to. If not specified, results are written to stdout.\",\n long,\n short\n )]\n output: Option,\n #[arg(help = \"Output format\", long, default_value = \"text\")]\n output_format: Option,\n #[arg(\n help = \"Time unit to be used. If not specified, the time unit is determined automatically. This option affects only text format.\",\n long,\n short = 'u'\n )]\n time_unit: Option,\n}\n\npub async fn run(mut opts: Opts) -> anyhow::Result<()> {\n let work_mode = opts.work_mode();\n let url = opts.url.expect(\"URL is required\");\n\n // Parse AWS credentials from basic auth if AWS signing is requested\n let aws_config = if let Some(signing_params) = opts.aws_sigv4 {\n if let Some(auth) = &opts.basic_auth {\n let parts: Vec<&str> = auth.split(':').collect();\n if parts.len() != 2 {\n anyhow::bail!(\"Invalid AWS credentials format. Expected access_key:secret_key\");\n }\n let access_key = parts[0];\n let secret_key = parts[1];\n let session_token = opts.aws_session.take();\n Some(AwsSignatureConfig::new(\n access_key,\n secret_key,\n &signing_params,\n session_token,\n )?)\n } else {\n anyhow::bail!(\"AWS credentials (--auth) required when using --aws-sigv4\");\n }\n } else {\n None\n };\n\n let parse_http_version = |is_http2: bool, version: Option<&str>| match (is_http2, version) {\n (true, Some(_)) => anyhow::bail!(\"--http2 and --http-version are exclusive\"),\n (true, None) => Ok(http::Version::HTTP_2),\n (false, Some(http_version)) => match http_version.trim() {\n \"0.9\" => Ok(http::Version::HTTP_09),\n \"1.0\" => Ok(http::Version::HTTP_10),\n \"1.1\" => Ok(http::Version::HTTP_11),\n \"2.0\" | \"2\" => Ok(http::Version::HTTP_2),\n #[cfg(feature = \"http3\")]\n \"3.0\" | \"3\" => Ok(http::Version::HTTP_3),\n #[cfg(not(feature = \"http3\"))]\n \"3.0\" | \"3\" => anyhow::bail!(\n \"Your Oha instance has not been built with HTTP/3 support. Try recompiling with the feature enabled.\"\n ),\n _ => anyhow::bail!(\"Unknown HTTP version. Valid versions are 0.9, 1.0, 1.1, 2, 3\"),\n },\n (false, None) => Ok(http::Version::HTTP_11),\n };\n\n let http_version: http::Version = parse_http_version(opts.http2, opts.http_version.as_deref())?;\n let proxy_http_version: http::Version =\n parse_http_version(opts.proxy_http2, opts.proxy_http_version.as_deref())?;\n\n let url_generator = if opts.rand_regex_url {\n // Almost URL has dot in domain, so disable dot in regex for convenience.\n let dot_disabled: String = url\n .chars()\n .map(|c| {\n if c == '.' {\n regex_syntax::escape(\".\")\n } else {\n c.to_string()\n }\n })\n .collect();\n UrlGenerator::new_dynamic(Regex::compile(&dot_disabled, opts.max_repeat)?)\n } else if opts.urls_from_file {\n let path = Path::new(url.as_str());\n let file = File::open(path)?;\n let reader = std::io::BufReader::new(file);\n\n let urls: Vec = reader\n .lines()\n .map_while(Result::ok)\n .filter(|line| !line.trim().is_empty())\n .map(|url_str| Url::parse(&url_str))\n .collect::, _>>()?;\n UrlGenerator::new_multi_static(urls)\n } else {\n UrlGenerator::new_static(Url::parse(&url)?)\n };\n\n if let Some(n) = opts.dump_urls {\n let mut rng = rand::rng();\n for _ in 0..n {\n let url = url_generator.generate(&mut rng)?;\n println!(\"{url}\");\n }\n return Ok(());\n }\n\n let url = url_generator.generate(&mut rand::rng())?;\n\n // Process form data or regular body first\n let has_form_data = !opts.form.is_empty();\n let (body_generator, form_content_type): (BodyGenerator, Option) = if has_form_data {\n let mut form = curl_compat::Form::new();\n\n for form_str in opts.form {\n let part: curl_compat::FormPart = form_str\n .parse()\n .with_context(|| format!(\"Failed to parse form data: {form_str}\"))?;\n form.add_part(part);\n }\n\n let form_body = form.body();\n let content_type = form.content_type();\n\n (BodyGenerator::Static(form_body.into()), Some(content_type))\n } else if let Some(body_string) = opts.body_string {\n (BodyGenerator::Static(body_string.into()), None)\n } else if let Some(body_path) = opts.body_path {\n let mut buf = Vec::new();\n std::fs::File::open(body_path)?.read_to_end(&mut buf)?;\n (BodyGenerator::Static(buf.into()), None)\n } else if let Some(body_path_lines) = opts.body_path_lines {\n let lines = BufReader::new(std::fs::File::open(body_path_lines)?)\n .lines()\n .map_while(Result::ok)\n .map(Bytes::from)\n .collect::>();\n\n (BodyGenerator::Random(lines), None)\n } else {\n (BodyGenerator::Static(Bytes::new()), None)\n };\n\n // Set method to POST if form data is used and method is GET\n let method = if has_form_data && opts.method == http::Method::GET {\n http::Method::POST\n } else {\n opts.method\n };\n\n let headers = {\n let mut headers: http::header::HeaderMap = Default::default();\n\n // Accept all\n headers.insert(\n http::header::ACCEPT,\n http::header::HeaderValue::from_static(\"*/*\"),\n );\n\n // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding\n if !opts.disable_compression {\n headers.insert(\n http::header::ACCEPT_ENCODING,\n http::header::HeaderValue::from_static(\"gzip, compress, deflate, br\"),\n );\n }\n\n // User agent\n headers\n .entry(http::header::USER_AGENT)\n .or_insert(HeaderValue::from_static(concat!(\n \"oha/\",\n env!(\"CARGO_PKG_VERSION\")\n )));\n\n if let Some(h) = opts.accept_header {\n headers.insert(http::header::ACCEPT, HeaderValue::from_bytes(h.as_bytes())?);\n }\n\n if let Some(h) = opts.content_type.or(form_content_type) {\n headers.insert(\n http::header::CONTENT_TYPE,\n HeaderValue::from_bytes(h.as_bytes())?,\n );\n }\n\n if let Some(h) = opts.host {\n headers.insert(http::header::HOST, HeaderValue::from_bytes(h.as_bytes())?);\n }\n\n if let Some(auth) = opts.basic_auth {\n let u_p = auth.splitn(2, ':').collect::>();\n anyhow::ensure!(u_p.len() == 2, anyhow::anyhow!(\"Parse auth\"));\n let mut header_value = b\"Basic \".to_vec();\n {\n use std::io::Write;\n let username = u_p[0];\n let password = if u_p[1].is_empty() {\n None\n } else {\n Some(u_p[1])\n };\n let mut encoder = base64::write::EncoderWriter::new(\n &mut header_value,\n &base64::engine::general_purpose::STANDARD,\n );\n // The unwraps here are fine because Vec::write* is infallible.\n write!(encoder, \"{username}:\").unwrap();\n if let Some(password) = password {\n write!(encoder, \"{password}\").unwrap();\n }\n }\n\n headers.insert(\n http::header::AUTHORIZATION,\n HeaderValue::from_bytes(&header_value)?,\n );\n }\n\n if opts.disable_keepalive && http_version == http::Version::HTTP_11 {\n headers.insert(http::header::CONNECTION, HeaderValue::from_static(\"close\"));\n }\n\n for (k, v) in opts.headers.into_iter() {\n headers.insert(k, v);\n }\n\n headers\n };\n\n let proxy_headers = opts.proxy_headers.into_iter().collect::>();\n\n let ip_strategy = match (opts.ipv4, opts.ipv6) {\n (false, false) => {\n if cfg!(target_os = \"macos\") && (url.host_str() == Some(\"localhost\")) {\n // #784\n // On macOS, localhost resolves to ::1 first, So web servers that bind to localhost tend to listen ipv6 only.\n // So prefer ipv6 on macos for localhost.\n\n hickory_resolver::config::LookupIpStrategy::Ipv6thenIpv4\n } else {\n Default::default()\n }\n }\n (true, false) => hickory_resolver::config::LookupIpStrategy::Ipv4Only,\n (false, true) => hickory_resolver::config::LookupIpStrategy::Ipv6Only,\n (true, true) => hickory_resolver::config::LookupIpStrategy::Ipv4AndIpv6,\n };\n let (config, mut resolver_opts) = system_resolv_conf()?;\n resolver_opts.ip_strategy = ip_strategy;\n let resolver = hickory_resolver::Resolver::builder_with_config(\n config,\n hickory_resolver::name_server::TokioConnectionProvider::default(),\n )\n .with_options(resolver_opts)\n .build();\n let cacert = opts.cacert.as_deref().map(std::fs::read).transpose()?;\n let client_auth = match (opts.cert, opts.key) {\n (Some(cert), Some(key)) => Some((std::fs::read(cert)?, std::fs::read(key)?)),\n (None, None) => None,\n // Not possible because of clap requires\n _ => anyhow::bail!(\"Both --cert and --key must be specified\"),\n };\n\n let url = url.into_owned();\n let client = Arc::new(client::Client {\n request_generator: RequestGenerator {\n url_generator,\n https: url.scheme() == \"https\",\n version: http_version,\n aws_config,\n method,\n headers,\n body_generator,\n http_proxy: if opts.proxy.is_some() && url.scheme() == \"http\" {\n Some(Proxy {\n headers: proxy_headers.clone(),\n version: proxy_http_version,\n })\n } else {\n None\n },\n },\n proxy_http_version,\n proxy_headers,\n dns: client::Dns {\n resolver,\n connect_to: opts.connect_to,\n },\n timeout: opts.timeout.map(|d| d.into()),\n connect_timeout: opts.connect_timeout.into(),\n redirect_limit: opts.redirect,\n disable_keepalive: opts.disable_keepalive,\n proxy_url: opts.proxy,\n #[cfg(unix)]\n unix_socket: opts.unix_socket,\n #[cfg(feature = \"vsock\")]\n vsock_addr: opts.vsock_addr,\n #[cfg(feature = \"rustls\")]\n rustls_configs: tls_config::RuslsConfigs::new(\n opts.insecure,\n cacert.as_deref(),\n client_auth\n .as_ref()\n .map(|(cert, key)| (cert.as_slice(), key.as_slice())),\n ),\n #[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\n native_tls_connectors: tls_config::NativeTlsConnectors::new(\n opts.insecure,\n cacert.as_deref(),\n client_auth\n .as_ref()\n .map(|(cert, key)| (cert.as_slice(), key.as_slice())),\n ),\n });\n\n if !opts.no_pre_lookup {\n client.pre_lookup().await?;\n }\n\n let no_tui = opts.no_tui || !std::io::stdout().is_tty() || opts.debug;\n\n let print_config = {\n let mode = opts.output_format.unwrap_or_default();\n\n let disable_style = opts.no_color || !std::io::stdout().is_tty() || opts.output.is_some();\n\n let output: Box = if let Some(output) = opts.output {\n Box::new(File::create(output)?)\n } else {\n Box::new(std::io::stdout())\n };\n\n PrintConfig {\n mode,\n output,\n disable_style,\n stats_success_breakdown: opts.stats_success_breakdown,\n time_unit: opts.time_unit,\n }\n };\n\n let run = std::time::SystemTime::now()\n .duration_since(std::time::UNIX_EPOCH)?\n .as_secs();\n let start = std::time::Instant::now();\n\n let data_collect_future: Pin>> =\n match work_mode {\n WorkMode::Debug => {\n let mut print_config = print_config;\n client::work_debug(&mut print_config.output, client).await?;\n return Ok(());\n }\n WorkMode::FixedNumber {\n n_requests,\n n_connections,\n n_http2_parallel,\n query_limit: None,\n latency_correction: _,\n } if no_tui => {\n // Use optimized worker of no_tui mode.\n let (result_tx, result_rx) = kanal::unbounded();\n\n client::fast::work(\n client.clone(),\n result_tx,\n n_requests,\n n_connections,\n n_http2_parallel,\n )\n .await;\n\n Box::pin(async move {\n let mut res = ResultData::default();\n for r in result_rx {\n res.merge(r);\n }\n (res, print_config)\n })\n }\n WorkMode::Until {\n duration,\n n_connections,\n n_http2_parallel,\n query_limit: None,\n latency_correction: _,\n wait_ongoing_requests_after_deadline,\n } if no_tui => {\n // Use optimized worker of no_tui mode.\n let (result_tx, result_rx) = kanal::unbounded();\n\n client::fast::work_until(\n client.clone(),\n result_tx,\n start + duration,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n\n Box::pin(async move {\n let mut res = ResultData::default();\n for r in result_rx {\n res.merge(r);\n }\n (res, print_config)\n })\n }\n mode => {\n let (result_tx, result_rx) = kanal::unbounded();\n let data_collector = if no_tui {\n // When `--no-tui` is enabled, just collect all data.\n\n let token = tokio_util::sync::CancellationToken::new();\n let result_rx_ctrl_c = result_rx.clone();\n let token_ctrl_c = token.clone();\n let ctrl_c = tokio::spawn(async move {\n tokio::select! {\n _ = tokio::signal::ctrl_c() => {\n let mut all: ResultData = Default::default();\n let mut buf = Vec::new();\n let _ = result_rx_ctrl_c.drain_into(&mut buf);\n for res in buf {\n all.push(res);\n }\n let _ = printer::print_result(print_config, start, &all, start.elapsed());\n std::process::exit(libc::EXIT_SUCCESS);\n }\n _ = token_ctrl_c.cancelled() => {\n print_config\n }\n\n }\n });\n\n Box::pin(async move {\n token.cancel();\n let config = ctrl_c.await.unwrap();\n let mut all = ResultData::default();\n while let Ok(res) = result_rx.recv() {\n all.push(res);\n }\n (all, config)\n })\n as Pin>>\n } else {\n // Spawn monitor future which draws realtime tui\n let join_handle = tokio::spawn(\n monitor::Monitor {\n print_config,\n end_line: opts\n .duration\n .map(|d| monitor::EndLine::Duration(d.into()))\n .unwrap_or(monitor::EndLine::NumQuery(opts.n_requests)),\n report_receiver: result_rx,\n start,\n fps: opts.fps,\n disable_color: opts.no_color,\n time_unit: opts.time_unit,\n }\n .monitor(),\n );\n\n Box::pin(async { join_handle.await.unwrap().unwrap() })\n as Pin>>\n };\n\n match mode {\n WorkMode::Debug => unreachable!(\"Must be already handled\"),\n WorkMode::FixedNumber {\n n_requests,\n n_connections,\n n_http2_parallel,\n query_limit,\n latency_correction,\n } => {\n if let Some(query_limit) = query_limit {\n if latency_correction {\n client::work_with_qps_latency_correction(\n client.clone(),\n result_tx,\n query_limit,\n n_requests,\n n_connections,\n n_http2_parallel,\n )\n .await;\n } else {\n client::work_with_qps(\n client.clone(),\n result_tx,\n query_limit,\n n_requests,\n n_connections,\n n_http2_parallel,\n )\n .await;\n }\n } else {\n client::work(\n client.clone(),\n result_tx,\n n_requests,\n n_connections,\n n_http2_parallel,\n )\n .await;\n }\n }\n WorkMode::Until {\n duration,\n n_connections,\n n_http2_parallel,\n query_limit,\n latency_correction,\n wait_ongoing_requests_after_deadline,\n } => {\n if let Some(query_limit) = query_limit {\n if latency_correction {\n client::work_until_with_qps_latency_correction(\n client.clone(),\n result_tx,\n query_limit,\n start,\n start + duration,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n } else {\n client::work_until_with_qps(\n client.clone(),\n result_tx,\n query_limit,\n start,\n start + duration,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n }\n } else {\n client::work_until(\n client.clone(),\n result_tx,\n start + duration,\n n_connections,\n n_http2_parallel,\n wait_ongoing_requests_after_deadline,\n )\n .await;\n }\n }\n }\n\n data_collector\n }\n };\n\n let duration = start.elapsed();\n let (res, print_config) = data_collect_future.await;\n\n printer::print_result(print_config, start, &res, duration)?;\n\n if let Some(db_url) = opts.db_url {\n eprintln!(\"Storing results to {db_url}\");\n db::store(&client, &db_url, start, res.success(), run as i64)?;\n }\n\n Ok(())\n}\n\npub(crate) fn system_resolv_conf() -> anyhow::Result<(ResolverConfig, ResolverOpts)> {\n // check if we are running in termux https://github.com/termux/termux-app\n #[cfg(unix)]\n if env::var(\"TERMUX_VERSION\").is_ok() {\n let prefix = env::var(\"PREFIX\")?;\n let path = format!(\"{prefix}/etc/resolv.conf\");\n return match std::fs::read(&path) {\n Ok(conf_data) => hickory_resolver::system_conf::parse_resolv_conf(conf_data)\n .context(format!(\"DNS: failed to parse {path}\")),\n Err(err) => {\n fallback_resolver_config(anyhow::anyhow!(\"DNS: failed to load {path}: {err}\"))\n }\n };\n }\n\n match hickory_resolver::system_conf::read_system_conf() {\n Ok(conf) => Ok(conf),\n Err(err) => fallback_resolver_config(anyhow::anyhow!(\n \"DNS: failed to load /etc/resolv.conf: {err}\"\n )),\n }\n}\n\nfn fallback_resolver_config(err: anyhow::Error) -> anyhow::Result<(ResolverConfig, ResolverOpts)> {\n // Notify the user that we had to fall back to a default resolver configuration.\n eprintln!(\"{err}\");\n\n let config = ResolverConfig::default();\n let opts = ResolverOpts::default();\n Ok((config, opts))\n}\n\nenum WorkMode {\n Debug,\n FixedNumber {\n n_requests: usize,\n n_connections: usize,\n n_http2_parallel: usize,\n query_limit: Option,\n // ignored when query_limit is None\n latency_correction: bool,\n },\n Until {\n duration: std::time::Duration,\n n_connections: usize,\n n_http2_parallel: usize,\n query_limit: Option,\n // ignored when query_limit is None\n latency_correction: bool,\n wait_ongoing_requests_after_deadline: bool,\n },\n}\n\nimpl Opts {\n fn work_mode(&self) -> WorkMode {\n if self.debug {\n WorkMode::Debug\n } else if let Some(duration) = self.duration {\n WorkMode::Until {\n duration: duration.into(),\n n_connections: self.n_connections,\n n_http2_parallel: self.n_http2_parallel,\n query_limit: match self.query_per_second {\n Some(0f64) | None => self.burst_duration.map(|burst_duration| {\n client::QueryLimit::Burst(\n burst_duration.into(),\n self.burst_requests.unwrap_or(1),\n )\n }),\n Some(qps) => Some(client::QueryLimit::Qps(qps)),\n },\n latency_correction: self.latency_correction,\n wait_ongoing_requests_after_deadline: self.wait_ongoing_requests_after_deadline,\n }\n } else {\n let mut n_connections = self.n_connections;\n let max_useful = self.n_requests.div_ceil(self.n_http2_parallel);\n if n_connections > max_useful {\n n_connections = max_useful;\n }\n\n WorkMode::FixedNumber {\n n_requests: self.n_requests,\n n_connections,\n n_http2_parallel: self.n_http2_parallel,\n query_limit: match self.query_per_second {\n Some(0f64) | None => self.burst_duration.map(|burst_duration| {\n client::QueryLimit::Burst(\n burst_duration.into(),\n self.burst_requests.unwrap_or(1),\n )\n }),\n Some(qps) => Some(client::QueryLimit::Qps(qps)),\n },\n latency_correction: self.latency_correction,\n }\n }\n }\n}\n" }, { "path": "src/main.rs", "content": "use clap::{CommandFactory, Parser};\nuse oha::{Opts, run};\n\nfn main() {\n let num_workers_threads = std::env::var(\"TOKIO_WORKER_THREADS\")\n .ok()\n .and_then(|s| s.parse().ok())\n // Prefer to use physical cores rather than logical one because it's more performant empirically.\n .unwrap_or(num_cpus::get_physical());\n\n let rt = tokio::runtime::Builder::new_multi_thread()\n .worker_threads(num_workers_threads)\n .enable_all()\n .build()\n .unwrap();\n\n let opts = Opts::parse();\n\n if let Some(shell) = opts.completions {\n clap_complete::generate(shell, &mut Opts::command(), \"oha\", &mut std::io::stdout());\n return;\n }\n\n if let Err(e) = rt.block_on(run(opts)) {\n eprintln!(\"Error: {e}\");\n std::process::exit(libc::EXIT_FAILURE);\n }\n}\n" }, { "path": "src/monitor.rs", "content": "use byte_unit::Byte;\nuse crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers};\nuse hyper::http;\nuse ratatui::{DefaultTerminal, crossterm};\nuse ratatui::{\n layout::{Constraint, Direction, Layout},\n style::{Color, Style},\n text::{Line, Span},\n widgets::{BarChart, Block, Borders, Gauge, Paragraph},\n};\nuse std::collections::BTreeMap;\n\nuse crate::{\n client::{ClientError, RequestResult},\n printer::PrintConfig,\n result_data::{MinMaxMean, ResultData},\n timescale::{TimeLabel, TimeScale},\n};\n\n/// When the monitor ends\npub enum EndLine {\n /// After a duration\n Duration(std::time::Duration),\n /// After n query done\n NumQuery(usize),\n}\n\nstruct ColorScheme {\n light_blue: Option,\n green: Option,\n yellow: Option,\n}\n\nimpl ColorScheme {\n fn new() -> ColorScheme {\n ColorScheme {\n light_blue: None,\n green: None,\n yellow: None,\n }\n }\n\n fn set_colors(&mut self) {\n self.light_blue = Some(Color::Cyan);\n self.green = Some(Color::Green);\n self.yellow = Some(Color::Yellow);\n }\n}\n\npub struct Monitor {\n pub print_config: PrintConfig,\n pub end_line: EndLine,\n /// All workers sends each result to this channel\n pub report_receiver: kanal::Receiver>,\n // When started\n pub start: std::time::Instant,\n // Frame per second of TUI\n pub fps: usize,\n pub disable_color: bool,\n pub time_unit: Option,\n}\n\nstruct IntoRawMode;\n\nimpl IntoRawMode {\n pub fn new() -> Result<(Self, DefaultTerminal), std::io::Error> {\n let terminal = ratatui::try_init()?;\n Ok((Self, terminal))\n }\n}\n\nimpl Drop for IntoRawMode {\n fn drop(&mut self) {\n ratatui::restore();\n }\n}\n\nimpl Monitor {\n pub async fn monitor(self) -> Result<(ResultData, PrintConfig), std::io::Error> {\n let (raw_mode, mut terminal) = IntoRawMode::new()?;\n\n // Return this when ends to application print summary\n // We must not read all data from this due to computational cost.\n let mut all: ResultData = Default::default();\n // stats for HTTP status\n let mut status_dist: BTreeMap = Default::default();\n\n #[cfg(unix)]\n // Limit for number open files. eg. ulimit -n\n let nofile_limit = rlimit::getrlimit(rlimit::Resource::NOFILE);\n\n // None means auto timescale which depends on how long it takes\n let mut timescale_auto = self.time_unit;\n\n let mut colors = ColorScheme::new();\n if !self.disable_color {\n colors.set_colors();\n }\n\n let mut buf = Vec::new();\n loop {\n let frame_start = std::time::Instant::now();\n let is_disconnected = self.report_receiver.is_disconnected();\n\n let _ = self.report_receiver.drain_into(&mut buf);\n for report in buf.drain(..) {\n if let Ok(report) = report.as_ref() {\n *status_dist.entry(report.status).or_default() += 1;\n }\n all.push(report);\n }\n\n if is_disconnected {\n break;\n }\n\n let now = std::time::Instant::now();\n let progress = match &self.end_line {\n EndLine::Duration(d) => {\n ((now - self.start).as_secs_f64() / d.as_secs_f64()).clamp(0.0, 1.0)\n }\n EndLine::NumQuery(n) => (all.len() as f64 / *n as f64).clamp(0.0, 1.0),\n };\n\n let count = 32;\n\n // Make ms smallest timescale viewable for TUI\n let timescale = (if let Some(timescale) = timescale_auto {\n timescale\n } else {\n TimeScale::from_elapsed(self.start.elapsed())\n })\n .max(TimeScale::Millisecond);\n\n let bin = timescale.as_secs_f64();\n\n let mut bar_num_req = vec![0u64; count];\n let short_bin = (now - self.start).as_secs_f64() % bin;\n for r in all.success().iter().rev() {\n let past = (now - r.end).as_secs_f64();\n let i = if past <= short_bin {\n 0\n } else {\n 1 + ((past - short_bin) / bin) as usize\n };\n if i >= bar_num_req.len() {\n break;\n }\n bar_num_req[i] += 1;\n }\n\n let cols = bar_num_req\n .iter()\n .map(|x| x.to_string().chars().count())\n .max()\n .unwrap_or(0);\n\n let bar_num_req: Vec<(String, u64)> = bar_num_req\n .into_iter()\n .enumerate()\n .map(|(i, n)| {\n (\n {\n let mut s = TimeLabel { x: i, timescale }.to_string();\n if cols > s.len() {\n for _ in 0..cols - s.len() {\n s.push(' ');\n }\n }\n s\n },\n n,\n )\n })\n .collect();\n\n let bar_num_req_str: Vec<(&str, u64)> =\n bar_num_req.iter().map(|(a, b)| (a.as_str(), *b)).collect();\n\n #[cfg(unix)]\n let nofile = std::fs::read_dir(\"/dev/fd\").map(|dir| dir.count());\n\n terminal.draw(|f| {\n let row4 = Layout::default()\n .direction(Direction::Vertical)\n .constraints(\n [\n Constraint::Length(3),\n Constraint::Length(8),\n Constraint::Length(all.error_distribution().len() as u16 + 2),\n Constraint::Fill(1),\n ]\n .as_ref(),\n )\n .split(f.area());\n\n let mid = Layout::default()\n .direction(Direction::Horizontal)\n .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())\n .split(row4[1]);\n\n let bottom = Layout::default()\n .direction(Direction::Horizontal)\n .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())\n .split(row4[3]);\n\n let gauge_label = match &self.end_line {\n EndLine::Duration(d) => format!(\n \"{} / {}\",\n humantime::Duration::from(std::time::Duration::from_secs(\n (now - self.start).as_secs_f64() as u64\n )),\n humantime::Duration::from(*d)\n ),\n EndLine::NumQuery(n) => format!(\"{} / {}\", all.len(), n),\n };\n let gauge = Gauge::default()\n .block(Block::default().title(\"Progress\").borders(Borders::ALL))\n .gauge_style(Style::default().fg(colors.light_blue.unwrap_or(Color::White)))\n .label(Span::raw(gauge_label))\n .ratio(progress);\n f.render_widget(gauge, row4[0]);\n\n let last_1_timescale = {\n let success = all.success();\n let index = match success.binary_search_by(|probe| {\n (now - probe.end)\n .as_secs_f64()\n .partial_cmp(×cale.as_secs_f64())\n // Should be fine\n .unwrap()\n .reverse()\n }) {\n Ok(i) => i,\n Err(i) => i,\n };\n\n &success[index..]\n };\n\n let last_1_minmaxmean: MinMaxMean = last_1_timescale\n .iter()\n .map(|r| r.duration().as_secs_f64())\n .collect();\n\n let stats_text = vec![\n Line::from(format!(\"Requests : {}\", last_1_timescale.len())),\n Line::from(vec![Span::styled(\n format!(\"Slowest: {:.4} secs\", last_1_minmaxmean.max(),),\n Style::default().fg(colors.yellow.unwrap_or(Color::Reset)),\n )]),\n Line::from(vec![Span::styled(\n format!(\"Fastest: {:.4} secs\", last_1_minmaxmean.min(),),\n Style::default().fg(colors.green.unwrap_or(Color::Reset)),\n )]),\n Line::from(vec![Span::styled(\n format!(\"Average: {:.4} secs\", last_1_minmaxmean.mean(),),\n Style::default().fg(colors.light_blue.unwrap_or(Color::Reset)),\n )]),\n Line::from(format!(\n \"Data: {:.2}\",\n Byte::from_u64(\n last_1_timescale\n .iter()\n .map(|r| r.len_bytes as u64)\n .sum::()\n )\n .get_appropriate_unit(byte_unit::UnitType::Binary)\n )),\n #[cfg(unix)]\n // Note: Windows can open 255 * 255 * 255 files. So not showing on windows is OK.\n Line::from(format!(\n \"Number of open files: {} / {}\",\n nofile\n .map(|c| c.to_string())\n .unwrap_or_else(|_| \"Error\".to_string()),\n nofile_limit\n .as_ref()\n .map(|(s, _h)| s.to_string())\n .unwrap_or_else(|_| \"Unknown\".to_string())\n )),\n ];\n let stats_title = format!(\"Stats for last {timescale}\");\n let stats = Paragraph::new(stats_text).block(\n Block::default()\n .title(Span::raw(stats_title))\n .borders(Borders::ALL),\n );\n f.render_widget(stats, mid[0]);\n\n let mut status_v: Vec<(http::StatusCode, usize)> =\n status_dist.clone().into_iter().collect();\n status_v.sort_by_key(|t| std::cmp::Reverse(t.1));\n\n let stats2_text = status_v\n .into_iter()\n .map(|(status, count)| {\n Line::from(format!(\"[{}] {} responses\", status.as_str(), count))\n })\n .collect::>();\n let stats2 = Paragraph::new(stats2_text).block(\n Block::default()\n .title(\"Status code distribution\")\n .borders(Borders::ALL),\n );\n f.render_widget(stats2, mid[1]);\n\n let mut error_v: Vec<(String, usize)> =\n all.error_distribution().clone().into_iter().collect();\n error_v.sort_by_key(|t| std::cmp::Reverse(t.1));\n let errors_text = error_v\n .into_iter()\n .map(|(e, count)| Line::from(format!(\"[{count}] {e}\")))\n .collect::>();\n let errors = Paragraph::new(errors_text).block(\n Block::default()\n .title(\"Error distribution\")\n .borders(Borders::ALL),\n );\n f.render_widget(errors, row4[2]);\n\n let title = format!(\n \"Requests / past {}{}. press -/+/a to change\",\n timescale,\n if timescale_auto.is_none() {\n \" (auto)\"\n } else {\n \"\"\n }\n );\n\n let barchart = BarChart::default()\n .block(\n Block::default()\n .title(Span::raw(title))\n .style(\n Style::default()\n .fg(colors.green.unwrap_or(Color::Reset))\n .bg(Color::Reset),\n )\n .borders(Borders::ALL),\n )\n .data(bar_num_req_str.as_slice())\n .bar_width(\n bar_num_req\n .iter()\n .map(|(s, _)| s.chars().count())\n .max()\n .map(|w| w + 2)\n .unwrap_or(1) as u16,\n );\n f.render_widget(barchart, bottom[0]);\n\n let resp_histo_width = 7;\n let resp_histo_data: Vec<(String, u64)> = {\n let bins = if bottom[1].width < 2 {\n 0\n } else {\n (bottom[1].width as usize - 2) / (resp_histo_width + 1)\n }\n .max(2);\n let values = last_1_timescale\n .iter()\n .map(|r| r.duration().as_secs_f64())\n .collect::>();\n\n let histo = crate::histogram::histogram(&values, bins);\n histo\n .into_iter()\n .map(|(label, v)| (format!(\"{label:.4}\"), v as u64))\n .collect()\n };\n\n let resp_histo_data_str: Vec<(&str, u64)> = resp_histo_data\n .iter()\n .map(|(l, v)| (l.as_str(), *v))\n .collect();\n\n let resp_histo = BarChart::default()\n .block(\n Block::default()\n .title(\"Response time histogram\")\n .style(\n Style::default()\n .fg(colors.yellow.unwrap_or(Color::Reset))\n .bg(Color::Reset),\n )\n .borders(Borders::ALL),\n )\n .data(resp_histo_data_str.as_slice())\n .bar_width(resp_histo_width as u16);\n f.render_widget(resp_histo, bottom[1]);\n })?;\n\n while crossterm::event::poll(std::time::Duration::from_secs(0))? {\n match crossterm::event::read()? {\n Event::Key(KeyEvent {\n code: KeyCode::Char('+'),\n ..\n }) => {\n // Make ms the smallest timescale viewable in TUI\n timescale_auto = Some(timescale.dec().max(TimeScale::Millisecond))\n }\n Event::Key(KeyEvent {\n code: KeyCode::Char('-'),\n ..\n }) => timescale_auto = Some(timescale.inc()),\n Event::Key(KeyEvent {\n code: KeyCode::Char('a'),\n ..\n }) => {\n if timescale_auto.is_some() {\n timescale_auto = None;\n } else {\n timescale_auto = Some(timescale)\n }\n }\n // User pressed q or ctrl-c\n Event::Key(KeyEvent {\n code: KeyCode::Char('q'),\n ..\n })\n | Event::Key(KeyEvent {\n code: KeyCode::Char('c'),\n modifiers: KeyModifiers::CONTROL,\n ..\n }) => {\n drop(terminal);\n drop(raw_mode);\n let _ = crate::printer::print_result(\n self.print_config,\n self.start,\n &all,\n now - self.start,\n );\n std::process::exit(libc::EXIT_SUCCESS);\n }\n _ => (),\n }\n }\n\n let per_frame = std::time::Duration::from_secs(1) / self.fps as u32;\n let elapsed = frame_start.elapsed();\n if per_frame > elapsed {\n tokio::time::sleep(per_frame - elapsed).await;\n }\n }\n Ok((all, self.print_config))\n }\n}\n" }, { "path": "src/pcg64si.rs", "content": "use std::convert::Infallible;\n\n// https://github.com/imneme/pcg-c\nuse rand::{SeedableRng, TryRng, rand_core::utils::fill_bytes_via_next_word};\n\n#[derive(Debug, Copy, Clone)]\n#[repr(transparent)]\npub struct Pcg64Si {\n state: u64,\n}\n\nimpl TryRng for Pcg64Si {\n type Error = Infallible;\n\n fn try_next_u32(&mut self) -> Result {\n Ok(self.try_next_u64()? as u32)\n }\n\n fn try_next_u64(&mut self) -> Result {\n let old_state = self.state;\n self.state = self\n .state\n .wrapping_mul(6364136223846793005)\n .wrapping_add(1442695040888963407);\n\n let word =\n ((old_state >> ((old_state >> 59) + 5)) ^ old_state).wrapping_mul(12605985483714917081);\n Ok((word >> 43) ^ word)\n }\n\n fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {\n fill_bytes_via_next_word(dest, || self.try_next_u64())\n }\n}\n\nimpl SeedableRng for Pcg64Si {\n type Seed = [u8; 8];\n\n fn from_seed(seed: Self::Seed) -> Pcg64Si {\n Pcg64Si {\n state: u64::from_le_bytes(seed),\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use rand::Rng;\n use std::collections::HashSet;\n\n // For a given seed the RNG is deterministic\n // thus we can perform some basic tests consistently\n #[test]\n fn test_rng_next() {\n let mut rng = Pcg64Si::from_seed([1, 2, 3, 4, 5, 6, 7, 8]);\n let mut values_set: HashSet = HashSet::new();\n // Generate 1000 values modulus 100 (so each value is between 0 and 99)\n for _ in 0..1000 {\n values_set.insert(rng.next_u32() % 100);\n }\n // Expect to generate every number between 0 and 99 (the generated values are somewhat evenly distributed)\n assert_eq!(values_set.len(), 100);\n }\n\n #[test]\n fn test_rng_from_seed() {\n // Different seeds should result in a different RNG state\n let rng1 = Pcg64Si::from_seed([1, 2, 3, 4, 5, 6, 7, 8]);\n let rng2 = Pcg64Si::from_seed([1, 2, 3, 4, 5, 6, 7, 7]);\n assert_ne!(rng1.state, rng2.state);\n }\n\n #[test]\n fn test_rng_fill_bytes() {\n // This uses the next_u64/u32 functions underneath, so don't need to test the pseudo randomness again\n let mut array: [u8; 8] = [0, 0, 0, 0, 0, 0, 0, 0];\n let mut rng = Pcg64Si::from_seed([1, 2, 3, 4, 5, 6, 7, 8]);\n rng.fill_bytes(&mut array);\n assert_ne!(array, [0, 0, 0, 0, 0, 0, 0, 0]);\n }\n}\n" }, { "path": "src/printer.rs", "content": "use crate::{result_data::ResultData, timescale::TimeScale};\nuse average::{Max, Min, Variance};\nuse byte_unit::Byte;\nuse crossterm::style::{StyledContent, Stylize};\nuse hyper::http::{self, StatusCode};\nuse ratatui::crossterm;\nuse std::{\n collections::BTreeMap,\n io::Write,\n time::{Duration, Instant},\n};\n\n#[derive(Clone, Copy)]\nstruct StyleScheme {\n style_enabled: bool,\n}\nimpl StyleScheme {\n fn no_style(self, text: &str) -> StyledContent<&str> {\n StyledContent::new(crossterm::style::ContentStyle::new(), text)\n }\n fn heading(self, text: &str) -> StyledContent<&str> {\n if self.style_enabled {\n text.bold().underlined()\n } else {\n self.no_style(text)\n }\n }\n fn success_rate(self, text: &str, success_rate: f64) -> StyledContent<&str> {\n if self.style_enabled {\n if success_rate >= 100.0 {\n text.green().bold()\n } else if success_rate >= 99.0 {\n text.yellow().bold()\n } else {\n text.red().bold()\n }\n } else {\n self.no_style(text)\n }\n }\n fn fastest(self, text: &str) -> StyledContent<&str> {\n if self.style_enabled {\n text.green()\n } else {\n self.no_style(text)\n }\n }\n fn slowest(self, text: &str) -> StyledContent<&str> {\n if self.style_enabled {\n text.yellow()\n } else {\n self.no_style(text)\n }\n }\n fn average(self, text: &str) -> StyledContent<&str> {\n if self.style_enabled {\n text.cyan()\n } else {\n self.no_style(text)\n }\n }\n\n fn latency_distribution(self, text: &str, label: f64) -> StyledContent<&str> {\n // See #609 for justification of these thresholds\n const LATENCY_YELLOW_THRESHOLD: f64 = 0.1;\n const LATENCY_RED_THRESHOLD: f64 = 0.4;\n\n if self.style_enabled {\n if label <= LATENCY_YELLOW_THRESHOLD {\n text.green()\n } else if label <= LATENCY_RED_THRESHOLD {\n text.yellow()\n } else {\n text.red()\n }\n } else {\n self.no_style(text)\n }\n }\n\n fn status_distribution(self, text: &str, status: StatusCode) -> StyledContent<&str> {\n if self.style_enabled {\n if status.is_success() {\n text.green()\n } else if status.is_client_error() {\n text.yellow()\n } else if status.is_server_error() {\n text.red()\n } else {\n text.white()\n }\n } else {\n self.no_style(text)\n }\n }\n}\n\n#[derive(Clone, Copy, Debug, Default, clap::ValueEnum)]\npub enum PrintMode {\n #[default]\n Text,\n Json,\n Csv,\n Quiet,\n}\n\npub struct PrintConfig {\n pub output: Box,\n pub mode: PrintMode,\n pub disable_style: bool,\n pub stats_success_breakdown: bool,\n pub time_unit: Option,\n}\n\npub fn print_result(\n mut config: PrintConfig,\n start: Instant,\n res: &ResultData,\n total_duration: Duration,\n) -> anyhow::Result<()> {\n match config.mode {\n PrintMode::Text => print_summary(\n &mut config.output,\n res,\n total_duration,\n config.disable_style,\n config.stats_success_breakdown,\n config.time_unit,\n )?,\n PrintMode::Json => print_json(\n &mut config.output,\n start,\n res,\n total_duration,\n config.stats_success_breakdown,\n )?,\n PrintMode::Csv => print_csv(&mut config.output, start, res)?,\n PrintMode::Quiet => { /* Do nothing */ }\n }\n Ok(())\n}\n\n/// Print all summary as JSON\nfn print_json(\n w: &mut W,\n start: Instant,\n res: &ResultData,\n total_duration: Duration,\n stats_success_breakdown: bool,\n) -> serde_json::Result<()> {\n use serde::Serialize;\n #[derive(Serialize)]\n struct Summary {\n #[serde(rename = \"successRate\")]\n success_rate: f64,\n total: f64,\n slowest: f64,\n fastest: f64,\n average: f64,\n #[serde(rename = \"requestsPerSec\")]\n requests_per_sec: f64,\n #[serde(rename = \"totalData\")]\n total_data: u64,\n #[serde(rename = \"sizePerRequest\")]\n size_per_request: Option,\n #[serde(rename = \"sizePerSec\")]\n size_per_sec: f64,\n }\n\n #[derive(Serialize)]\n struct Triple {\n average: f64,\n fastest: f64,\n slowest: f64,\n }\n\n #[derive(Serialize)]\n struct Details {\n #[serde(rename = \"DNSDialup\")]\n dns_dialup: Triple,\n #[serde(rename = \"DNSLookup\")]\n dns_lookup: Triple,\n #[serde(rename = \"firstByte\")]\n first_byte: Triple,\n }\n\n #[derive(Serialize)]\n struct Rps {\n mean: f64,\n stddev: f64,\n max: f64,\n min: f64,\n percentiles: BTreeMap,\n }\n\n #[derive(Serialize)]\n struct Result {\n summary: Summary,\n #[serde(rename = \"responseTimeHistogram\")]\n response_time_histogram: BTreeMap,\n #[serde(rename = \"latencyPercentiles\")]\n latency_percentiles: BTreeMap,\n #[serde(rename = \"firstByteHistogram\")]\n first_byte_histogram: BTreeMap,\n #[serde(rename = \"firstBytePercentiles\")]\n first_byte_percentiles: BTreeMap,\n #[serde(\n rename = \"responseTimeHistogramSuccessful\",\n skip_serializing_if = \"Option::is_none\"\n )]\n response_time_histogram_successful: Option>,\n #[serde(\n rename = \"latencyPercentilesSuccessful\",\n skip_serializing_if = \"Option::is_none\"\n )]\n latency_percentiles_successful: Option>,\n #[serde(\n rename = \"responseTimeHistogramNotSuccessful\",\n skip_serializing_if = \"Option::is_none\"\n )]\n response_time_histogram_not_successful: Option>,\n #[serde(\n rename = \"latencyPercentilesNotSuccessful\",\n skip_serializing_if = \"Option::is_none\"\n )]\n latency_percentiles_not_successful: Option>,\n #[serde(rename = \"rps\")]\n rps: Rps,\n details: Details,\n #[serde(rename = \"statusCodeDistribution\")]\n status_code_distribution: BTreeMap,\n #[serde(rename = \"errorDistribution\")]\n error_distribution: BTreeMap,\n }\n\n let latency_stat = res.latency_stat();\n\n let summary = Summary {\n success_rate: res.success_rate(),\n total: total_duration.as_secs_f64(),\n slowest: latency_stat.max(),\n fastest: latency_stat.min(),\n average: latency_stat.mean(),\n requests_per_sec: res.len() as f64 / total_duration.as_secs_f64(),\n total_data: res.total_data() as u64,\n size_per_request: res.size_per_request(),\n size_per_sec: res.total_data() as f64 / total_duration.as_secs_f64(),\n };\n\n let durations_statistics = res.duration_all_statistics();\n\n let response_time_histogram = durations_statistics\n .histogram\n .into_iter()\n .map(|(k, v)| (k.to_string(), v))\n .collect();\n\n let latency_percentiles = durations_statistics\n .percentiles\n .into_iter()\n .map(|(p, v)| (format!(\"p{p}\"), v))\n .collect();\n\n let first_byte_statistics = res.first_byte_all_statistics();\n\n let first_byte_histogram = first_byte_statistics\n .histogram\n .into_iter()\n .map(|(k, v)| (k.to_string(), v))\n .collect();\n\n let first_byte_percentiles = first_byte_statistics\n .percentiles\n .into_iter()\n .map(|(p, v)| (format!(\"p{p}\"), v))\n .collect();\n\n let mut response_time_histogram_successful: Option> = None;\n let mut latency_percentiles_successful: Option> = None;\n let mut response_time_histogram_not_successful: Option> = None;\n let mut latency_percentiles_not_successful: Option> = None;\n\n if stats_success_breakdown {\n let durations_successful_statistics = res.duration_successful_statistics();\n\n response_time_histogram_successful = Some(\n durations_successful_statistics\n .histogram\n .into_iter()\n .map(|(k, v)| (k.to_string(), v))\n .collect(),\n );\n\n latency_percentiles_successful = Some(\n durations_successful_statistics\n .percentiles\n .into_iter()\n .map(|(p, v)| (format!(\"p{p}\"), v))\n .collect(),\n );\n\n let durations_not_successful_statistics = res.duration_not_successful_statistics();\n\n response_time_histogram_not_successful = Some(\n durations_not_successful_statistics\n .histogram\n .into_iter()\n .map(|(k, v)| (k.to_string(), v))\n .collect(),\n );\n\n latency_percentiles_not_successful = Some(\n durations_not_successful_statistics\n .percentiles\n .into_iter()\n .map(|(p, v)| (format!(\"p{p}\"), v))\n .collect(),\n );\n }\n\n let mut ends = res\n .end_times_from_start(start)\n .map(|d| d.as_secs_f64())\n .collect::>();\n ends.push(0.0);\n float_ord::sort(&mut ends);\n\n let mut rps: Vec = Vec::new();\n // 10ms\n const INTERVAL: f64 = 0.01;\n let mut r = 0;\n loop {\n let prev_r = r;\n\n // increment at least 1\n if r + 1 < ends.len() {\n r += 1;\n }\n\n while r + 1 < ends.len() && ends[prev_r] + INTERVAL > ends[r + 1] {\n r += 1;\n }\n\n if r == prev_r {\n break;\n }\n\n let n = r - prev_r;\n let t = ends[r] - ends[prev_r];\n rps.push(n as f64 / t);\n }\n\n let rps_percentiles = percentiles(&mut rps);\n\n let variance = rps.iter().collect::();\n let rps = Rps {\n mean: variance.mean(),\n stddev: variance.sample_variance().sqrt(),\n max: rps.iter().collect::().max(),\n min: rps.iter().collect::().min(),\n percentiles: rps_percentiles,\n };\n\n let status_code_distribution = res.status_code_distribution();\n\n let dns_dialup_stat = res.dns_dialup_stat();\n let dns_lookup_stat = res.dns_lookup_stat();\n let first_byte_stat = res.first_byte_stat();\n\n let details = Details {\n dns_dialup: Triple {\n average: dns_dialup_stat.mean(),\n fastest: dns_dialup_stat.min(),\n slowest: dns_dialup_stat.max(),\n },\n dns_lookup: Triple {\n average: dns_lookup_stat.mean(),\n fastest: dns_lookup_stat.min(),\n slowest: dns_lookup_stat.max(),\n },\n first_byte: Triple {\n average: first_byte_stat.mean(),\n fastest: first_byte_stat.min(),\n slowest: first_byte_stat.max(),\n },\n };\n\n serde_json::to_writer_pretty(\n w,\n &Result {\n summary,\n response_time_histogram,\n latency_percentiles,\n first_byte_histogram,\n first_byte_percentiles,\n response_time_histogram_successful,\n latency_percentiles_successful,\n response_time_histogram_not_successful,\n latency_percentiles_not_successful,\n rps,\n details,\n status_code_distribution: status_code_distribution\n .into_iter()\n .map(|(k, v)| (k.as_u16().to_string(), v))\n .collect(),\n error_distribution: res.error_distribution().clone(),\n },\n )\n}\n\nfn print_csv(w: &mut W, start: Instant, res: &ResultData) -> std::io::Result<()> {\n // csv header\n writeln!(\n w,\n \"request-start,DNS,DNS+dialup,Response-delay,request-duration,bytes,status\"\n )?;\n\n let mut success_requests = res.success().to_vec();\n success_requests.sort_by_key(|r| r.start);\n\n for request in success_requests {\n let dns_and_dialup = match request.connection_time {\n Some(connection_time) => (connection_time.dns_lookup, connection_time.dialup),\n None => (std::time::Duration::ZERO, std::time::Duration::ZERO),\n };\n let first_byte = match request.first_byte {\n Some(first_byte) => first_byte - request.start,\n None => std::time::Duration::ZERO,\n };\n writeln!(\n w,\n \"{},{},{},{},{},{},{}\",\n (request.start - start).as_secs_f64(),\n dns_and_dialup.0.as_secs_f64(),\n dns_and_dialup.1.as_secs_f64(),\n first_byte.as_secs_f64(),\n request.duration().as_secs_f64(),\n request.len_bytes,\n request.status.as_u16(),\n )?;\n }\n Ok(())\n}\n\n/// Print all summary as Text\nfn print_summary(\n w: &mut W,\n res: &ResultData,\n total_duration: Duration,\n disable_style: bool,\n stats_success_breakdown: bool,\n time_unit: Option,\n) -> std::io::Result<()> {\n let style = StyleScheme {\n style_enabled: !disable_style,\n };\n writeln!(w, \"{}\", style.heading(\"Summary:\"))?;\n let success_rate = 100.0 * res.success_rate();\n writeln!(\n w,\n \"{}\",\n style.success_rate(\n &format!(\" Success rate:\\t{success_rate:.2}%\"),\n success_rate\n )\n )?;\n let latency_stat = res.latency_stat();\n // Determine timescale automatically\n let timescale = if let Some(timescale) = time_unit {\n timescale\n } else {\n // Use max latency (slowest request)\n TimeScale::from_f64(latency_stat.max())\n };\n writeln!(\n w,\n \" Total:\\t{:.4} {timescale}\",\n total_duration.as_secs_f64() / timescale.as_secs_f64()\n )?;\n writeln!(\n w,\n \"{}\",\n style.slowest(&format!(\n \" Slowest:\\t{:.4} {timescale}\",\n latency_stat.max() / timescale.as_secs_f64()\n ))\n )?;\n writeln!(\n w,\n \"{}\",\n style.fastest(&format!(\n \" Fastest:\\t{:.4} {timescale}\",\n latency_stat.min() / timescale.as_secs_f64()\n ))\n )?;\n writeln!(\n w,\n \"{}\",\n style.average(&format!(\n \" Average:\\t{:.4} {timescale}\",\n latency_stat.mean() / timescale.as_secs_f64()\n ))\n )?;\n writeln!(\n w,\n \" Requests/sec:\\t{:.4}\",\n res.len() as f64 / total_duration.as_secs_f64()\n )?;\n writeln!(w)?;\n writeln!(\n w,\n \" Total data:\\t{:.2}\",\n Byte::from_u64(res.total_data() as u64).get_appropriate_unit(byte_unit::UnitType::Binary)\n )?;\n if let Some(size) = res\n .size_per_request()\n .map(|n| Byte::from_u64(n).get_appropriate_unit(byte_unit::UnitType::Binary))\n {\n writeln!(w, \" Size/request:\\t{size:.2}\")?;\n } else {\n writeln!(w, \" Size/request:\\tNaN\")?;\n }\n writeln!(\n w,\n \" Size/sec:\\t{:.2}\",\n Byte::from_u64((res.total_data() as f64 / total_duration.as_secs_f64()) as u64)\n .get_appropriate_unit(byte_unit::UnitType::Binary)\n )?;\n writeln!(w)?;\n\n let duration_all_statistics = res.duration_all_statistics();\n\n writeln!(w, \"{}\", style.heading(\"Response time histogram:\"))?;\n print_histogram(w, &duration_all_statistics.histogram, style, timescale)?;\n writeln!(w)?;\n\n writeln!(w, \"{}\", style.heading(\"Response time distribution:\"))?;\n print_distribution(w, &duration_all_statistics.percentiles, style, timescale)?;\n writeln!(w)?;\n\n if stats_success_breakdown {\n let durations_successful_statics = res.duration_successful_statistics();\n\n writeln!(\n w,\n \"{}\",\n style.heading(\"Response time histogram (2xx only):\")\n )?;\n print_histogram(w, &durations_successful_statics.histogram, style, timescale)?;\n writeln!(w)?;\n\n writeln!(\n w,\n \"{}\",\n style.heading(\"Response time distribution (2xx only):\")\n )?;\n print_distribution(\n w,\n &durations_successful_statics.percentiles,\n style,\n timescale,\n )?;\n writeln!(w)?;\n\n let durations_not_successful = res.duration_not_successful_statistics();\n\n writeln!(\n w,\n \"{}\",\n style.heading(\"Response time histogram (4xx + 5xx only):\")\n )?;\n print_histogram(w, &durations_not_successful.histogram, style, timescale)?;\n writeln!(w)?;\n\n writeln!(\n w,\n \"{}\",\n style.heading(\"Response time distribution (4xx + 5xx only):\")\n )?;\n print_distribution(w, &durations_not_successful.percentiles, style, timescale)?;\n writeln!(w)?;\n }\n writeln!(w)?;\n\n let dns_dialup_stat = res.dns_dialup_stat();\n let dns_lookup_stat = res.dns_lookup_stat();\n\n writeln!(\n w,\n \"{}\",\n style.heading(\"Details (average, fastest, slowest):\")\n )?;\n\n writeln!(\n w,\n \" DNS+dialup:\\t{:.4} {timescale}, {:.4} {timescale}, {:.4} {timescale}\",\n dns_dialup_stat.mean() / timescale.as_secs_f64(),\n dns_dialup_stat.min() / timescale.as_secs_f64(),\n dns_dialup_stat.max() / timescale.as_secs_f64()\n )?;\n writeln!(\n w,\n \" DNS-lookup:\\t{:.4} {timescale}, {:.4} {timescale}, {:.4} {timescale}\",\n dns_lookup_stat.mean() / timescale.as_secs_f64(),\n dns_lookup_stat.min() / timescale.as_secs_f64(),\n dns_lookup_stat.max() / timescale.as_secs_f64()\n )?;\n writeln!(w)?;\n\n let status_dist: BTreeMap = res.status_code_distribution();\n\n let mut status_v: Vec<(http::StatusCode, usize)> = status_dist.into_iter().collect();\n status_v.sort_by_key(|t| std::cmp::Reverse(t.1));\n\n writeln!(w, \"{}\", style.heading(\"Status code distribution:\"))?;\n\n for (status, count) in status_v {\n writeln!(\n w,\n \"{}\",\n style.status_distribution(\n &format!(\" [{}] {} responses\", status.as_str(), count),\n status\n )\n )?;\n }\n\n let mut error_v: Vec<(String, usize)> = res\n .error_distribution()\n .iter()\n .map(|(k, v)| (k.clone(), *v))\n .collect();\n error_v.sort_by_key(|t| std::cmp::Reverse(t.1));\n\n if !error_v.is_empty() {\n writeln!(w)?;\n writeln!(w, \"Error distribution:\")?;\n for (error, count) in error_v {\n writeln!(w, \" [{count}] {error}\")?;\n }\n }\n\n Ok(())\n}\n\n/// This is used to print histogram of response time.\nfn print_histogram(\n w: &mut W,\n data: &[(f64, usize)],\n style: StyleScheme,\n timescale: TimeScale,\n) -> std::io::Result<()> {\n let max_bar = data.iter().map(|t| t.1).max().unwrap();\n let str_len_max = max_bar.to_string().len();\n let width = data\n .iter()\n .map(|t| ((t.0 / timescale.as_secs_f64()) as u64).to_string().len())\n .max()\n .unwrap()\n + 4;\n\n for (label, b) in data.iter() {\n let indent = str_len_max - b.to_string().len();\n write!(\n w,\n \"{}\",\n style.latency_distribution(\n &format!(\n \" {:>width$.3} {timescale} [{}]{} |\",\n label / timescale.as_secs_f64(),\n b,\n \" \".repeat(indent),\n width = width\n ),\n *label\n )\n )?;\n bar(w, *b as f64 / max_bar as f64, style, *label)?;\n writeln!(w)?;\n }\n Ok(())\n}\n\n// Print Bar like ■■■■■■■■■\nfn bar(w: &mut W, ratio: f64, style: StyleScheme, label: f64) -> std::io::Result<()> {\n // TODO: Use more block element code to show more precise bar\n let width = 32;\n for _ in 0..(width as f64 * ratio) as usize {\n write!(w, \"{}\", style.latency_distribution(\"■\", label))?;\n }\n Ok(())\n}\n\nfn percentile_iter(values: &mut [f64]) -> impl Iterator + '_ {\n float_ord::sort(values);\n\n [10.0, 25.0, 50.0, 75.0, 90.0, 95.0, 99.0, 99.9, 99.99]\n .iter()\n .map(move |&p| {\n let i = (p / 100.0 * values.len() as f64) as usize;\n (p, *values.get(i).unwrap_or(&f64::NAN))\n })\n}\n\n/// Print distribution of collection of f64\nfn print_distribution(\n w: &mut W,\n percentiles: &[(f64, f64)],\n style: StyleScheme,\n timescale: TimeScale,\n) -> std::io::Result<()> {\n for (p, v) in percentiles {\n writeln!(\n w,\n \"{}\",\n style.latency_distribution(\n &format!(\n \" {p:.2}% in {:.4} {timescale}\",\n v / timescale.as_secs_f64()\n ),\n *v\n )\n )?;\n }\n\n Ok(())\n}\n\nfn percentiles(values: &mut [f64]) -> BTreeMap {\n percentile_iter(values)\n .map(|(p, v)| (format!(\"p{p}\"), v))\n .collect()\n}\n\n#[cfg(test)]\nmod tests {\n use float_cmp::assert_approx_eq;\n\n use super::*;\n\n #[test]\n fn test_percentile_iter() {\n let mut values: [f64; 40] = [\n 5.0, 5.0, 5.0, 5.0, 5.0, 10.0, 10.0, 10.0, 10.0, 10.0, 11.0, 11.0, 11.0, 11.0, 11.0,\n 11.0, 11.0, 11.0, 11.0, 11.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0,\n 12.0, 15.0, 15.0, 15.0, 15.0, 15.0, 20.0, 20.0, 20.0, 25.0, 30.0,\n ];\n let result: Vec<(f64, f64)> = percentile_iter(&mut values).collect();\n assert_approx_eq!(&[f64], &[result[0].0, result[0].1], &[10.0, 5_f64]);\n assert_approx_eq!(&[f64], &[result[1].0, result[1].1], &[25.0, 11_f64]);\n assert_approx_eq!(&[f64], &[result[2].0, result[2].1], &[50.0, 12_f64]);\n assert_approx_eq!(&[f64], &[result[3].0, result[3].1], &[75.0, 15_f64]);\n assert_approx_eq!(&[f64], &[result[4].0, result[4].1], &[90.0, 20_f64]);\n assert_approx_eq!(&[f64], &[result[5].0, result[5].1], &[95.0, 25_f64]);\n assert_approx_eq!(&[f64], &[result[6].0, result[6].1], &[99.0, 30_f64]);\n assert_approx_eq!(&[f64], &[result[7].0, result[7].1], &[99.9, 30_f64]);\n assert_approx_eq!(&[f64], &[result[8].0, result[8].1], &[99.99, 30_f64]);\n }\n}\n" }, { "path": "src/request_generator.rs", "content": "use std::borrow::Cow;\n\nuse bytes::Bytes;\nuse http_body_util::Full;\nuse hyper::http;\nuse hyper::{HeaderMap, Method, Version};\nuse rand::Rng;\nuse rand::seq::IndexedRandom;\nuse thiserror::Error;\nuse url::Url;\n\nuse crate::aws_auth::{self, AwsSignatureConfig};\nuse crate::url_generator;\n\npub struct Proxy {\n pub headers: HeaderMap,\n pub version: Version,\n}\n\npub enum BodyGenerator {\n Static(Bytes),\n Random(Vec),\n}\n\npub struct RequestGenerator {\n pub url_generator: url_generator::UrlGenerator,\n pub https: bool,\n // Only if http with proxy\n pub http_proxy: Option,\n pub method: Method,\n pub version: Version,\n pub headers: HeaderMap,\n pub body_generator: BodyGenerator,\n pub aws_config: Option,\n}\n\n#[derive(Error, Debug)]\npub enum RequestGenerationError {\n #[error(\"URL generation error: {0}\")]\n UrlGeneration(#[from] url_generator::UrlGeneratorError),\n #[error(\"Request building error: {0}\")]\n RequestBuild(#[from] http::Error),\n #[error(\"AWS Signature error: {0}\")]\n AwsSignature(#[from] aws_auth::AwsSignatureError),\n}\n\nimpl RequestGenerator {\n #[inline]\n fn is_http1(&self) -> bool {\n self.version <= Version::HTTP_11\n }\n\n fn generate_body(&self, rng: &mut R) -> Bytes {\n match &self.body_generator {\n BodyGenerator::Static(b) => b.clone(),\n BodyGenerator::Random(choices) => choices.choose(rng).cloned().unwrap_or_default(),\n }\n }\n\n pub fn generate(\n &self,\n rng: &mut R,\n ) -> Result<(Cow<'_, Url>, hyper::Request>), RequestGenerationError> {\n let url = self.url_generator.generate(rng)?;\n let body = self.generate_body(rng);\n\n let mut builder = hyper::Request::builder()\n .uri(if !self.is_http1() || self.http_proxy.is_some() {\n &url[..]\n } else {\n &url[url::Position::BeforePath..]\n })\n .method(self.method.clone())\n .version(\n self.http_proxy\n .as_ref()\n .map(|p| p.version)\n .unwrap_or(self.version),\n );\n\n let mut headers = self.headers.clone();\n\n // Apply AWS SigV4 if configured\n if let Some(aws_config) = &self.aws_config {\n aws_config.sign_request(self.method.as_str(), &mut headers, &url, &body)?;\n }\n\n if let Some(proxy) = &self.http_proxy {\n for (key, value) in proxy.headers.iter() {\n headers.insert(key, value.clone());\n }\n }\n\n if self.version < Version::HTTP_2 {\n headers\n .entry(http::header::HOST)\n .or_insert_with(|| http::header::HeaderValue::from_str(url.authority()).unwrap());\n }\n\n *builder.headers_mut().unwrap() = headers;\n\n let req = builder.body(Full::new(body))?;\n Ok((url, req))\n }\n}\n" }, { "path": "src/result_data.rs", "content": "use std::{\n collections::BTreeMap,\n time::{Duration, Instant},\n};\n\nuse average::{Estimate, Max, Mean, Min, concatenate};\nuse hyper::StatusCode;\n\nuse crate::{\n client::{ClientError, RequestResult},\n histogram::histogram,\n};\n\n/// Data container for the results of the all requests\n/// When a request is successful, the result is pushed to the `success` vector and the memory consumption will not be a problem because the number of successful requests is limited by network overhead.\n/// When a request fails, the error message is pushed to the `error` map because the number of error messages may huge.\n#[derive(Debug, Default)]\npub struct ResultData {\n success: Vec,\n error_distribution: BTreeMap,\n}\n\nconcatenate!(pub MinMaxMean, [Min, min], [Max, max], [Mean, mean]);\n\npub struct Statistics {\n pub percentiles: Vec<(f64, f64)>,\n pub histogram: Vec<(f64, usize)>,\n}\n\nimpl Statistics {\n /* private */\n fn new(data: &mut [f64]) -> Self {\n float_ord::sort(data);\n\n Self {\n percentiles: percentile_iter(data).collect(),\n histogram: histogram(data, 11),\n }\n }\n}\n\nfn percentile_iter(values: &mut [f64]) -> impl Iterator + '_ {\n float_ord::sort(values);\n\n [10.0, 25.0, 50.0, 75.0, 90.0, 95.0, 99.0, 99.9, 99.99]\n .iter()\n .map(move |&p| {\n let i = (p / 100.0 * values.len() as f64) as usize;\n (p, *values.get(i).unwrap_or(&f64::NAN))\n })\n}\n\nimpl ResultData {\n #[inline]\n pub fn push(&mut self, result: Result) {\n match result {\n Ok(result) => self.success.push(result),\n Err(err) => {\n let count = self.error_distribution.entry(err.to_string()).or_insert(0);\n *count += 1;\n }\n }\n }\n\n pub fn len(&self) -> usize {\n self.success.len() + self.error_distribution.values().sum::()\n }\n\n pub fn merge(&mut self, other: ResultData) {\n self.success.extend(other.success);\n for (k, v) in other.error_distribution {\n let count = self.error_distribution.entry(k).or_insert(0);\n *count += v;\n }\n }\n\n // An existence of this method doesn't prevent us to using hdrhistogram.\n // Because this is only called from `monitor` and `monitor` can collect own data.\n pub fn success(&self) -> &[RequestResult] {\n &self.success\n }\n\n // It's very happy if you can provide all below methods without array (= non liner memory consumption) and fast `push` runtime.\n\n pub fn success_rate(&self) -> f64 {\n let dead_line = ClientError::Deadline.to_string();\n // We ignore deadline errors which are because of `-z` option, not because of the server\n let denominator = self.success.len()\n + self\n .error_distribution\n .iter()\n .filter_map(|(k, v)| if k == &dead_line { None } else { Some(v) })\n .sum::();\n let numerator = self.success.len();\n\n numerator as f64 / denominator as f64\n }\n\n pub fn latency_stat(&self) -> MinMaxMean {\n self.success\n .iter()\n .map(|result| result.duration().as_secs_f64())\n .collect()\n }\n\n pub fn error_distribution(&self) -> &BTreeMap {\n &self.error_distribution\n }\n\n pub fn end_times_from_start(&self, start: Instant) -> impl Iterator + '_ {\n self.success.iter().map(move |result| result.end - start)\n }\n\n pub fn status_code_distribution(&self) -> BTreeMap {\n let mut dist = BTreeMap::new();\n for result in &self.success {\n let count = dist.entry(result.status).or_insert(0);\n *count += 1;\n }\n dist\n }\n\n pub fn dns_dialup_stat(&self) -> MinMaxMean {\n self.success\n .iter()\n .filter_map(|r| r.connection_time.map(|ct| ct.dialup.as_secs_f64()))\n .collect()\n }\n\n pub fn dns_lookup_stat(&self) -> MinMaxMean {\n self.success\n .iter()\n .filter_map(|r| r.connection_time.map(|ct| ct.dns_lookup.as_secs_f64()))\n .collect()\n }\n\n pub fn first_byte_stat(&self) -> MinMaxMean {\n self.success\n .iter()\n .filter_map(|r| r.first_byte.map(|fb| (fb - r.start).as_secs_f64()))\n .collect()\n }\n\n pub fn total_data(&self) -> usize {\n self.success.iter().map(|r| r.len_bytes).sum()\n }\n\n pub fn size_per_request(&self) -> Option {\n self.success\n .iter()\n .map(|r| r.len_bytes as u64)\n .sum::()\n .checked_div(self.success.len() as u64)\n }\n\n pub fn duration_all_statistics(&self) -> Statistics {\n let mut data = self\n .success\n .iter()\n .map(|r| r.duration().as_secs_f64())\n .collect::>();\n Statistics::new(&mut data)\n }\n\n pub fn first_byte_all_statistics(&self) -> Statistics {\n let mut data = self\n .success\n .iter()\n .filter_map(|r| r.first_byte.map(|fb| (fb - r.start).as_secs_f64()))\n .collect::>();\n Statistics::new(&mut data)\n }\n\n pub fn duration_successful_statistics(&self) -> Statistics {\n let mut data = self\n .success\n .iter()\n .filter(|r| r.status.is_success())\n .map(|r| r.duration().as_secs_f64())\n .collect::>();\n\n Statistics::new(&mut data)\n }\n\n pub fn duration_not_successful_statistics(&self) -> Statistics {\n let mut data = self\n .success\n .iter()\n .filter(|r| !r.status.is_success())\n .map(|r| r.duration().as_secs_f64())\n .collect::>();\n\n Statistics::new(&mut data)\n }\n}\n\n#[cfg(test)]\nmod tests {\n use float_cmp::assert_approx_eq;\n use rand::SeedableRng;\n\n use super::*;\n use crate::client::{ClientError, ConnectionTime, RequestResult};\n use std::time::{Duration, Instant};\n\n fn build_mock_request_result(\n status: StatusCode,\n request_time: u64,\n connection_time_dns_lookup: u64,\n connection_time_dialup: u64,\n first_byte: u64,\n size: usize,\n ) -> Result {\n let now = Instant::now();\n Ok(RequestResult {\n rng: SeedableRng::seed_from_u64(0),\n start_latency_correction: None,\n start: now,\n connection_time: Some(ConnectionTime {\n dns_lookup: Duration::from_millis(connection_time_dns_lookup),\n dialup: Duration::from_millis(connection_time_dialup),\n }),\n first_byte: Some(now.checked_add(Duration::from_millis(first_byte)).unwrap()),\n end: now\n .checked_add(Duration::from_millis(request_time))\n .unwrap(),\n status,\n len_bytes: size,\n })\n }\n\n fn build_mock_request_results() -> ResultData {\n let mut results = ResultData::default();\n\n results.push(build_mock_request_result(\n StatusCode::OK,\n 1000,\n 200,\n 50,\n 300,\n 100,\n ));\n results.push(build_mock_request_result(\n StatusCode::BAD_REQUEST,\n 100000,\n 250,\n 100,\n 400,\n 200,\n ));\n results.push(build_mock_request_result(\n StatusCode::INTERNAL_SERVER_ERROR,\n 1000000,\n 300,\n 150,\n 500,\n 300,\n ));\n results\n }\n\n #[test]\n fn test_calculate_success_rate() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.success_rate(), 1.0);\n }\n\n #[test]\n fn test_calculate_slowest_request() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.latency_stat().max(), 1000.0);\n }\n\n #[test]\n fn test_calculate_average_request() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.latency_stat().mean(), 367.0);\n }\n\n #[test]\n fn test_calculate_total_data() {\n let res = build_mock_request_results();\n assert_eq!(res.total_data(), 600);\n }\n\n #[test]\n fn test_calculate_size_per_request() {\n let res = build_mock_request_results();\n assert_eq!(res.size_per_request(), Some(200));\n }\n\n #[test]\n fn test_calculate_connection_times_dns_dialup_average() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_dialup_stat().mean(), 0.1);\n }\n\n #[test]\n fn test_calculate_connection_times_dns_dialup_fastest() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_dialup_stat().min(), 0.05);\n }\n\n #[test]\n fn test_calculate_connection_times_dns_dialup_slowest() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_dialup_stat().max(), 0.15);\n }\n\n #[test]\n fn test_calculate_connection_times_dns_lookup_average() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_lookup_stat().mean(), 0.25);\n }\n\n #[test]\n fn test_calculate_connection_times_dns_lookup_fastest() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_lookup_stat().min(), 0.2);\n }\n\n #[test]\n fn test_calculate_connection_times_dns_lookup_slowest() {\n let res = build_mock_request_results();\n assert_approx_eq!(f64, res.dns_lookup_stat().max(), 0.3);\n }\n}\n" }, { "path": "src/timescale.rs", "content": "use std::{fmt, time::Duration};\n\n#[derive(Clone, Copy, PartialEq, Eq, Debug, PartialOrd, Ord)]\npub enum TimeScale {\n Nanosecond, // 1e-9\n Microsecond, // 1e-6\n Millisecond, // 1e-3\n Second, // 1\n TenSeconds, // 10\n Minute, // 60\n TenMinutes, // 600\n Hour, // 3600\n}\n\n#[derive(Clone, Copy, PartialEq, Eq, Debug)]\npub struct TimeLabel {\n pub x: usize,\n pub timescale: TimeScale,\n}\n\nimpl fmt::Display for TimeScale {\n fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {\n match self {\n TimeScale::Nanosecond => write!(f, \"ns\"),\n TimeScale::Microsecond => write!(f, \"us\"),\n TimeScale::Millisecond => write!(f, \"ms\"),\n TimeScale::Second => write!(f, \"sec\"),\n TimeScale::TenSeconds => write!(f, \"10 sec\"),\n TimeScale::Minute => write!(f, \"min\"),\n TimeScale::TenMinutes => write!(f, \"10 min\"),\n TimeScale::Hour => write!(f, \"hr\"),\n }\n }\n}\n\nimpl fmt::Display for TimeLabel {\n fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {\n match self {\n TimeLabel {\n x,\n timescale: TimeScale::Nanosecond,\n } => write!(f, \"{x}ns\"),\n TimeLabel {\n x,\n timescale: TimeScale::Microsecond,\n } => write!(f, \"{x}us\"),\n TimeLabel {\n x,\n timescale: TimeScale::Millisecond,\n } => write!(f, \"{x}ms\"),\n TimeLabel {\n x,\n timescale: TimeScale::Second,\n } => write!(f, \"{x}s\"),\n TimeLabel {\n x,\n timescale: TimeScale::TenSeconds,\n } => write!(f, \"{}s\", 10 * x),\n TimeLabel {\n x,\n timescale: TimeScale::Minute,\n } => write!(f, \"{x}m\"),\n TimeLabel {\n x,\n timescale: TimeScale::TenMinutes,\n } => write!(f, \"{}m\", 10 * x),\n TimeLabel {\n x,\n timescale: TimeScale::Hour,\n } => write!(f, \"{x}h\"),\n }\n }\n}\n\nimpl clap::ValueEnum for TimeScale {\n fn value_variants<'a>() -> &'a [Self] {\n &[\n Self::Nanosecond,\n Self::Microsecond,\n Self::Millisecond,\n Self::Second,\n Self::Minute,\n Self::Hour,\n ]\n }\n\n fn to_possible_value(&self) -> Option {\n match self {\n TimeScale::Nanosecond => Some(clap::builder::PossibleValue::new(\"ns\")),\n TimeScale::Microsecond => Some(clap::builder::PossibleValue::new(\"us\")),\n TimeScale::Millisecond => Some(clap::builder::PossibleValue::new(\"ms\")),\n TimeScale::Second => Some(clap::builder::PossibleValue::new(\"s\")),\n TimeScale::Minute => Some(clap::builder::PossibleValue::new(\"m\")),\n TimeScale::Hour => Some(clap::builder::PossibleValue::new(\"h\")),\n TimeScale::TenSeconds | TimeScale::TenMinutes => None,\n }\n }\n}\n\nimpl TimeScale {\n pub fn as_secs_f64(&self) -> f64 {\n match self {\n TimeScale::Nanosecond => 1e-9,\n TimeScale::Microsecond => 1e-6,\n TimeScale::Millisecond => 1e-3,\n TimeScale::Second => 1.0,\n TimeScale::TenSeconds => 10.0,\n TimeScale::Minute => 60.0,\n TimeScale::TenMinutes => 10.0 * 60.0,\n TimeScale::Hour => 60.0 * 60.0,\n }\n }\n\n /// From seconds as f64\n pub fn from_f64(seconds: f64) -> Self {\n for ts in &[\n TimeScale::Hour,\n TimeScale::TenMinutes,\n TimeScale::Minute,\n TimeScale::TenSeconds,\n TimeScale::Second,\n TimeScale::Millisecond,\n TimeScale::Microsecond,\n TimeScale::Nanosecond,\n ] {\n if seconds > ts.as_secs_f64() {\n return *ts;\n }\n }\n TimeScale::Nanosecond\n }\n\n pub fn from_elapsed(duration: Duration) -> Self {\n Self::from_f64(duration.as_secs_f64())\n }\n\n pub fn inc(&self) -> Self {\n match self {\n TimeScale::Nanosecond => TimeScale::Microsecond,\n TimeScale::Microsecond => TimeScale::Millisecond,\n TimeScale::Millisecond => TimeScale::Second,\n TimeScale::Second => TimeScale::TenSeconds,\n TimeScale::TenSeconds => TimeScale::Minute,\n TimeScale::Minute => TimeScale::TenMinutes,\n TimeScale::TenMinutes => TimeScale::Hour,\n TimeScale::Hour => TimeScale::Hour,\n }\n }\n\n pub fn dec(&self) -> Self {\n match self {\n TimeScale::Nanosecond => TimeScale::Nanosecond,\n TimeScale::Microsecond => TimeScale::Nanosecond,\n TimeScale::Millisecond => TimeScale::Microsecond,\n TimeScale::Second => TimeScale::Millisecond,\n TimeScale::TenSeconds => TimeScale::Second,\n TimeScale::Minute => TimeScale::TenSeconds,\n TimeScale::TenMinutes => TimeScale::Minute,\n TimeScale::Hour => TimeScale::TenMinutes,\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n fn assert_timescale_correct_for_seconds_range(\n range: [f64; 2],\n expected_timescale: TimeScale,\n expected_timescale_str: &str,\n expected_timescale_as_secs: f64,\n ) {\n for durations in range {\n let timescale = TimeScale::from_elapsed(Duration::from_secs_f64(durations));\n assert_eq!(timescale, expected_timescale);\n assert_eq!(format!(\"{timescale}\"), expected_timescale_str);\n assert_eq!(timescale.as_secs_f64(), expected_timescale_as_secs);\n }\n }\n\n #[test]\n fn test_timescale_ranges() {\n assert_timescale_correct_for_seconds_range(\n [f64::MIN_POSITIVE, 1e-6],\n TimeScale::Nanosecond,\n \"ns\",\n 1e-9,\n );\n assert_timescale_correct_for_seconds_range(\n [0.000_001_1, 1e-3],\n TimeScale::Microsecond,\n \"us\",\n 1e-6,\n );\n assert_timescale_correct_for_seconds_range(\n [0.001_1, 1.0],\n TimeScale::Millisecond,\n \"ms\",\n 1e-3,\n );\n assert_timescale_correct_for_seconds_range([1.1, 10.0], TimeScale::Second, \"sec\", 1.0);\n assert_timescale_correct_for_seconds_range(\n [10.1, 60.0],\n TimeScale::TenSeconds,\n \"10 sec\",\n 10.0,\n );\n assert_timescale_correct_for_seconds_range([60.1, 600.0], TimeScale::Minute, \"min\", 60.0);\n assert_timescale_correct_for_seconds_range(\n [600.1, 3600.0],\n TimeScale::TenMinutes,\n \"10 min\",\n 600.0,\n );\n assert_timescale_correct_for_seconds_range(\n [3600.1, 31536000.0],\n TimeScale::Hour,\n \"hr\",\n 3600.0,\n );\n }\n\n #[test]\n fn test_timescale_inc() {\n let timescale = TimeScale::from_elapsed(Duration::from_secs_f64(1e-10));\n let timescale_microsecond = timescale.inc();\n assert_eq!(timescale_microsecond, TimeScale::Microsecond);\n let timescale_millisecond = timescale_microsecond.inc();\n assert_eq!(timescale_millisecond, TimeScale::Millisecond);\n let timescale_second = timescale_millisecond.inc();\n assert_eq!(timescale_second, TimeScale::Second);\n let timescale_ten_seconds = timescale_second.inc();\n assert_eq!(timescale_ten_seconds, TimeScale::TenSeconds);\n let timescale_minute = timescale_ten_seconds.inc();\n assert_eq!(timescale_minute, TimeScale::Minute);\n let timescale_ten_minutes = timescale_minute.inc();\n assert_eq!(timescale_ten_minutes, TimeScale::TenMinutes);\n let timescale_hour = timescale_ten_minutes.inc();\n assert_eq!(timescale_hour, TimeScale::Hour);\n }\n\n #[test]\n fn test_timescale_dec() {\n let timescale = TimeScale::from_elapsed(Duration::from_secs_f64(31536000.0));\n let timescale_ten_minutes = timescale.dec();\n assert_eq!(timescale_ten_minutes, TimeScale::TenMinutes);\n let timescale_minute = timescale_ten_minutes.dec();\n assert_eq!(timescale_minute, TimeScale::Minute);\n let timescale_ten_seconds = timescale_minute.dec();\n assert_eq!(timescale_ten_seconds, TimeScale::TenSeconds);\n let timescale_second = timescale_ten_seconds.dec();\n assert_eq!(timescale_second, TimeScale::Second);\n let timescale_millisecond = timescale_second.dec();\n assert_eq!(timescale_millisecond, TimeScale::Millisecond);\n let timescale_microsecond = timescale_millisecond.dec();\n assert_eq!(timescale_microsecond, TimeScale::Microsecond);\n let timescale_nanosecond = timescale_microsecond.dec();\n assert_eq!(timescale_nanosecond, TimeScale::Nanosecond);\n }\n}\n" }, { "path": "src/tls_config.rs", "content": "#[cfg(feature = \"rustls\")]\npub struct RuslsConfigs {\n no_alpn: std::sync::Arc,\n alpn_h2: std::sync::Arc,\n alpn_h3: std::sync::Arc,\n}\n\n#[cfg(feature = \"rustls\")]\nimpl RuslsConfigs {\n pub fn new(\n insecure: bool,\n cacert_pem: Option<&[u8]>,\n client_auth: Option<(&[u8], &[u8])>,\n ) -> Self {\n use rustls_pki_types::pem::PemObject;\n use std::sync::Arc;\n\n let mut root_cert_store = rustls::RootCertStore::empty();\n for cert in rustls_native_certs::load_native_certs().expect(\"could not load platform certs\")\n {\n root_cert_store.add(cert).unwrap();\n }\n\n if let Some(cacert_pem) = cacert_pem {\n for der in rustls_pki_types::CertificateDer::pem_slice_iter(cacert_pem) {\n root_cert_store.add(der.unwrap()).unwrap();\n }\n }\n\n let _ = rustls::crypto::CryptoProvider::install_default(\n rustls::crypto::aws_lc_rs::default_provider(),\n );\n let builder = rustls::ClientConfig::builder().with_root_certificates(root_cert_store);\n\n let mut config = if let Some((cert, key)) = client_auth {\n let certs = rustls_pki_types::CertificateDer::pem_slice_iter(cert)\n .collect::, _>>()\n .unwrap();\n let key = rustls_pki_types::PrivateKeyDer::from_pem_slice(key).unwrap();\n\n builder.with_client_auth_cert(certs, key).unwrap()\n } else {\n builder.with_no_client_auth()\n };\n if insecure {\n config\n .dangerous()\n .set_certificate_verifier(Arc::new(AcceptAnyServerCert));\n }\n\n let mut no_alpn = config.clone();\n no_alpn.alpn_protocols = vec![];\n let mut alpn_h2 = config.clone();\n alpn_h2.alpn_protocols = vec![b\"h2\".to_vec()];\n let mut alpn_h3 = config;\n alpn_h3.alpn_protocols = vec![b\"h3\".to_vec()];\n alpn_h3.enable_early_data = true;\n Self {\n no_alpn: Arc::new(no_alpn),\n alpn_h2: Arc::new(alpn_h2),\n alpn_h3: Arc::new(alpn_h3),\n }\n }\n\n pub fn config(&self, http: hyper::http::Version) -> &std::sync::Arc {\n use hyper::http;\n match http {\n http::Version::HTTP_09 | http::Version::HTTP_10 | http::Version::HTTP_11 => {\n &self.no_alpn\n }\n http::Version::HTTP_2 => &self.alpn_h2,\n http::Version::HTTP_3 => &self.alpn_h3,\n _ => panic!(\"nonsupported HTTP version\"),\n }\n }\n}\n\n#[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\npub struct NativeTlsConnectors {\n pub no_alpn: tokio_native_tls::TlsConnector,\n pub alpn_h2: tokio_native_tls::TlsConnector,\n}\n\n#[cfg(all(feature = \"native-tls\", not(feature = \"rustls\")))]\nimpl NativeTlsConnectors {\n pub fn new(\n insecure: bool,\n cacert_pem: Option<&[u8]>,\n client_auth: Option<(&[u8], &[u8])>,\n ) -> Self {\n let new = |is_http2: bool| {\n let mut connector_builder = native_tls::TlsConnector::builder();\n\n if let Some(cacert_pem) = cacert_pem {\n let cert = native_tls::Certificate::from_pem(cacert_pem)\n .expect(\"Failed to parse cacert_pem\");\n connector_builder.add_root_certificate(cert);\n }\n\n if insecure {\n connector_builder\n .danger_accept_invalid_certs(true)\n .danger_accept_invalid_hostnames(true);\n }\n\n if let Some((cert, key)) = client_auth {\n let cert = native_tls::Identity::from_pkcs8(cert, key)\n .expect(\"Failed to parse client_auth cert/key\");\n connector_builder.identity(cert);\n }\n\n if is_http2 {\n connector_builder.request_alpns(&[\"h2\"]);\n }\n\n connector_builder\n .build()\n .expect(\"Failed to build native_tls::TlsConnector\")\n .into()\n };\n\n Self {\n no_alpn: new(false),\n alpn_h2: new(true),\n }\n }\n\n pub fn connector(&self, is_http2: bool) -> &tokio_native_tls::TlsConnector {\n if is_http2 {\n &self.alpn_h2\n } else {\n &self.no_alpn\n }\n }\n}\n\n/// A server certificate verifier that accepts any certificate.\n#[cfg(feature = \"rustls\")]\n#[derive(Debug)]\npub struct AcceptAnyServerCert;\n\n#[cfg(feature = \"rustls\")]\nimpl rustls::client::danger::ServerCertVerifier for AcceptAnyServerCert {\n fn verify_server_cert(\n &self,\n _end_entity: &rustls_pki_types::CertificateDer<'_>,\n _intermediates: &[rustls_pki_types::CertificateDer<'_>],\n _server_name: &rustls_pki_types::ServerName<'_>,\n _ocsp_response: &[u8],\n _now: rustls_pki_types::UnixTime,\n ) -> Result {\n Ok(rustls::client::danger::ServerCertVerified::assertion())\n }\n\n fn verify_tls12_signature(\n &self,\n _message: &[u8],\n _cert: &rustls_pki_types::CertificateDer<'_>,\n _dss: &rustls::DigitallySignedStruct,\n ) -> Result {\n Ok(rustls::client::danger::HandshakeSignatureValid::assertion())\n }\n\n fn verify_tls13_signature(\n &self,\n _message: &[u8],\n _cert: &rustls_pki_types::CertificateDer<'_>,\n _dss: &rustls::DigitallySignedStruct,\n ) -> Result {\n Ok(rustls::client::danger::HandshakeSignatureValid::assertion())\n }\n\n fn supported_verify_schemes(&self) -> Vec {\n rustls::crypto::CryptoProvider::get_default()\n .unwrap()\n .signature_verification_algorithms\n .supported_schemes()\n }\n}\n" }, { "path": "src/url_generator.rs", "content": "use std::{borrow::Cow, string::FromUtf8Error};\n\nuse rand::prelude::*;\nuse rand_regex::Regex;\nuse thiserror::Error;\nuse url::{ParseError, Url};\n\n#[derive(Clone, Debug)]\npub enum UrlGenerator {\n Static(Url),\n MultiStatic(Vec),\n Dynamic(Regex),\n}\n\n#[derive(Error, Debug)]\npub enum UrlGeneratorError {\n #[error(\"{0}, generated url: {1}\")]\n Parse(ParseError, String),\n #[error(transparent)]\n FromUtf8(#[from] FromUtf8Error),\n #[error(\"No valid URLs found\")]\n NoURLs(),\n #[error(transparent)]\n Io(#[from] std::io::Error),\n}\n\nimpl UrlGenerator {\n pub fn new_static(url: Url) -> Self {\n Self::Static(url)\n }\n\n pub fn new_multi_static(urls: Vec) -> Self {\n assert!(!urls.is_empty());\n Self::MultiStatic(urls)\n }\n\n pub fn new_dynamic(regex: Regex) -> Self {\n Self::Dynamic(regex)\n }\n\n pub fn generate(&self, rng: &mut R) -> Result, UrlGeneratorError> {\n match self {\n Self::Static(url) => Ok(Cow::Borrowed(url)),\n Self::MultiStatic(urls) => {\n if let Some(random_url) = urls.choose(rng) {\n Ok(Cow::Borrowed(random_url))\n } else {\n Err(UrlGeneratorError::NoURLs())\n }\n }\n Self::Dynamic(regex) => {\n let generated = Distribution::>::sample(regex, rng)?;\n Ok(Cow::Owned(\n Url::parse(generated.as_str())\n .map_err(|e| UrlGeneratorError::Parse(e, generated))?,\n ))\n }\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use crate::pcg64si::Pcg64Si;\n\n use super::*;\n use rand_regex::Regex as RandRegex;\n use regex::Regex;\n use std::net::Ipv4Addr;\n use url::{Host, Url};\n\n #[test]\n fn test_url_generator_static() {\n let url_generator = UrlGenerator::new_static(Url::parse(\"http://127.0.0.1/test\").unwrap());\n let url = url_generator.generate(&mut rand::rng()).unwrap();\n assert_eq!(url.host(), Some(Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1))));\n assert_eq!(url.path(), \"/test\");\n }\n\n #[test]\n fn test_url_generator_multistatic() {\n let urls = [\n \"http://127.0.0.1/a1\",\n \"http://127.0.0.1/b2\",\n \"http://127.0.0.1/c3\",\n ];\n\n let url_generator =\n UrlGenerator::new_multi_static(urls.iter().map(|u| Url::parse(u).unwrap()).collect());\n\n for _ in 0..10 {\n let url = url_generator.generate(&mut rand::rng()).unwrap();\n assert_eq!(url.host(), Some(Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1))));\n assert!(urls.contains(&url.as_str()));\n }\n }\n\n #[test]\n fn test_url_generator_dynamic() {\n let path_regex = \"/[a-z][a-z][0-9]\";\n let url_generator = UrlGenerator::new_dynamic(\n RandRegex::compile(&format!(r\"http://127\\.0\\.0\\.1{path_regex}\"), 4).unwrap(),\n );\n let url = url_generator.generate(&mut rand::rng()).unwrap();\n assert_eq!(url.host(), Some(Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1))));\n assert!(\n Regex::new(path_regex)\n .unwrap()\n .captures(url.path())\n .is_some()\n );\n }\n\n #[test]\n fn test_url_generator_dynamic_consistency() {\n let url_generator = UrlGenerator::new_dynamic(\n RandRegex::compile(r\"http://127\\.0\\.0\\.1/[a-z][a-z][0-9]\", 4).unwrap(),\n );\n\n for _ in 0..100 {\n let rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n\n assert_eq!(\n url_generator.generate(&mut rng.clone()).unwrap(),\n url_generator.generate(&mut rng.clone()).unwrap()\n );\n }\n }\n\n #[test]\n fn test_url_generator_multi_consistency() {\n let urls = [\n \"http://example.com/a1\",\n \"http://example.com/a2\",\n \"http://example.com/a3\",\n \"http://example.com/a4\",\n \"http://example.com/a5\",\n ];\n let url_generator =\n UrlGenerator::new_multi_static(urls.iter().map(|u| Url::parse(u).unwrap()).collect());\n\n for _ in 0..100 {\n let rng: Pcg64Si = SeedableRng::from_rng(&mut rand::rng());\n\n assert_eq!(\n url_generator.generate(&mut rng.clone()).unwrap(),\n url_generator.generate(&mut rng.clone()).unwrap()\n );\n }\n }\n}\n" }, { "path": "tests/common/mod.rs", "content": "use std::{net::SocketAddr, sync::Arc};\n\nuse bytes::{Buf, Bytes};\nuse http::{Request, Response};\nuse kanal::Sender;\nuse rustls::pki_types::{CertificateDer, PrivateKeyDer};\n\nuse h3::{quic::BidiStream, server::RequestStream};\nuse h3_quinn::quinn::{self, crypto::rustls::QuicServerConfig};\n\nstatic ALPN: &[u8] = b\"h3\";\n\n// This would be much cleaner if it took `process_request` as a callback, similar to the hyper service_fn.\npub async fn h3_server(\n tx: Sender>,\n port: u16,\n) -> Result<(), Box> {\n let listen = SocketAddr::new(\"127.0.0.1\".parse().unwrap(), port);\n\n // Get the directory of the current file\n let current_file = file!();\n let current_dir = std::path::Path::new(current_file)\n .parent()\n .unwrap_or_else(|| std::path::Path::new(\"\"));\n\n // Construct paths to cert and key files\n let cert_path = current_dir.join(\"server.cert\");\n let key_path = current_dir.join(\"server.key\");\n\n // both cert and key must be DER-encoded\n let cert = CertificateDer::from(std::fs::read(&cert_path)?);\n let key = PrivateKeyDer::try_from(std::fs::read(&key_path)?)?;\n\n let _ = rustls::crypto::CryptoProvider::install_default(\n rustls::crypto::aws_lc_rs::default_provider(),\n );\n let mut tls_config = rustls::ServerConfig::builder()\n .with_no_client_auth()\n .with_single_cert(vec![cert], key)?;\n\n tls_config.max_early_data_size = u32::MAX;\n tls_config.alpn_protocols = vec![ALPN.into()];\n\n let server_config =\n quinn::ServerConfig::with_crypto(Arc::new(QuicServerConfig::try_from(tls_config)?));\n let endpoint = quinn::Endpoint::server(server_config, listen)?;\n\n // handle incoming connections and requests\n while let Some(new_conn) = endpoint.accept().await {\n let tx = tx.clone();\n\n let _ = tokio::spawn(async move {\n match new_conn.await {\n Ok(conn) => {\n let mut h3_conn = h3::server::Connection::new(h3_quinn::Connection::new(conn))\n .await\n .unwrap();\n\n let tx = tx.clone();\n match h3_conn.accept().await {\n Ok(Some(request_resolver)) => {\n let (req, stream) = request_resolver.resolve_request().await.unwrap();\n process_request(req, stream, tx).await\n }\n\n // indicating no more streams to be received\n Ok(None) => Ok(()),\n\n Err(_err) => {\n unimplemented!()\n // error!(\"error on accept {}\", err);\n /*\n match err.get_error_level() {\n ErrorLevel::ConnectionError => break,\n ErrorLevel::StreamError => continue,\n }\n */\n }\n }\n }\n Err(_err) => Ok(()),\n }\n })\n .await?;\n }\n\n // shut down gracefully\n // wait for connections to be closed before exiting\n endpoint.wait_idle().await;\n\n Ok(())\n}\n\nasync fn process_request(\n req: Request<()>,\n mut stream: RequestStream,\n tx: Sender>,\n) -> Result<(), h3::error::StreamError>\nwhere\n T: BidiStream,\n{\n let (parts, _) = req.into_parts();\n let mut body_bytes = bytes::BytesMut::new();\n\n while let Some(mut chunk) = stream.recv_data().await? {\n let bytes = chunk.copy_to_bytes(chunk.remaining());\n body_bytes.extend_from_slice(&bytes);\n }\n let body = body_bytes.freeze();\n let req = Request::from_parts(parts, body);\n\n tx.send(req).unwrap();\n let resp = Response::new(());\n stream.send_response(resp).await?;\n stream.send_data(\"Hello world\".into()).await?;\n stream.finish().await\n}\n" }, { "path": "tests/tests.rs", "content": "use std::{\n convert::Infallible,\n error::Error as StdError,\n fs::File,\n future::Future,\n io::Write,\n net::{Ipv6Addr, SocketAddr},\n str::FromStr,\n sync::{Arc, OnceLock, atomic::AtomicU16},\n};\n\nuse axum::{Router, extract::Path, response::Redirect, routing::get};\nuse bytes::Bytes;\nuse clap::Parser;\nuse http::{HeaderMap, Request, Response};\nuse http_body_util::BodyExt;\nuse http_mitm_proxy::MitmProxy;\nuse hyper::{\n body::{Body, Incoming},\n http,\n service::{HttpService, service_fn},\n};\nuse hyper_util::rt::{TokioExecutor, TokioIo};\nuse rstest::rstest;\nuse rstest_reuse::{self, *};\n#[cfg(feature = \"http3\")]\nmod common;\n\nasync fn run<'a>(args: impl Iterator) {\n let opts = oha::Opts::parse_from(\n [\"oha\", \"--no-tui\", \"--output-format\", \"quiet\"]\n .into_iter()\n .chain(args),\n );\n\n oha::run(opts).await.unwrap();\n}\n\n// Port 5111- is reserved for testing\nstatic PORT: AtomicU16 = AtomicU16::new(5111);\n\nfn next_port() -> u16 {\n PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)\n}\n\n#[ctor::ctor]\nfn install_crypto_provider() {\n static INSTALL: OnceLock<()> = OnceLock::new();\n INSTALL.get_or_init(|| {\n let _ = rustls::crypto::CryptoProvider::install_default(\n rustls::crypto::aws_lc_rs::default_provider(),\n );\n });\n}\n\nasync fn bind_port(port: u16) -> tokio::net::TcpListener {\n let addr = SocketAddr::new(\"127.0.0.1\".parse().unwrap(), port);\n\n tokio::net::TcpListener::bind(addr).await.unwrap()\n}\n\nasync fn bind_port_and_increment() -> (tokio::net::TcpListener, u16) {\n let port = next_port();\n let listener = bind_port(port).await;\n (listener, port)\n}\n\nasync fn bind_port_ipv6(port: u16) -> tokio::net::TcpListener {\n let addr = SocketAddr::new(std::net::IpAddr::V6(Ipv6Addr::LOCALHOST), port);\n\n tokio::net::TcpListener::bind(addr).await.unwrap()\n}\n\n#[derive(Clone, Copy, PartialEq)]\nenum HttpWorkType {\n H1,\n H2,\n #[cfg(feature = \"http3\")]\n H3,\n}\n\nfn http_work_type(args: &[&str]) -> HttpWorkType {\n // Check for HTTP/2\n if args.contains(&\"--http2\") || args.windows(2).any(|w| w == [\"--http-version\", \"2\"]) {\n return HttpWorkType::H2;\n }\n\n // Check for HTTP/3 when the feature is enabled\n #[cfg(feature = \"http3\")]\n if args.contains(&\"--http3\") || args.windows(2).any(|w| w == [\"--http-version\", \"3\"]) {\n return HttpWorkType::H3;\n }\n\n // Default to HTTP/1.1\n HttpWorkType::H1\n}\n\n#[cfg(feature = \"http3\")]\n#[template]\n#[rstest]\n#[case(\"1.1\")]\n#[case(\"2\")]\n#[case(\"3\")]\nfn test_all_http_versions(#[case] http_version_param: &str) {}\n\n#[cfg(not(feature = \"http3\"))]\n#[template]\n#[rstest]\n#[case(\"1.1\")]\n#[case(\"2\")]\nfn test_all_http_versions(#[case] http_version_param: &str) {}\n\nasync fn get_req(path: &str, args: &[&str]) -> Request {\n let (tx, rx) = kanal::unbounded();\n\n let port = next_port();\n\n let work_type = http_work_type(args);\n let listener = bind_port(port).await;\n\n tokio::spawn(async move {\n match work_type {\n HttpWorkType::H2 => loop {\n let (tcp, _) = listener.accept().await.unwrap();\n let tx = tx.clone();\n let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new())\n .serve_connection(\n TokioIo::new(tcp),\n service_fn(move |req: Request| {\n let tx = tx.clone();\n async move {\n let (parts, body) = req.into_parts();\n let body_bytes = body.collect().await.unwrap().to_bytes();\n let req = Request::from_parts(parts, body_bytes);\n tx.send(req).unwrap();\n Ok::<_, Infallible>(Response::new(\"Hello World\".to_string()))\n }\n }),\n )\n .await;\n },\n HttpWorkType::H1 => {\n let (tcp, _) = listener.accept().await.unwrap();\n hyper::server::conn::http1::Builder::new()\n .serve_connection(\n TokioIo::new(tcp),\n service_fn(move |req: Request| {\n let tx = tx.clone();\n\n async move {\n let (parts, body) = req.into_parts();\n let body_bytes = body.collect().await.unwrap().to_bytes();\n let req = Request::from_parts(parts, body_bytes);\n tx.send(req).unwrap();\n Ok::<_, Infallible>(Response::new(\"Hello World\".to_string()))\n }\n }),\n )\n .await\n .unwrap();\n }\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => {\n drop(listener);\n common::h3_server(tx, port).await.unwrap();\n }\n }\n });\n\n let mut args = args.iter().map(|s| s.to_string()).collect::>();\n args.push(\"-n\".to_string());\n args.push(\"1\".to_string());\n match work_type {\n HttpWorkType::H1 | HttpWorkType::H2 => {\n args.push(format!(\"http://127.0.0.1:{port}{path}\"));\n }\n #[cfg(feature = \"http3\")]\n HttpWorkType::H3 => {\n args.push(\"--insecure\".to_string());\n args.push(format!(\"https://127.0.0.1:{port}{path}\"));\n }\n }\n\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap()\n}\n\nasync fn redirect(n: usize, is_relative: bool, limit: usize) -> bool {\n let (tx, rx) = kanal::unbounded();\n\n let (listener, port) = bind_port_and_increment().await;\n\n let app = Router::new().route(\n \"/{n}\",\n get(move |Path(x): Path| async move {\n Ok::<_, Infallible>(if x == n {\n tx.send(()).unwrap();\n Redirect::permanent(\"/end\")\n } else if is_relative {\n Redirect::permanent(&format!(\"/{}\", x + 1))\n } else {\n Redirect::permanent(&format!(\"http://localhost:{}/{}\", port, x + 1))\n })\n }),\n );\n\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n \"--redirect\".to_string(),\n limit.to_string(),\n format!(\"http://127.0.0.1:{port}/0\"),\n ];\n\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().is_some()\n}\n\nasync fn get_host_with_connect_to(host: &'static str) -> String {\n let (tx, rx) = kanal::unbounded();\n\n let app = Router::new().route(\n \"/\",\n get(|header: HeaderMap| async move {\n tx.send(header.get(\"host\").unwrap().to_str().unwrap().to_string())\n .unwrap();\n \"Hello World\"\n }),\n );\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n format!(\"http://{host}/\"),\n \"--connect-to\".to_string(),\n format!(\"{host}:80:localhost:{port}\"),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap()\n}\n\nasync fn get_host_with_connect_to_ipv6_target(host: &'static str) -> String {\n let (tx, rx) = kanal::unbounded();\n let app = Router::new().route(\n \"/\",\n get(|header: HeaderMap| async move {\n tx.send(header.get(\"host\").unwrap().to_str().unwrap().to_string())\n .unwrap();\n \"Hello World\"\n }),\n );\n\n let port = next_port();\n let listener = bind_port_ipv6(port).await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n format!(\"http://{host}/\"),\n \"--connect-to\".to_string(),\n format!(\"{host}:80:[::1]:{port}\"),\n ];\n\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap()\n}\n\nasync fn get_host_with_connect_to_ipv6_requested() -> String {\n let (tx, rx) = kanal::unbounded();\n let app = Router::new().route(\n \"/\",\n get(|header: HeaderMap| async move {\n tx.send(header.get(\"host\").unwrap().to_str().unwrap().to_string())\n .unwrap();\n \"Hello World\"\n }),\n );\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n \"http://[::1]/\".to_string(),\n \"--connect-to\".to_string(),\n format!(\"[::1]:80:localhost:{port}\"),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap()\n}\n\nasync fn get_host_with_connect_to_redirect(host: &'static str) -> String {\n let (tx, rx) = kanal::unbounded();\n\n let app = Router::new()\n .route(\n \"/source\",\n get(move || async move { Redirect::permanent(&format!(\"http://{host}/destination\")) }),\n )\n .route(\n \"/destination\",\n get(move || async move {\n tx.send(host.to_string()).unwrap();\n \"Hello World\"\n }),\n );\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n format!(\"http://{host}/source\"),\n \"--connect-to\".to_string(),\n format!(\"{host}:80:localhost:{port}\"),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap()\n}\n\nasync fn test_request_count(args: &[&str]) -> usize {\n let (tx, rx) = kanal::unbounded();\n\n let app = Router::new().route(\n \"/\",\n get(|| async move {\n tx.send(()).unwrap();\n \"Success\"\n }),\n );\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let mut args: Vec = args.iter().map(|s| s.to_string()).collect();\n args.push(format!(\"http://127.0.0.1:{port}\"));\n run(args.iter().map(|s| s.as_str())).await;\n\n let mut count = 0;\n while let Ok(Some(())) = rx.try_recv() {\n count += 1;\n }\n count\n}\n\n// Randomly spread 100 requests on two matching --connect-to targets, and return a count for each\nasync fn distribution_on_two_matching_connect_to(host: &'static str) -> (i32, i32) {\n let (tx1, rx1) = kanal::unbounded();\n let (tx2, rx2) = kanal::unbounded();\n\n let app1 = Router::new().route(\n \"/\",\n get(move || async move {\n tx1.send(()).unwrap();\n \"Success1\"\n }),\n );\n\n let app2 = Router::new().route(\n \"/\",\n get(move || async move {\n tx2.send(()).unwrap();\n \"Success2\"\n }),\n );\n\n let (listener1, port1) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener1, app1).await });\n\n let (listener2, port2) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener2, app2).await });\n\n let args = [\n \"--disable-keepalive\".to_string(),\n \"-n\".to_string(),\n \"100\".to_string(),\n format!(\"http://{host}/\"),\n \"--connect-to\".to_string(),\n format!(\"{host}:80:localhost:{port1}\"),\n \"--connect-to\".to_string(),\n format!(\"{host}:80:localhost:{port2}\"),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n let mut count1 = 0;\n let mut count2 = 0;\n loop {\n if rx1.try_recv().unwrap().is_some() {\n count1 += 1;\n } else if rx2.try_recv().unwrap().is_some() {\n count2 += 1;\n } else {\n break;\n }\n }\n (count1, count2)\n}\n\n#[apply(test_all_http_versions)]\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_enable_compression_default(http_version_param: &str) {\n let req = get_req(\"/\", &[\"--http-version\", http_version_param]).await;\n let accept_encoding: Vec<&str> = req\n .headers()\n .get(\"accept-encoding\")\n .unwrap()\n .to_str()\n .unwrap()\n .split(\", \")\n .collect();\n\n assert!(accept_encoding.contains(&\"gzip\"));\n assert!(accept_encoding.contains(&\"br\"));\n}\n\n#[apply(test_all_http_versions)]\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_setting_custom_header(http_version_param: &str) {\n let req = get_req(\n \"/\",\n &[\"--http-version\", http_version_param, \"-H\", \"foo: bar\"],\n )\n .await;\n assert_eq!(req.headers().get(\"foo\").unwrap().to_str().unwrap(), \"bar\");\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\n#[apply(test_all_http_versions)]\nasync fn test_setting_accept_header(http_version_param: &str) {\n let req = get_req(\n \"/\",\n &[\"-A\", \"text/html\", \"--http-version\", http_version_param],\n )\n .await;\n assert_eq!(\n req.headers().get(\"accept\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n let req = get_req(\n \"/\",\n &[\n \"-H\",\n \"accept:text/html\",\n \"--http-version\",\n http_version_param,\n ],\n )\n .await;\n assert_eq!(\n req.headers().get(\"accept\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\n#[apply(test_all_http_versions)]\nasync fn test_setting_body(http_version_param: &str) {\n let req = get_req(\n \"/\",\n &[\"-d\", \"hello body\", \"--http-version\", http_version_param],\n )\n .await;\n assert_eq!(\n req.into_body(),\n &b\"hello body\"[..] /* This looks dirty... Any suggestion? */\n );\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_setting_content_type_header() {\n let req = get_req(\"/\", &[\"-T\", \"text/html\"]).await;\n assert_eq!(\n req.headers().get(\"content-type\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n let req = get_req(\"/\", &[\"-H\", \"content-type:text/html\"]).await;\n assert_eq!(\n req.headers().get(\"content-type\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n\n let req = get_req(\"/\", &[\"--http2\", \"-T\", \"text/html\"]).await;\n assert_eq!(\n req.headers().get(\"content-type\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n let req = get_req(\"/\", &[\"--http2\", \"-H\", \"content-type:text/html\"]).await;\n assert_eq!(\n req.headers().get(\"content-type\").unwrap().to_str().unwrap(),\n \"text/html\"\n );\n}\n\n#[apply(test_all_http_versions)]\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_setting_basic_auth(http_version_param: &str) {\n let req = get_req(\n \"/\",\n &[\"-a\", \"hatoo:pass\", \"--http-version\", http_version_param],\n )\n .await;\n assert_eq!(\n req.headers()\n .get(\"authorization\")\n .unwrap()\n .to_str()\n .unwrap(),\n \"Basic aGF0b286cGFzcw==\"\n );\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_setting_host() {\n let req = get_req(\"/\", &[\"--host\", \"hatoo.io\"]).await;\n assert_eq!(\n req.headers().get(\"host\").unwrap().to_str().unwrap(),\n \"hatoo.io\"\n );\n\n let req = get_req(\"/\", &[\"-H\", \"host:hatoo.io\"]).await;\n assert_eq!(\n req.headers().get(\"host\").unwrap().to_str().unwrap(),\n \"hatoo.io\"\n );\n\n // You shouldn't set host header when using HTTP/2\n // Use --connect-to instead\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_setting_method() {\n assert_eq!(get_req(\"/\", &[]).await.method(), http::method::Method::GET);\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"GET\"]).await.method(),\n http::method::Method::GET\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"POST\"]).await.method(),\n http::method::Method::POST\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"CONNECT\"]).await.method(),\n http::method::Method::CONNECT\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"DELETE\"]).await.method(),\n http::method::Method::DELETE\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"HEAD\"]).await.method(),\n http::method::Method::HEAD\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"OPTIONS\"]).await.method(),\n http::method::Method::OPTIONS\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"PATCH\"]).await.method(),\n http::method::Method::PATCH\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"PUT\"]).await.method(),\n http::method::Method::PUT\n );\n assert_eq!(\n get_req(\"/\", &[\"-m\", \"TRACE\"]).await.method(),\n http::method::Method::TRACE\n );\n\n assert_eq!(\n get_req(\"/\", &[\"--http2\"]).await.method(),\n http::method::Method::GET\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"GET\"]).await.method(),\n http::method::Method::GET\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"POST\"]).await.method(),\n http::method::Method::POST\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"DELETE\"]).await.method(),\n http::method::Method::DELETE\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"HEAD\"]).await.method(),\n http::method::Method::HEAD\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"OPTIONS\"]).await.method(),\n http::method::Method::OPTIONS\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"PATCH\"]).await.method(),\n http::method::Method::PATCH\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"PUT\"]).await.method(),\n http::method::Method::PUT\n );\n assert_eq!(\n get_req(\"/\", &[\"--http2\", \"-m\", \"TRACE\"]).await.method(),\n http::method::Method::TRACE\n );\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_query() {\n assert_eq!(\n get_req(\"/index?a=b&c=d\", &[]).await.uri().to_string(),\n \"/index?a=b&c=d\".to_string()\n );\n\n assert_eq!(\n get_req(\"/index?a=b&c=d\", &[\"--http2\"])\n .await\n .uri()\n .to_string()\n .split('/')\n .next_back()\n .unwrap(),\n \"index?a=b&c=d\".to_string()\n );\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_query_rand_regex() {\n let req = get_req(\"/[a-z][0-9][a-z]\", &[\"--rand-regex-url\"]).await;\n let chars = req\n .uri()\n .to_string()\n .trim_start_matches('/')\n .chars()\n .collect::>();\n assert_eq!(chars.len(), 3);\n assert!(chars[0].is_ascii_lowercase());\n assert!(chars[1].is_ascii_digit());\n assert!(chars[2].is_ascii_lowercase());\n\n let req = get_req(\"/[a-z][0-9][a-z]\", &[\"--http2\", \"--rand-regex-url\"]).await;\n let chars = req\n .uri()\n .to_string()\n .split('/')\n .next_back()\n .unwrap()\n .chars()\n .collect::>();\n assert_eq!(chars.len(), 3);\n assert!(chars[0].is_ascii_lowercase());\n assert!(chars[1].is_ascii_digit());\n assert!(chars[2].is_ascii_lowercase());\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_redirect() {\n for n in 1..=5 {\n assert!(redirect(n, true, 10).await);\n assert!(redirect(n, false, 10).await);\n }\n for n in 11..=15 {\n assert!(!redirect(n, true, 10).await);\n assert!(!redirect(n, false, 10).await);\n }\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to() {\n assert_eq!(\n get_host_with_connect_to(\"invalid.example.org\").await,\n \"invalid.example.org\"\n )\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_randomness() {\n let (count1, count2) = distribution_on_two_matching_connect_to(\"invalid.example.org\").await;\n assert!(count1 + count2 == 100);\n assert!(count1 >= 10 && count2 >= 10); // should not be too flaky with 100 coin tosses\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_ipv6_target() {\n assert_eq!(\n get_host_with_connect_to_ipv6_target(\"invalid.example.org\").await,\n \"invalid.example.org\"\n )\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_ipv6_requested() {\n assert_eq!(get_host_with_connect_to_ipv6_requested().await, \"[::1]\")\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_redirect() {\n assert_eq!(\n get_host_with_connect_to_redirect(\"invalid.example.org\").await,\n \"invalid.example.org\"\n )\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_http_proxy_override() {\n let (tx, rx) = kanal::unbounded();\n let proxy_port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n\n let listener = tokio::net::TcpListener::bind((\"127.0.0.1\", proxy_port))\n .await\n .unwrap();\n\n tokio::spawn(async move {\n let (stream, _) = listener.accept().await.unwrap();\n let tx = tx.clone();\n\n hyper::server::conn::http1::Builder::new()\n .preserve_header_case(true)\n .title_case_headers(true)\n .serve_connection(\n TokioIo::new(stream),\n service_fn(move |req: Request| {\n let tx = tx.clone();\n async move {\n let authority = req\n .uri()\n .authority()\n .map(|a| a.to_string())\n .expect(\"proxy received origin-form request\");\n let host = req\n .headers()\n .get(\"host\")\n .and_then(|v| v.to_str().ok())\n .map(|s| s.to_string())\n .unwrap_or_default();\n tx.send((authority, host)).unwrap();\n Ok::<_, Infallible>(Response::new(\"proxy\".to_string()))\n }\n }),\n )\n .await\n .unwrap();\n });\n\n let override_port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n \"-x\".to_string(),\n format!(\"http://127.0.0.1:{proxy_port}\"),\n \"--connect-to\".to_string(),\n format!(\"example.test:80:127.0.0.1:{override_port}\"),\n \"http://example.test/\".to_string(),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n let (authority, host) = rx.try_recv().unwrap().unwrap();\n assert_eq!(authority, format!(\"127.0.0.1:{override_port}\"));\n assert_eq!(host, \"example.test\");\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_connect_to_https_proxy_connect_override() {\n let (connect_tx, connect_rx) = kanal::unbounded();\n let (host_tx, host_rx) = kanal::unbounded();\n\n let service = service_fn(move |req: Request| {\n let host_tx = host_tx.clone();\n async move {\n let host = req\n .headers()\n .get(\"host\")\n .and_then(|h| h.to_str().ok())\n .map(|s| s.to_string())\n .unwrap_or_default();\n host_tx.send(host).unwrap();\n Ok::<_, Infallible>(Response::new(\"Hello World\".to_string()))\n }\n });\n\n let (proxy_port, proxy_serve) =\n bind_proxy_with_recorder(service, false, connect_tx.clone()).await;\n\n tokio::spawn(proxy_serve);\n\n let override_port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n let args = vec![\n \"-n\".to_string(),\n \"1\".to_string(),\n \"--insecure\".to_string(),\n \"-x\".to_string(),\n format!(\"http://127.0.0.1:{proxy_port}\"),\n \"--proxy-header\".to_string(),\n \"proxy-authorization: test\".to_string(),\n \"--connect-to\".to_string(),\n format!(\"example.test:443:127.0.0.1:{override_port}\"),\n \"https://example.test/\".to_string(),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n let connect_target = connect_rx.try_recv().unwrap().unwrap();\n assert_eq!(connect_target, format!(\"127.0.0.1:{override_port}\"));\n let host_header = host_rx.try_recv().unwrap().unwrap();\n assert_eq!(host_header, \"example.test\");\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_ipv6() {\n let (tx, rx) = kanal::unbounded();\n\n let app = Router::new().route(\n \"/\",\n get(|| async move {\n tx.send(()).unwrap();\n \"Hello World\"\n }),\n );\n\n let port = next_port();\n let listener = bind_port_ipv6(port).await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n format!(\"http://[::1]:{port}/\"),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap();\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_query_limit() {\n // burst 10 requests with delay of 2s and rate of 4\n let mut args = vec![\"-n\", \"10\", \"--burst-delay\", \"2s\", \"--burst-rate\", \"4\"];\n assert_eq!(test_request_count(args.as_slice()).await, 10);\n args.push(\"--http2\");\n assert_eq!(test_request_count(args.as_slice()).await, 10);\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_query_limit_with_time_limit() {\n // 1.75 qps for 2sec = expect 4 requests at times 0, 0.571, 1.142, 1,714sec\n assert_eq!(test_request_count(&[\"-z\", \"2s\", \"-q\", \"1.75\"]).await, 4);\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_http_versions() {\n assert_eq!(get_req(\"/\", &[]).await.version(), http::Version::HTTP_11);\n assert_eq!(\n get_req(\"/\", &[\"--http2\"]).await.version(),\n http::Version::HTTP_2\n );\n assert_eq!(\n get_req(\"/\", &[\"--http-version\", \"2\"]).await.version(),\n http::Version::HTTP_2\n );\n #[cfg(feature = \"http3\")]\n assert_eq!(\n get_req(\"/\", &[\"--http-version\", \"3\"]).await.version(),\n http::Version::HTTP_3\n );\n}\n\n#[cfg(unix)]\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_unix_socket() {\n let (tx, rx) = kanal::unbounded();\n\n let tmp = tempfile::tempdir().unwrap();\n let path = tmp.path().join(\"socket\");\n\n let listener = std::os::unix::net::UnixListener::bind(&path).unwrap();\n tokio::spawn(async move {\n actix_web::HttpServer::new(move || {\n let tx = actix_web::web::Data::new(tx.clone());\n actix_web::App::new().service(actix_web::web::resource(\"/\").to(move || {\n let tx = tx.clone();\n async move {\n tx.send(()).unwrap();\n \"Hello World\"\n }\n }))\n })\n .listen_uds(listener)\n .unwrap()\n .run()\n .await\n .unwrap();\n });\n\n let args = [\n \"-n\".to_string(),\n \"1\".to_string(),\n \"--unix-socket\".to_string(),\n path.to_str().unwrap().to_string(),\n \"http://unix-socket.invalid-tld/\".to_string(),\n ];\n run(args.iter().map(|s| s.as_str())).await;\n\n rx.try_recv().unwrap().unwrap();\n}\n\nfn make_root_issuer() -> rcgen::Issuer<'static, rcgen::KeyPair> {\n let mut params = rcgen::CertificateParams::default();\n\n params.distinguished_name = rcgen::DistinguishedName::new();\n params.distinguished_name.push(\n rcgen::DnType::CommonName,\n rcgen::DnValue::Utf8String(\"\".to_string()),\n );\n params.key_usages = vec![\n rcgen::KeyUsagePurpose::KeyCertSign,\n rcgen::KeyUsagePurpose::CrlSign,\n ];\n params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);\n\n let signing_key = rcgen::KeyPair::generate().unwrap();\n\n rcgen::Issuer::new(params, signing_key)\n}\n\nasync fn bind_proxy(service: S, http2: bool) -> (u16, impl Future)\nwhere\n S: HttpService + Clone + Send + 'static,\n S::Error: Into>,\n S::ResBody: Send + Sync + 'static,\n ::Data: Send,\n ::Error: Into>,\n S::Future: Send,\n{\n let port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n let tcp_listener = tokio::net::TcpListener::bind((\"127.0.0.1\", port))\n .await\n .unwrap();\n\n let issuer = make_root_issuer();\n let proxy = Arc::new(http_mitm_proxy::MitmProxy::new(Some(issuer), None));\n\n let serve = async move {\n let (stream, _) = tcp_listener.accept().await.unwrap();\n\n let proxy = proxy.clone();\n let service = service.clone();\n\n let outer = service_fn(move |req| {\n // Test --proxy-header option\n assert_eq!(\n req.headers()\n .get(\"proxy-authorization\")\n .unwrap()\n .to_str()\n .unwrap(),\n \"test\"\n );\n\n MitmProxy::wrap_service(proxy.clone(), service.clone()).call(req)\n });\n\n tokio::spawn(async move {\n if http2 {\n let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new())\n .serve_connection(TokioIo::new(stream), outer)\n .await;\n } else {\n let _ = hyper::server::conn::http1::Builder::new()\n .preserve_header_case(true)\n .title_case_headers(true)\n .serve_connection(TokioIo::new(stream), outer)\n .with_upgrades()\n .await;\n }\n });\n };\n\n (port, serve)\n}\n\nasync fn bind_proxy_with_recorder(\n service: S,\n http2: bool,\n recorder: kanal::Sender,\n) -> (u16, impl Future)\nwhere\n S: HttpService + Clone + Send + 'static,\n S::Error: Into>,\n S::ResBody: Send + Sync + 'static,\n ::Data: Send,\n ::Error: Into>,\n S::Future: Send,\n{\n let port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n let tcp_listener = tokio::net::TcpListener::bind((\"127.0.0.1\", port))\n .await\n .unwrap();\n\n let issuer = make_root_issuer();\n let proxy = Arc::new(http_mitm_proxy::MitmProxy::new(Some(issuer), None));\n\n let serve = async move {\n let (stream, _) = tcp_listener.accept().await.unwrap();\n\n let proxy = proxy.clone();\n let service = service.clone();\n let recorder = recorder.clone();\n\n let outer = service_fn(move |req| {\n let recorder = recorder.clone();\n if req.method() == hyper::Method::CONNECT {\n recorder.send(req.uri().to_string()).unwrap();\n }\n\n assert_eq!(\n req.headers()\n .get(\"proxy-authorization\")\n .unwrap()\n .to_str()\n .unwrap(),\n \"test\"\n );\n\n MitmProxy::wrap_service(proxy.clone(), service.clone()).call(req)\n });\n\n tokio::spawn(async move {\n if http2 {\n let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new())\n .serve_connection(TokioIo::new(stream), outer)\n .await;\n } else {\n let _ = hyper::server::conn::http1::Builder::new()\n .preserve_header_case(true)\n .title_case_headers(true)\n .serve_connection(TokioIo::new(stream), outer)\n .with_upgrades()\n .await;\n }\n });\n };\n\n (port, serve)\n}\n\nasync fn test_proxy_with_setting(https: bool, http2: bool, proxy_http2: bool) {\n let (proxy_port, proxy_serve) = bind_proxy(\n service_fn(|_req| async {\n let res = Response::new(\"Hello World\".to_string());\n Ok::<_, Infallible>(res)\n }),\n proxy_http2,\n )\n .await;\n\n tokio::spawn(proxy_serve);\n\n let mut args = Vec::new();\n\n let scheme = if https { \"https\" } else { \"http\" };\n args.extend(\n [\n \"--no-tui\",\n \"-n\",\n \"1\",\n \"--no-tui\",\n \"--output-format\",\n \"quiet\",\n \"--insecure\",\n \"-x\",\n ]\n .into_iter()\n .map(|s| s.to_string()),\n );\n args.push(format!(\"http://127.0.0.1:{proxy_port}/\"));\n args.extend(\n [\"--proxy-header\", \"proxy-authorization: test\"]\n .into_iter()\n .map(|s| s.to_string()),\n );\n args.push(format!(\"{scheme}://example.com/\"));\n if http2 {\n args.push(\"--http2\".to_string());\n }\n if proxy_http2 {\n args.push(\"--proxy-http2\".to_string());\n }\n\n use clap::Parser;\n let opts = oha::Opts::try_parse_from(args).unwrap();\n oha::run(opts).await.unwrap();\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_proxy() {\n for https in [false, true] {\n for http2 in [false, true] {\n for proxy_http2 in [false, true] {\n test_proxy_with_setting(https, http2, proxy_http2).await;\n }\n }\n }\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_google() {\n let temp_path = tempfile::NamedTempFile::new().unwrap().into_temp_path();\n let args = vec![\n \"oha\".to_string(),\n \"--no-tui\".to_string(),\n \"-n\".to_string(),\n \"1\".to_string(),\n \"https://www.google.com/\".to_string(),\n \"--output\".to_string(),\n temp_path.to_str().unwrap().to_string(),\n ];\n let opts = oha::Opts::try_parse_from(args).unwrap();\n oha::run(opts).await.unwrap();\n\n let output = std::fs::read_to_string(&temp_path).unwrap();\n assert!(output.contains(\"[200] 1 responses\\n\"));\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_json_schema() {\n let app = Router::new().route(\"/\", get(|| async move { \"Hello World\" }));\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n const SCHEMA: &str = include_str!(\"../schema.json\");\n let schema_value: serde_json::Value = serde_json::from_str(SCHEMA).unwrap();\n let validator = jsonschema::validator_for(&schema_value).unwrap();\n\n let temp_path = tempfile::NamedTempFile::new().unwrap().into_temp_path();\n let args = vec![\n \"oha\".to_string(),\n \"--no-tui\".to_string(),\n \"-n\".to_string(),\n \"10\".to_string(),\n \"--output-format\".to_string(),\n \"json\".to_string(),\n format!(\"http://127.0.0.1:{port}/\"),\n \"--output\".to_string(),\n temp_path.to_str().unwrap().to_string(),\n ];\n let opts = oha::Opts::try_parse_from(args).unwrap();\n oha::run(opts).await.unwrap();\n\n let output_json = std::fs::read_to_string(&temp_path).unwrap();\n\n let temp_path = tempfile::NamedTempFile::new().unwrap().into_temp_path();\n let args = vec![\n \"oha\".to_string(),\n \"--no-tui\".to_string(),\n \"-n\".to_string(),\n \"10\".to_string(),\n \"--output-format\".to_string(),\n \"json\".to_string(),\n \"--stats-success-breakdown\".to_string(),\n format!(\"http://127.0.0.1:{port}/\"),\n \"--output\".to_string(),\n temp_path.to_str().unwrap().to_string(),\n ];\n let opts = oha::Opts::try_parse_from(args).unwrap();\n oha::run(opts).await.unwrap();\n let output_json_stats_success_breakdown = std::fs::read_to_string(&temp_path).unwrap();\n\n let value: serde_json::Value = serde_json::from_str(&output_json).unwrap();\n let value_stats_success_breakdown: serde_json::Value =\n serde_json::from_str(&output_json_stats_success_breakdown).unwrap();\n\n if validator.validate(&value).is_err() {\n for error in validator.iter_errors(&value) {\n eprintln!(\"{error}\");\n }\n panic!(\"JSON schema validation failed\\n{output_json}\");\n }\n\n if validator.validate(&value_stats_success_breakdown).is_err() {\n for error in validator.iter_errors(&value_stats_success_breakdown) {\n eprintln!(\"{error}\");\n }\n panic!(\"JSON schema validation failed\\n{output_json_stats_success_breakdown}\");\n }\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_csv_output() {\n let app = Router::new().route(\"/\", get(|| async move { \"Hello World\" }));\n\n let (listener, port) = bind_port_and_increment().await;\n tokio::spawn(async { axum::serve(listener, app).await });\n\n let temp_path = tempfile::NamedTempFile::new().unwrap().into_temp_path();\n let args = vec![\n \"oha\".to_string(),\n \"--no-tui\".to_string(),\n \"-n\".to_string(),\n \"5\".to_string(),\n \"--output-format\".to_string(),\n \"csv\".to_string(),\n format!(\"http://127.0.0.1:{port}/\"),\n \"--output\".to_string(),\n temp_path.to_str().unwrap().to_string(),\n ];\n let opts = oha::Opts::try_parse_from(args).unwrap();\n oha::run(opts).await.unwrap();\n let output_csv = std::fs::read_to_string(&temp_path).unwrap();\n\n // Validate that we get CSV output in following format,\n // header and one row for each request:\n // request-start,DNS,DNS+dialup,Response-delay,request-duration,bytes,status\n // 0.002211678,0.000374078,0.001148565,0.002619327,0.002626127,11,200\n // ...\n\n let lines: Vec<&str> = output_csv.lines().collect();\n assert_eq!(lines.len(), 6);\n assert_eq!(\n lines[0],\n \"request-start,DNS,DNS+dialup,Response-delay,request-duration,bytes,status\"\n );\n let mut latest_start = 0f64;\n for line in lines.iter().skip(1) {\n let parts: Vec<&str> = line.split(\",\").collect();\n assert_eq!(parts.len(), 7);\n // validate that the requests are in ascending time order\n let current_start = f64::from_str(parts[0]).unwrap();\n assert!(current_start >= latest_start);\n latest_start = current_start;\n assert!(f64::from_str(parts[1]).unwrap() >= 0f64);\n assert!(f64::from_str(parts[2]).unwrap() >= 0f64);\n assert!(f64::from_str(parts[3]).unwrap() > 0f64);\n assert!(f64::from_str(parts[4]).unwrap() > 0f64);\n assert_eq!(usize::from_str(parts[5]).unwrap(), 11);\n assert_eq!(u16::from_str(parts[6]).unwrap(), 200);\n }\n}\n\nfn setup_mtls_server(\n dir: std::path::PathBuf,\n) -> (u16, impl Future>) {\n let port = PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);\n let addr = SocketAddr::from(([127, 0, 0, 1], port));\n\n // build our application with a route\n let app = Router::new()\n // `GET /` goes to `root`\n .route(\"/\", get(|| async { \"Hello, World\" }));\n\n let make_cert = || {\n // Workaround for mac & native-tls\n // https://github.com/sfackler/rust-native-tls/issues/225\n let key_pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_RSA_SHA256).unwrap();\n let params = rcgen::CertificateParams::new(vec![\"localhost\".to_string()]).unwrap();\n\n let cert = params.self_signed(&key_pair).unwrap();\n (cert, key_pair)\n };\n\n let server_cert = make_cert();\n let client_cert = make_cert();\n\n let mut roots = rustls::RootCertStore::empty();\n roots.add(client_cert.0.der().clone()).unwrap();\n let _ = rustls::crypto::CryptoProvider::install_default(\n rustls::crypto::aws_lc_rs::default_provider(),\n );\n let verifier = rustls::server::WebPkiClientVerifier::builder(Arc::new(roots))\n .build()\n .unwrap();\n\n let config = rustls::ServerConfig::builder()\n .with_client_cert_verifier(verifier)\n .with_single_cert(\n vec![server_cert.0.der().clone()],\n rustls::pki_types::PrivateKeyDer::Pkcs8(rustls::pki_types::PrivatePkcs8KeyDer::from(\n server_cert.1.serialize_der(),\n )),\n )\n .unwrap();\n\n let config = axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(config));\n\n File::create(dir.join(\"server.crt\"))\n .unwrap()\n .write_all(server_cert.0.pem().as_bytes())\n .unwrap();\n\n File::create(dir.join(\"client.crt\"))\n .unwrap()\n .write_all(client_cert.0.pem().as_bytes())\n .unwrap();\n\n File::create(dir.join(\"client.key\"))\n .unwrap()\n .write_all(client_cert.1.serialize_pem().as_bytes())\n .unwrap();\n\n (\n port,\n axum_server::bind_rustls(addr, config).serve(app.into_make_service()),\n )\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_mtls() {\n let dir = tempfile::tempdir().unwrap();\n let (port, server) = setup_mtls_server(dir.path().to_path_buf());\n\n tokio::spawn(server);\n\n let args = vec![\n \"-n\".to_string(),\n \"1\".to_string(),\n \"--cacert\".to_string(),\n dir.path().join(\"server.crt\").to_string_lossy().to_string(),\n \"--cert\".to_string(),\n dir.path().join(\"client.crt\").to_string_lossy().to_string(),\n \"--key\".to_string(),\n dir.path().join(\"client.key\").to_string_lossy().to_string(),\n format!(\"https://localhost:{port}/\"),\n ];\n\n run(args.iter().map(|s| s.as_str())).await;\n}\n\n#[tokio::test(flavor = \"multi_thread\")]\nasync fn test_body_path_lines() {\n let body = \"0\\n1\\n2\";\n let mut tmp = tempfile::NamedTempFile::new().unwrap();\n tmp.write_all(body.as_bytes()).unwrap();\n\n let tmp_path = tmp.path().to_str().unwrap();\n\n let mut counts = [0; 3];\n for _ in 0..32 {\n let req = get_req(\"/\", [\"-Z\", tmp_path, \"-m\", \"POST\"].as_slice()).await;\n\n let req_body = req.into_body();\n let line = std::str::from_utf8(&req_body).unwrap();\n counts[line.parse::().unwrap()] += 1;\n }\n\n // test failure rate should be very low\n assert!(counts.iter().all(|&c| c > 0));\n}\n" } ]