Skip to content

Commit

Permalink
Merge pull request #181 from 56quarters/zero
Browse files Browse the repository at this point in the history 
Disallow zero values for CLI arguments where they don't make sense
  • Loading branch information
@56quarters
56quarters authored Aug 18, 2024
2 parents 51d781a + bee4e2a commit 3cf2192
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 40 deletions.
11 changes: 6 additions & 5 deletions mtop/src/bin/dns.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use mtop_client::dns::{DnsClient, Flags, Message, MessageId, Name, Question, Rec
use std::fmt::Write;
use std::io::Cursor;
use std::net::SocketAddr;
use std::num::NonZeroU64;
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::atomic::AtomicBool;
Expand All @@ -16,7 +17,7 @@ use tracing::{Instrument, Level};

const DEFAULT_RECORD_TYPE: RecordType = RecordType::A;
const DEFAULT_RECORD_CLASS: RecordClass = RecordClass::INET;
const MIN_PING_INTERVAL_SECS: f64 = 0.01;
const MIN_PING_INTERVAL_SECS: f64 = 0.1;

/// dns: Make DNS queries or read/write binary format DNS messages
#[derive(Debug, Parser)]
Expand Down Expand Up @@ -68,7 +69,7 @@ struct PingCommand {
/// Timeout for DNS network operations in seconds, overriding whatever timeout is
/// configured in resolv.conf.
#[arg(long)]
nameserver_timeout_secs: Option<u64>,
nameserver_timeout_secs: Option<NonZeroU64>,

/// Type of record to request. Supported: A, AAAA, CNAME, NS, SOA, SRV, TXT.
#[arg(long, default_value_t = DEFAULT_RECORD_TYPE)]
Expand Down Expand Up @@ -107,7 +108,7 @@ struct QueryCommand {
/// Timeout for DNS network operations in seconds, overriding whatever timeout is
/// configured in resolv.conf.
#[arg(long)]
nameserver_timeout_secs: Option<u64>,
nameserver_timeout_secs: Option<NonZeroU64>,

/// Output query results in raw binary format instead of human-readable
/// text. NOTE, this may break your terminal and so should probably be piped
Expand Down Expand Up @@ -179,7 +180,7 @@ async fn run_ping(cmd: &PingCommand) -> ExitCode {
let client = mtop::dns::new_client(
&cmd.resolv_conf,
cmd.nameserver,
cmd.nameserver_timeout_secs.map(Duration::from_secs),
cmd.nameserver_timeout_secs.map(|v| Duration::from_secs(v.get())),
)
.instrument(tracing::span!(Level::INFO, "dns.new_client"))
.await;
Expand All @@ -198,7 +199,7 @@ async fn run_query(cmd: &QueryCommand) -> ExitCode {
let client = mtop::dns::new_client(
&cmd.resolv_conf,
cmd.nameserver,
cmd.nameserver_timeout_secs.map(Duration::from_secs),
cmd.nameserver_timeout_secs.map(|v| Duration::from_secs(v.get())),
)
.instrument(tracing::span!(Level::INFO, "dns.new_client"))
.await;
Expand Down
65 changes: 33 additions & 32 deletions mtop/src/bin/mc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use mtop_client::{
Timeout, TlsConfig, Value,
};
use rustls_pki_types::{InvalidDnsNameError, ServerName};
use std::num::{NonZeroU64, NonZeroUsize};
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::atomic::AtomicBool;
Expand Down Expand Up @@ -36,12 +37,12 @@ struct McConfig {
host: String,

/// Timeout for Memcached network operations, in seconds.
#[arg(long, env = "MC_TIMEOUT_SECS", default_value_t = 30)]
timeout_secs: u64,
#[arg(long, env = "MC_TIMEOUT_SECS", default_value_t = NonZeroU64::new(30).unwrap())]
timeout_secs: NonZeroU64,

/// Maximum number of idle connections to maintain per host.
#[arg(long, env = "MC_CONNECTIONS", default_value_t = 4)]
connections: u64,
#[arg(long, env = "MC_CONNECTIONS", default_value_t = NonZeroU64::new(4).unwrap())]
connections: NonZeroU64,

/// Output pprof protobuf profile data to this file if profiling support was enabled
/// at build time.
Expand Down Expand Up @@ -129,8 +130,8 @@ struct AddCommand {
#[derive(Debug, Args)]
struct BenchCommand {
/// How long to run the benchmark for in seconds.
#[arg(long, env = "MC_BENCH_TIME_SECS", default_value_t = 60)]
time_secs: u64,
#[arg(long, env = "MC_BENCH_TIME_SECS", default_value_t = NonZeroU64::new(60).unwrap())]
time_secs: NonZeroU64,

/// How many writes to the cache as a percentage of reads from the cache, 0 to 1.
///
Expand All @@ -143,17 +144,17 @@ struct BenchCommand {
/// How many workers to run at once, performing gets and sets against the cache.
///
/// Each worker does 10,000 gets and 500 sets per second in the default configuration.
#[arg(long, env = "MC_BENCH_CONCURRENCY", default_value_t = 1)]
concurrency: usize,
#[arg(long, env = "MC_BENCH_CONCURRENCY", default_value_t = NonZeroUsize::new(1).unwrap())]
concurrency: NonZeroUsize,

/// How long to wait between each batch of gets and sets performed against the cache.
///
/// Each batch is 1000 gets and 50 sets by default. With a delay of 100ms this means
/// there will be 10,000 gets and 500 sets per second. To increase the number of gets
/// and sets performed by a worker, reduce this number. To decrease the number of gets
/// and sets performed by a worker, increase this number.
#[arg(long, env = "MC_BENCH_DELAY_MILLIS", default_value_t = 100)]
delay_millis: u64,
#[arg(long, env = "MC_BENCH_DELAY_MILLIS", default_value_t = NonZeroU64::new(100).unwrap())]
delay_millis: NonZeroU64,

/// TTL to use for test values stored in the cache in seconds.
#[arg(long, env = "MC_BENCH_TTL_SECS", default_value_t = 300)]
Expand All @@ -169,12 +170,12 @@ struct BenchCommand {
#[derive(Debug, Args)]
struct CheckCommand {
/// How long to run the checks for in seconds.
#[arg(long, env = "MC_CHECK_TIME_SECS", default_value_t = 60)]
time_secs: u64,
#[arg(long, env = "MC_CHECK_TIME_SECS", default_value_t = NonZeroU64::new(60).unwrap())]
time_secs: NonZeroU64,

/// How long to wait between each health check in milliseconds.
#[arg(long, env = "MC_CHECK_DELAY_MILLIS", default_value_t = 100)]
delay_millis: u64,
#[arg(long, env = "MC_CHECK_DELAY_MILLIS", default_value_t = NonZeroU64::new(100).unwrap())]
delay_millis: NonZeroU64,
}

/// Decrement the value of an item in the cache.
Expand Down Expand Up @@ -290,7 +291,7 @@ async fn main() -> ExitCode {

let dns_client = mtop::dns::new_client(&opts.resolv_conf, None, None).await;
let discovery = Discovery::new(dns_client);
let timeout = Duration::from_secs(opts.timeout_secs);
let timeout = Duration::from_secs(opts.timeout_secs.get());
let servers = match discovery
.resolve_by_proto(&opts.host)
.timeout(timeout, "discovery.resolve_by_proto")
Expand Down Expand Up @@ -349,7 +350,7 @@ async fn new_client(opts: &McConfig, servers: &[Server]) -> Result<MemcachedClie
};

let cfg = MemcachedClientConfig {
pool_max_idle: opts.connections,
pool_max_idle: opts.connections.get(),
};

let selector = RendezvousSelector::new(servers.to_vec());
Expand Down Expand Up @@ -382,7 +383,7 @@ async fn run_add(opts: &McConfig, cmd: &AddCommand, client: &MemcachedClient) ->

if let Err(e) = client
.add(&cmd.key, 0, cmd.ttl, &buf)
.timeout(Duration::from_secs(opts.timeout_secs), "client.add")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.add")
.instrument(tracing::span!(Level::INFO, "client.add"))
.await
{
Expand All @@ -400,15 +401,15 @@ async fn run_bench(opts: &McConfig, cmd: &BenchCommand, client: MemcachedClient)
let bencher = Bencher::new(
client,
Handle::current(),
Duration::from_millis(cmd.delay_millis),
Duration::from_secs(opts.timeout_secs),
Duration::from_millis(cmd.delay_millis.get()),
Duration::from_secs(opts.timeout_secs.get()),
Duration::from_secs(cmd.ttl_secs as u64),
cmd.write_percent,
cmd.concurrency,
cmd.concurrency.get(),
stop.clone(),
);

let measurements = bencher.run(Duration::from_secs(cmd.time_secs)).await;
let measurements = bencher.run(Duration::from_secs(cmd.time_secs.into())).await;
print_bench_results(&measurements);

ExitCode::SUCCESS
Expand All @@ -421,11 +422,11 @@ async fn run_check(opts: &McConfig, cmd: &CheckCommand, client: MemcachedClient,
let checker = Checker::new(
client,
resolver,
Duration::from_millis(cmd.delay_millis),
Duration::from_secs(opts.timeout_secs),
Duration::from_millis(cmd.delay_millis.get()),
Duration::from_secs(opts.timeout_secs.get()),
stop.clone(),
);
let results = checker.run(&opts.host, Duration::from_secs(cmd.time_secs)).await;
let results = checker.run(&opts.host, Duration::from_secs(cmd.time_secs.get())).await;
print_check_results(&results);

if results.failures.total > 0 {
Expand All @@ -438,7 +439,7 @@ async fn run_check(opts: &McConfig, cmd: &CheckCommand, client: MemcachedClient,
async fn run_decr(opts: &McConfig, cmd: &DecrCommand, client: &MemcachedClient) -> ExitCode {
if let Err(e) = client
.decr(&cmd.key, cmd.delta)
.timeout(Duration::from_secs(opts.timeout_secs), "client.decr")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.decr")
.instrument(tracing::span!(Level::INFO, "client.decr"))
.await
{
Expand All @@ -452,7 +453,7 @@ async fn run_decr(opts: &McConfig, cmd: &DecrCommand, client: &MemcachedClient)
async fn run_delete(opts: &McConfig, cmd: &DeleteCommand, client: &MemcachedClient) -> ExitCode {
if let Err(e) = client
.delete(&cmd.key)
.timeout(Duration::from_secs(opts.timeout_secs), "client.delete")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.delete")
.instrument(tracing::span!(Level::INFO, "client.delete"))
.await
{
Expand All @@ -466,7 +467,7 @@ async fn run_delete(opts: &McConfig, cmd: &DeleteCommand, client: &MemcachedClie
async fn run_get(opts: &McConfig, cmd: &GetCommand, client: &MemcachedClient) -> ExitCode {
let response = match client
.get(&[cmd.key.clone()])
.timeout(Duration::from_secs(opts.timeout_secs), "client.get")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.get")
.instrument(tracing::span!(Level::INFO, "client.get"))
.await
{
Expand Down Expand Up @@ -497,7 +498,7 @@ async fn run_get(opts: &McConfig, cmd: &GetCommand, client: &MemcachedClient) ->
async fn run_incr(opts: &McConfig, cmd: &IncrCommand, client: &MemcachedClient) -> ExitCode {
if let Err(e) = client
.incr(&cmd.key, cmd.delta)
.timeout(Duration::from_secs(opts.timeout_secs), "client.incr")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.incr")
.instrument(tracing::span!(Level::INFO, "client.incr"))
.await
{
Expand All @@ -511,7 +512,7 @@ async fn run_incr(opts: &McConfig, cmd: &IncrCommand, client: &MemcachedClient)
async fn run_keys(opts: &McConfig, cmd: &KeysCommand, client: &MemcachedClient) -> ExitCode {
let response = match client
.metas()
.timeout(Duration::from_secs(opts.timeout_secs), "client.metas")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.metas")
.instrument(tracing::span!(Level::INFO, "client.metas"))
.await
{
Expand Down Expand Up @@ -552,7 +553,7 @@ async fn run_replace(opts: &McConfig, cmd: &ReplaceCommand, client: &MemcachedCl

if let Err(e) = client
.replace(&cmd.key, 0, cmd.ttl, &buf)
.timeout(Duration::from_secs(opts.timeout_secs), "client.replace")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.replace")
.instrument(tracing::span!(Level::INFO, "client.replace"))
.await
{
Expand All @@ -574,7 +575,7 @@ async fn run_set(opts: &McConfig, cmd: &SetCommand, client: &MemcachedClient) ->

if let Err(e) = client
.set(&cmd.key, 0, cmd.ttl, &buf)
.timeout(Duration::from_secs(opts.timeout_secs), "client.set")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.set")
.instrument(tracing::span!(Level::INFO, "client.set"))
.await
{
Expand All @@ -588,7 +589,7 @@ async fn run_set(opts: &McConfig, cmd: &SetCommand, client: &MemcachedClient) ->
async fn run_touch(opts: &McConfig, cmd: &TouchCommand, client: &MemcachedClient) -> ExitCode {
if let Err(e) = client
.touch(&cmd.key, cmd.ttl)
.timeout(Duration::from_secs(opts.timeout_secs), "client.touch")
.timeout(Duration::from_secs(opts.timeout_secs.get()), "client.touch")
.instrument(tracing::span!(Level::INFO, "client.touch"))
.await
{
Expand Down
7 changes: 4 additions & 3 deletions mtop/src/bin/mtop.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use mtop_client::{
};
use rustls_pki_types::{InvalidDnsNameError, ServerName};
use std::env;
use std::num::NonZeroU64;
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::Arc;
Expand Down Expand Up @@ -35,8 +36,8 @@ struct MtopConfig {
resolv_conf: PathBuf,

/// Timeout for connecting to Memcached and fetching statistics, in seconds.
#[arg(long, env = "MTOP_TIMEOUT_SECS", default_value_t = 5)]
timeout_secs: u64,
#[arg(long, env = "MTOP_TIMEOUT_SECS", default_value_t = NonZeroU64::new(5).unwrap())]
timeout_secs: NonZeroU64,

/// File to log errors to since they cannot be logged to the console. If the path is not
/// writable, mtop will not start.
Expand Down Expand Up @@ -113,7 +114,7 @@ async fn main() -> ExitCode {
}
};

let timeout = Duration::from_secs(opts.timeout_secs);
let timeout = Duration::from_secs(opts.timeout_secs.get());
let measurements = Arc::new(StatsQueue::new(NUM_MEASUREMENTS));
let dns_client = mtop::dns::new_client(&opts.resolv_conf, None, None).await;
let discovery = Discovery::new(dns_client);
Expand Down

0 comments on commit 3cf2192

Please sign in to comment.