use anyhow::Context;
use aws_auth::AwsSignatureConfig;
use bytes::Bytes;
use clap::Parser;
use crossterm::tty::IsTty;
use hickory_resolver::config::{ResolverConfig, ResolverOpts};
use humantime::Duration;
use hyper::{
HeaderMap,
http::{self, header::HeaderName, header::HeaderValue},
};
use printer::{PrintConfig, PrintMode};
use rand_regex::Regex;
use ratatui::crossterm;
use result_data::ResultData;
use std::{
env,
fs::File,
io::{BufRead, BufReader, Read},
path::{Path, PathBuf},
pin::Pin,
sync::Arc,
};
use timescale::TimeScale;
use url::Url;
use url_generator::UrlGenerator;
mod aws_auth;
mod cli;
mod client;
#[cfg(feature = "http3")]
mod client_h3;
mod curl_compat;
mod db;
mod histogram;
mod monitor;
mod pcg64si;
mod printer;
mod request_generator;
mod result_data;
mod timescale;
mod tls_config;
mod url_generator;
#[cfg(not(target_env = "msvc"))]
use tikv_jemallocator::Jemalloc;
use crate::{
cli::{ConnectToEntry, parse_header},
request_generator::{BodyGenerator, Proxy, RequestGenerator},
};
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: Jemalloc = Jemalloc;
#[derive(Parser)]
#[command(version, about, long_about = None)]
#[command(arg_required_else_help(true))]
#[command(styles = clap_cargo::style::CLAP_STYLING)]
pub struct Opts {
#[arg(help = "Target URL or file with multiple URLs.")]
url: String,
#[arg(
help = "Number of requests to run. Accepts plain numbers or suffixes: k = 1,000, m = 1,000,000 (e.g. 10k, 1m).",
short = 'n',
default_value = "200",
conflicts_with = "duration",
value_parser = cli::parse_n_requests
)]
n_requests: usize,
#[arg(
help = "Number of connections to run concurrently. You may should increase limit to number of open files for larger `-c`.",
short = 'c',
default_value = "50"
)]
n_connections: usize,
#[arg(
help = "Number of parallel requests to send on HTTP/2. `oha` will run c * p concurrent workers in total.",
short = 'p',
default_value = "1"
)]
n_http2_parallel: usize,
#[arg(
help = "Duration of application to send requests.
On HTTP/1, When the duration is reached, ongoing requests are aborted and counted as \"aborted due to deadline\"
You can change this behavior with `-w` option.
Currently, on HTTP/2, When the duration is reached, ongoing requests are waited. `-w` option is ignored.
Examples: -z 10s -z 3m.",
short = 'z',
conflicts_with = "n_requests"
)]
duration: Option<Duration>,
#[arg(
help = "When the duration is reached, ongoing requests are waited",
short,
long,
default_value = "false",
requires = "duration"
)]
wait_ongoing_requests_after_deadline: bool,
#[arg(help = "Rate limit for all, in queries per second (QPS)", short = 'q', conflicts_with_all = ["burst_duration", "burst_requests"])]
query_per_second: Option<f64>,
#[arg(
help = "Introduce delay between a predefined number of requests.
Note: If qps is specified, burst will be ignored",
long = "burst-delay",
requires = "burst_requests",
conflicts_with = "query_per_second"
)]
burst_duration: Option<Duration>,
#[arg(
help = "Rates of requests for burst. Default is 1
Note: If qps is specified, burst will be ignored",
long = "burst-rate",
requires = "burst_duration",
conflicts_with = "query_per_second"
)]
burst_requests: Option<usize>,
#[arg(
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.",
default_value = "false",
long
)]
rand_regex_url: bool,
#[arg(
help = "Read the URLs to query from a file",
default_value = "false",
long
)]
urls_from_file: bool,
#[arg(
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.",
default_value = "4",
long,
requires = "rand_regex_url"
)]
max_repeat: u32,
#[arg(
help = "Dump target Urls <DUMP_URLS> times to debug --rand-regex-url",
long
)]
dump_urls: Option<usize>,
#[arg(
help = "Correct latency to avoid coordinated omission problem. It's ignored if -q is not set.",
long = "latency-correction"
)]
latency_correction: bool,
#[arg(help = "No realtime tui", long = "no-tui")]
no_tui: bool,
#[arg(help = "Frame per second for tui.", default_value = "16", long = "fps")]
fps: usize,
#[arg(
help = "HTTP method",
short = 'm',
long = "method",
default_value = "GET"
)]
method: http::Method,
#[arg(help = "Custom HTTP header. Examples: -H \"foo: bar\"", short = 'H', value_parser = parse_header)]
headers: Vec<(HeaderName, HeaderValue)>,
#[arg(
help = "Custom Proxy HTTP header. Examples: --proxy-header \"foo: bar\"",
long = "proxy-header",
value_parser = parse_header
)]
proxy_headers: Vec<(HeaderName, HeaderValue)>,
#[arg(help = "Timeout for each request. Default to infinite.", short = 't')]
timeout: Option<humantime::Duration>,
#[arg(
help = "Timeout for establishing a new connection. Default to 5s.",
long = "connect-timeout",
default_value = "5s"
)]
connect_timeout: humantime::Duration,
#[arg(help = "HTTP Accept Header.", short = 'A')]
accept_header: Option<String>,
#[arg(help = "HTTP request body.", short = 'd', conflicts_with_all = ["body_path", "body_path_lines", "form"])]
body_string: Option<String>,
#[arg(help = "HTTP request body from file.", short = 'D', conflicts_with_all = ["body_string", "body_path_lines", "form"])]
body_path: Option<std::path::PathBuf>,
#[arg(help = "HTTP request body from file line by line.", short = 'Z', conflicts_with_all = ["body_string", "body_path", "form"])]
body_path_lines: Option<std::path::PathBuf>,
#[arg(
help = "Specify HTTP multipart POST data (curl compatible). Examples: -F 'name=value' -F 'file=@path/to/file'",
short = 'F',
long = "form",
conflicts_with_all = ["body_string", "body_path", "body_path_lines"]
)]
form: Vec<String>,
#[arg(help = "Content-Type.", short = 'T')]
content_type: Option<String>,
#[arg(
help = "Basic authentication (username:password), or AWS credentials (access_key:secret_key)",
short = 'a'
)]
basic_auth: Option<String>,
#[arg(help = "AWS session token", long = "aws-session")]
aws_session: Option<String>,
#[arg(
help = "AWS SigV4 signing params (format: aws:amz:region:service)",
long = "aws-sigv4"
)]
aws_sigv4: Option<String>,
#[arg(help = "HTTP proxy", short = 'x')]
proxy: Option<Url>,
#[arg(
help = "HTTP version to connect to proxy. Available values 0.9, 1.0, 1.1, 2.",
long = "proxy-http-version"
)]
proxy_http_version: Option<String>,
#[arg(
help = "Use HTTP/2 to connect to proxy. Shorthand for --proxy-http-version=2",
long = "proxy-http2"
)]
proxy_http2: bool,
#[arg(
help = "HTTP version. Available values 0.9, 1.0, 1.1, 2, 3",
long = "http-version"
)]
http_version: Option<String>,
#[arg(help = "Use HTTP/2. Shorthand for --http-version=2", long = "http2")]
http2: bool,
#[arg(help = "HTTP Host header", long = "host")]
host: Option<String>,
#[arg(help = "Disable compression.", long = "disable-compression")]
disable_compression: bool,
#[arg(
help = "Limit for number of Redirect. Set 0 for no redirection. Redirection isn't supported for HTTP/2.",
default_value = "10",
short = 'r',
long = "redirect"
)]
redirect: usize,
#[arg(
help = "Disable keep-alive, prevents re-use of TCP connections between different HTTP requests. This isn't supported for HTTP/2.",
long = "disable-keepalive"
)]
disable_keepalive: bool,
#[arg(
help = "*Not* perform a DNS lookup at beginning to cache it",
long = "no-pre-lookup",
default_value = "false"
)]
no_pre_lookup: bool,
#[arg(help = "Lookup only ipv6.", long = "ipv6")]
ipv6: bool,
#[arg(help = "Lookup only ipv4.", long = "ipv4")]
ipv4: bool,
#[arg(
help = "(TLS) Use the specified certificate file to verify the peer. Native certificate store is used even if this argument is specified.",
long
)]
cacert: Option<PathBuf>,
#[arg(
help = "(TLS) Use the specified client certificate file. --key must be also specified",
long,
requires = "key"
)]
cert: Option<PathBuf>,
#[arg(
help = "(TLS) Use the specified client key file. --cert must be also specified",
long,
requires = "cert"
)]
key: Option<PathBuf>,
#[arg(help = "Accept invalid certs.", long = "insecure")]
insecure: bool,
#[arg(
help = "Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'
Note: if used several times for the same host:port:target_host:target_port, a random choice is made",
long = "connect-to"
)]
connect_to: Vec<ConnectToEntry>,
#[arg(
help = "Disable the color scheme.",
alias = "disable-color",
long = "no-color",
env = "NO_COLOR"
)]
no_color: bool,
#[cfg(unix)]
#[arg(
help = "Connect to a unix socket instead of the domain in the URL. Only for non-HTTPS URLs.",
long = "unix-socket",
group = "socket-type"
)]
unix_socket: Option<std::path::PathBuf>,
#[cfg(feature = "vsock")]
#[arg(
help = "Connect to a VSOCK socket using 'cid:port' instead of the domain in the URL. Only for non-HTTPS URLs.",
long = "vsock-addr",
value_parser = cli::parse_vsock_addr,
group = "socket-type"
)]
vsock_addr: Option<tokio_vsock::VsockAddr>,
#[arg(
help = "Include a response status code successful or not successful breakdown for the time histogram and distribution statistics",
long = "stats-success-breakdown"
)]
stats_success_breakdown: bool,
#[arg(
help = "Write succeeded requests to sqlite database url E.G test.db",
long = "db-url"
)]
db_url: Option<String>,
#[arg(
long,
help = "Perform a single request and dump the request and response"
)]
debug: bool,
#[arg(
help = "Output file to write the results to. If not specified, results are written to stdout.",
long,
short
)]
output: Option<PathBuf>,
#[arg(help = "Output format", long, default_value = "text")]
output_format: Option<PrintMode>,
#[arg(
help = "Time unit to be used. If not specified, the time unit is determined automatically. This option affects only text format.",
long,
short = 'u'
)]
time_unit: Option<TimeScale>,
}
pub async fn run(mut opts: Opts) -> anyhow::Result<()> {
let work_mode = opts.work_mode();
let aws_config = if let Some(signing_params) = opts.aws_sigv4 {
if let Some(auth) = &opts.basic_auth {
let parts: Vec<&str> = auth.split(':').collect();
if parts.len() != 2 {
anyhow::bail!("Invalid AWS credentials format. Expected access_key:secret_key");
}
let access_key = parts[0];
let secret_key = parts[1];
let session_token = opts.aws_session.take();
Some(AwsSignatureConfig::new(
access_key,
secret_key,
&signing_params,
session_token,
)?)
} else {
anyhow::bail!("AWS credentials (--auth) required when using --aws-sigv4");
}
} else {
None
};
let parse_http_version = |is_http2: bool, version: Option<&str>| match (is_http2, version) {
(true, Some(_)) => anyhow::bail!("--http2 and --http-version are exclusive"),
(true, None) => Ok(http::Version::HTTP_2),
(false, Some(http_version)) => match http_version.trim() {
"0.9" => Ok(http::Version::HTTP_09),
"1.0" => Ok(http::Version::HTTP_10),
"1.1" => Ok(http::Version::HTTP_11),
"2.0" | "2" => Ok(http::Version::HTTP_2),
#[cfg(feature = "http3")]
"3.0" | "3" => Ok(http::Version::HTTP_3),
#[cfg(not(feature = "http3"))]
"3.0" | "3" => anyhow::bail!(
"Your Oha instance has not been built with HTTP/3 support. Try recompiling with the feature enabled."
),
_ => anyhow::bail!("Unknown HTTP version. Valid versions are 0.9, 1.0, 1.1, 2, 3"),
},
(false, None) => Ok(http::Version::HTTP_11),
};
let http_version: http::Version = parse_http_version(opts.http2, opts.http_version.as_deref())?;
let proxy_http_version: http::Version =
parse_http_version(opts.proxy_http2, opts.proxy_http_version.as_deref())?;
let url_generator = if opts.rand_regex_url {
let dot_disabled: String = opts
.url
.chars()
.map(|c| {
if c == '.' {
regex_syntax::escape(".")
} else {
c.to_string()
}
})
.collect();
UrlGenerator::new_dynamic(Regex::compile(&dot_disabled, opts.max_repeat)?)
} else if opts.urls_from_file {
let path = Path::new(opts.url.as_str());
let file = File::open(path)?;
let reader = std::io::BufReader::new(file);
let urls: Vec<Url> = reader
.lines()
.map_while(Result::ok)
.filter(|line| !line.trim().is_empty())
.map(|url_str| Url::parse(&url_str))
.collect::<Result<Vec<_>, _>>()?;
UrlGenerator::new_multi_static(urls)
} else {
UrlGenerator::new_static(Url::parse(&opts.url)?)
};
if let Some(n) = opts.dump_urls {
let mut rng = rand::rng();
for _ in 0..n {
let url = url_generator.generate(&mut rng)?;
println!("{url}");
}
return Ok(());
}
let url = url_generator.generate(&mut rand::rng())?;
let has_form_data = !opts.form.is_empty();
let (body_generator, form_content_type): (BodyGenerator, Option<String>) = if has_form_data {
let mut form = curl_compat::Form::new();
for form_str in opts.form {
let part: curl_compat::FormPart = form_str
.parse()
.with_context(|| format!("Failed to parse form data: {form_str}"))?;
form.add_part(part);
}
let form_body = form.body();
let content_type = form.content_type();
(BodyGenerator::Static(form_body.into()), Some(content_type))
} else if let Some(body_string) = opts.body_string {
(BodyGenerator::Static(body_string.into()), None)
} else if let Some(body_path) = opts.body_path {
let mut buf = Vec::new();
std::fs::File::open(body_path)?.read_to_end(&mut buf)?;
(BodyGenerator::Static(buf.into()), None)
} else if let Some(body_path_lines) = opts.body_path_lines {
let lines = BufReader::new(std::fs::File::open(body_path_lines)?)
.lines()
.map_while(Result::ok)
.map(Bytes::from)
.collect::<Vec<_>>();
(BodyGenerator::Random(lines), None)
} else {
(BodyGenerator::Static(Bytes::new()), None)
};
let method = if has_form_data && opts.method == http::Method::GET {
http::Method::POST
} else {
opts.method
};
let headers = {
let mut headers: http::header::HeaderMap = Default::default();
headers.insert(
http::header::ACCEPT,
http::header::HeaderValue::from_static("*/*"),
);
if !opts.disable_compression {
headers.insert(
http::header::ACCEPT_ENCODING,
http::header::HeaderValue::from_static("gzip, compress, deflate, br"),
);
}
headers
.entry(http::header::USER_AGENT)
.or_insert(HeaderValue::from_static(concat!(
"oha/",
env!("CARGO_PKG_VERSION")
)));
if let Some(h) = opts.accept_header {
headers.insert(http::header::ACCEPT, HeaderValue::from_bytes(h.as_bytes())?);
}
if let Some(h) = opts.content_type.or(form_content_type) {
headers.insert(
http::header::CONTENT_TYPE,
HeaderValue::from_bytes(h.as_bytes())?,
);
}
if let Some(h) = opts.host {
headers.insert(http::header::HOST, HeaderValue::from_bytes(h.as_bytes())?);
}
if let Some(auth) = opts.basic_auth {
let u_p = auth.splitn(2, ':').collect::<Vec<_>>();
anyhow::ensure!(u_p.len() == 2, anyhow::anyhow!("Parse auth"));
let mut header_value = b"Basic ".to_vec();
{
use std::io::Write;
let username = u_p[0];
let password = if u_p[1].is_empty() {
None
} else {
Some(u_p[1])
};
let mut encoder = base64::write::EncoderWriter::new(
&mut header_value,
&base64::engine::general_purpose::STANDARD,
);
write!(encoder, "{username}:").unwrap();
if let Some(password) = password {
write!(encoder, "{password}").unwrap();
}
}
headers.insert(
http::header::AUTHORIZATION,
HeaderValue::from_bytes(&header_value)?,
);
}
if opts.disable_keepalive && http_version == http::Version::HTTP_11 {
headers.insert(http::header::CONNECTION, HeaderValue::from_static("close"));
}
for (k, v) in opts.headers.into_iter() {
headers.insert(k, v);
}
headers
};
let proxy_headers = opts.proxy_headers.into_iter().collect::<HeaderMap<_>>();
let ip_strategy = match (opts.ipv4, opts.ipv6) {
(false, false) => {
if cfg!(target_os = "macos") && (url.host_str() == Some("localhost")) {
hickory_resolver::config::LookupIpStrategy::Ipv6thenIpv4
} else {
Default::default()
}
}
(true, false) => hickory_resolver::config::LookupIpStrategy::Ipv4Only,
(false, true) => hickory_resolver::config::LookupIpStrategy::Ipv6Only,
(true, true) => hickory_resolver::config::LookupIpStrategy::Ipv4AndIpv6,
};
let (config, mut resolver_opts) = system_resolv_conf()?;
resolver_opts.ip_strategy = ip_strategy;
let resolver = hickory_resolver::Resolver::builder_with_config(
config,
hickory_resolver::name_server::TokioConnectionProvider::default(),
)
.with_options(resolver_opts)
.build();
let cacert = opts.cacert.as_deref().map(std::fs::read).transpose()?;
let client_auth = match (opts.cert, opts.key) {
(Some(cert), Some(key)) => Some((std::fs::read(cert)?, std::fs::read(key)?)),
(None, None) => None,
_ => anyhow::bail!("Both --cert and --key must be specified"),
};
let url = url.into_owned();
let client = Arc::new(client::Client {
request_generator: RequestGenerator {
url_generator,
https: url.scheme() == "https",
version: http_version,
aws_config,
method,
headers,
body_generator,
http_proxy: if opts.proxy.is_some() && url.scheme() == "http" {
Some(Proxy {
headers: proxy_headers.clone(),
version: proxy_http_version,
})
} else {
None
},
},
proxy_http_version,
proxy_headers,
dns: client::Dns {
resolver,
connect_to: opts.connect_to,
},
timeout: opts.timeout.map(|d| d.into()),
connect_timeout: opts.connect_timeout.into(),
redirect_limit: opts.redirect,
disable_keepalive: opts.disable_keepalive,
proxy_url: opts.proxy,
#[cfg(unix)]
unix_socket: opts.unix_socket,
#[cfg(feature = "vsock")]
vsock_addr: opts.vsock_addr,
#[cfg(feature = "rustls")]
rustls_configs: tls_config::RuslsConfigs::new(
opts.insecure,
cacert.as_deref(),
client_auth
.as_ref()
.map(|(cert, key)| (cert.as_slice(), key.as_slice())),
),
#[cfg(all(feature = "native-tls", not(feature = "rustls")))]
native_tls_connectors: tls_config::NativeTlsConnectors::new(
opts.insecure,
cacert.as_deref(),
client_auth
.as_ref()
.map(|(cert, key)| (cert.as_slice(), key.as_slice())),
),
});
if !opts.no_pre_lookup {
client.pre_lookup().await?;
}
let no_tui = opts.no_tui || !std::io::stdout().is_tty() || opts.debug;
let print_config = {
let mode = opts.output_format.unwrap_or_default();
let disable_style = opts.no_color || !std::io::stdout().is_tty() || opts.output.is_some();
let output: Box<dyn std::io::Write + Send + 'static> = if let Some(output) = opts.output {
Box::new(File::create(output)?)
} else {
Box::new(std::io::stdout())
};
PrintConfig {
mode,
output,
disable_style,
stats_success_breakdown: opts.stats_success_breakdown,
time_unit: opts.time_unit,
}
};
let run = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)?
.as_secs();
let start = std::time::Instant::now();
let data_collect_future: Pin<Box<dyn std::future::Future<Output = (ResultData, PrintConfig)>>> =
match work_mode {
WorkMode::Debug => {
let mut print_config = print_config;
client::work_debug(&mut print_config.output, client).await?;
return Ok(());
}
WorkMode::FixedNumber {
n_requests,
n_connections,
n_http2_parallel,
query_limit: None,
latency_correction: _,
} if no_tui => {
let (result_tx, result_rx) = kanal::unbounded();
client::fast::work(
client.clone(),
result_tx,
n_requests,
n_connections,
n_http2_parallel,
)
.await;
Box::pin(async move {
let mut res = ResultData::default();
for r in result_rx {
res.merge(r);
}
(res, print_config)
})
}
WorkMode::Until {
duration,
n_connections,
n_http2_parallel,
query_limit: None,
latency_correction: _,
wait_ongoing_requests_after_deadline,
} if no_tui => {
let (result_tx, result_rx) = kanal::unbounded();
client::fast::work_until(
client.clone(),
result_tx,
start + duration,
n_connections,
n_http2_parallel,
wait_ongoing_requests_after_deadline,
)
.await;
Box::pin(async move {
let mut res = ResultData::default();
for r in result_rx {
res.merge(r);
}
(res, print_config)
})
}
mode => {
let (result_tx, result_rx) = kanal::unbounded();
let data_collector = if no_tui {
let token = tokio_util::sync::CancellationToken::new();
let result_rx_ctrl_c = result_rx.clone();
let token_ctrl_c = token.clone();
let ctrl_c = tokio::spawn(async move {
tokio::select! {
_ = tokio::signal::ctrl_c() => {
let mut all: ResultData = Default::default();
let mut buf = Vec::new();
let _ = result_rx_ctrl_c.drain_into(&mut buf);
for res in buf {
all.push(res);
}
let _ = printer::print_result(print_config, start, &all, start.elapsed());
std::process::exit(libc::EXIT_SUCCESS);
}
_ = token_ctrl_c.cancelled() => {
print_config
}
}
});
Box::pin(async move {
token.cancel();
let config = ctrl_c.await.unwrap();
let mut all = ResultData::default();
while let Ok(res) = result_rx.recv() {
all.push(res);
}
(all, config)
})
as Pin<Box<dyn std::future::Future<Output = (ResultData, PrintConfig)>>>
} else {
let join_handle = tokio::spawn(
monitor::Monitor {
print_config,
end_line: opts
.duration
.map(|d| monitor::EndLine::Duration(d.into()))
.unwrap_or(monitor::EndLine::NumQuery(opts.n_requests)),
report_receiver: result_rx,
start,
fps: opts.fps,
disable_color: opts.no_color,
time_unit: opts.time_unit,
}
.monitor(),
);
Box::pin(async { join_handle.await.unwrap().unwrap() })
as Pin<Box<dyn std::future::Future<Output = (ResultData, PrintConfig)>>>
};
match mode {
WorkMode::Debug => unreachable!("Must be already handled"),
WorkMode::FixedNumber {
n_requests,
n_connections,
n_http2_parallel,
query_limit,
latency_correction,
} => {
if let Some(query_limit) = query_limit {
if latency_correction {
client::work_with_qps_latency_correction(
client.clone(),
result_tx,
query_limit,
n_requests,
n_connections,
n_http2_parallel,
)
.await;
} else {
client::work_with_qps(
client.clone(),
result_tx,
query_limit,
n_requests,
n_connections,
n_http2_parallel,
)
.await;
}
} else {
client::work(
client.clone(),
result_tx,
n_requests,
n_connections,
n_http2_parallel,
)
.await;
}
}
WorkMode::Until {
duration,
n_connections,
n_http2_parallel,
query_limit,
latency_correction,
wait_ongoing_requests_after_deadline,
} => {
if let Some(query_limit) = query_limit {
if latency_correction {
client::work_until_with_qps_latency_correction(
client.clone(),
result_tx,
query_limit,
start,
start + duration,
n_connections,
n_http2_parallel,
wait_ongoing_requests_after_deadline,
)
.await;
} else {
client::work_until_with_qps(
client.clone(),
result_tx,
query_limit,
start,
start + duration,
n_connections,
n_http2_parallel,
wait_ongoing_requests_after_deadline,
)
.await;
}
} else {
client::work_until(
client.clone(),
result_tx,
start + duration,
n_connections,
n_http2_parallel,
wait_ongoing_requests_after_deadline,
)
.await;
}
}
}
data_collector
}
};
let duration = start.elapsed();
let (res, print_config) = data_collect_future.await;
printer::print_result(print_config, start, &res, duration)?;
if let Some(db_url) = opts.db_url {
eprintln!("Storing results to {db_url}");
db::store(&client, &db_url, start, res.success(), run as i64)?;
}
Ok(())
}
pub(crate) fn system_resolv_conf() -> anyhow::Result<(ResolverConfig, ResolverOpts)> {
#[cfg(unix)]
if env::var("TERMUX_VERSION").is_ok() {
let prefix = env::var("PREFIX")?;
let path = format!("{prefix}/etc/resolv.conf");
return match std::fs::read(&path) {
Ok(conf_data) => hickory_resolver::system_conf::parse_resolv_conf(conf_data)
.context(format!("DNS: failed to parse {path}")),
Err(err) => {
fallback_resolver_config(anyhow::anyhow!("DNS: failed to load {path}: {err}"))
}
};
}
match hickory_resolver::system_conf::read_system_conf() {
Ok(conf) => Ok(conf),
Err(err) => fallback_resolver_config(anyhow::anyhow!(
"DNS: failed to load /etc/resolv.conf: {err}"
)),
}
}
fn fallback_resolver_config(err: anyhow::Error) -> anyhow::Result<(ResolverConfig, ResolverOpts)> {
eprintln!("{err}");
let config = ResolverConfig::default();
let opts = ResolverOpts::default();
Ok((config, opts))
}
enum WorkMode {
Debug,
FixedNumber {
n_requests: usize,
n_connections: usize,
n_http2_parallel: usize,
query_limit: Option<client::QueryLimit>,
latency_correction: bool,
},
Until {
duration: std::time::Duration,
n_connections: usize,
n_http2_parallel: usize,
query_limit: Option<client::QueryLimit>,
latency_correction: bool,
wait_ongoing_requests_after_deadline: bool,
},
}
impl Opts {
fn work_mode(&self) -> WorkMode {
if self.debug {
WorkMode::Debug
} else if let Some(duration) = self.duration {
WorkMode::Until {
duration: duration.into(),
n_connections: self.n_connections,
n_http2_parallel: self.n_http2_parallel,
query_limit: match self.query_per_second {
Some(0f64) | None => self.burst_duration.map(|burst_duration| {
client::QueryLimit::Burst(
burst_duration.into(),
self.burst_requests.unwrap_or(1),
)
}),
Some(qps) => Some(client::QueryLimit::Qps(qps)),
},
latency_correction: self.latency_correction,
wait_ongoing_requests_after_deadline: self.wait_ongoing_requests_after_deadline,
}
} else {
WorkMode::FixedNumber {
n_requests: self.n_requests,
n_connections: self.n_connections,
n_http2_parallel: self.n_http2_parallel,
query_limit: match self.query_per_second {
Some(0f64) | None => self.burst_duration.map(|burst_duration| {
client::QueryLimit::Burst(
burst_duration.into(),
self.burst_requests.unwrap_or(1),
)
}),
Some(qps) => Some(client::QueryLimit::Qps(qps)),
},
latency_correction: self.latency_correction,
}
}
}
}