refactor: use NonZero types

Signed-off-by: Wenxuan Zhang <[email protected]>

Author: wfxr

src/cli.rs

@@ -81,7 +81,10 @@
 //!
 //!   -h, --help
 //!           Print help (see a summary with '-h')
-use std::io::stdout;
+use std::{
+    io::stdout,
+    num::{NonZeroU32, NonZeroU64, NonZeroU8},
+};
 
 use clap::{Parser, ValueEnum};
 use crossterm::tty::IsTty;
@@ -102,14 +105,14 @@ use crate::{
 #[allow(missing_docs)]
 pub struct BenchCli {
     /// Number of workers to run concurrently
-    #[clap(long, short = 'c', default_value = "1", value_parser = clap::value_parser!(u32).range(1..))]
-    pub concurrency: u32,
+    #[clap(long, short = 'c', default_value = "1")]
+    pub concurrency: NonZeroU32,
 
     /// Number of iterations
     ///
     /// When set, benchmark stops after reaching the number of iterations.
-    #[clap(long, short = 'n', value_parser = clap::value_parser!(u64).range(1..))]
-    pub iterations: Option<u64>,
+    #[clap(long, short = 'n')]
+    pub iterations: Option<NonZeroU64>,
 
     /// Duration to run the benchmark
     ///
@@ -122,8 +125,8 @@ pub struct BenchCli {
     /// Rate limit for benchmarking, in iterations per second (ips)
     ///
     /// When set, benchmark will try to run at the specified rate.
-    #[clap(long, short = 'r', value_parser = clap::value_parser!(u32).range(1..))]
-    pub rate: Option<u32>,
+    #[clap(long, short = 'r')]
+    pub rate: Option<NonZeroU32>,
 
     /// Run benchmark in quiet mode
     ///
@@ -136,8 +139,8 @@ pub struct BenchCli {
     pub collector: Option<Collector>,
 
     /// Refresh rate for the tui collector, in frames per second (fps)
-    #[clap(long, default_value = "32", value_parser = clap::value_parser!(u8).range(1..))]
-    pub fps: u8,
+    #[clap(long, default_value = "32")]
+    pub fps: NonZeroU8,
 
     /// Output format for the report
     #[clap(short, long, value_enum, default_value_t = ReportFormat::Text, ignore_case = true)]
@@ -148,8 +151,8 @@ impl BenchCli {
     pub(crate) fn bench_opts(&self, clock: Clock) -> BenchOpts {
         BenchOpts {
             clock,
-            concurrency: self.concurrency,
-            iterations: self.iterations,
+            concurrency: self.concurrency.get(),
+            iterations: self.iterations.map(|n| n.get()),
             duration: self.duration.map(|d| d.into()),
             rate: self.rate,
         }

src/collector/tui.rs

@@ -6,6 +6,7 @@ use crossterm::{
     terminal, ExecutableCommand,
 };
 use itertools::Itertools;
+use nonzero_ext::nonzero;
 use ratatui::{
     backend::CrosstermBackend,
     layout::{Constraint, Direction, Layout, Margin, Rect},
@@ -14,7 +15,7 @@ use ratatui::{
     widgets::{block::Title, BarChart, Block, Borders, Clear, Gauge, Padding, Paragraph},
     CompletedFrame, Frame,
 };
-use std::{collections::HashMap, fmt, io, time::Duration};
+use std::{collections::HashMap, fmt, io, num::NonZeroU8, time::Duration};
 use tokio::{
     sync::{mpsc, watch},
     time::MissedTickBehavior,
@@ -47,7 +48,7 @@ pub struct TuiCollector {
     /// The benchmark options.
     pub bench_opts: BenchOpts,
     /// Refresh rate for the tui collector, in frames per second (fps)
-    pub fps: u8,
+    pub fps: NonZeroU8,
     /// The receiver for iteration reports.
     pub res_rx: mpsc::UnboundedReceiver<Result<IterReport>>,
     /// The sender for pausing the benchmark runner.
@@ -63,7 +64,7 @@ impl TuiCollector {
     /// Create a new TUI report collector.
     pub fn new(
         bench_opts: BenchOpts,
-        fps: u8,
+        fps: NonZeroU8,
         res_rx: mpsc::UnboundedReceiver<Result<IterReport>>,
         pause: watch::Sender<bool>,
         cancel: CancellationToken,
@@ -131,13 +132,13 @@ impl ReportCollector for TuiCollector {
 
         let mut clock = self.bench_opts.clock.clone();
 
-        let mut latest_iters = RotateWindowGroup::new(60);
+        let mut latest_iters = RotateWindowGroup::new(nonzero!(60usize));
         let mut latest_iters_ticker = clock.ticker(SECOND);
 
-        let mut latest_stats = RotateDiffWindowGroup::new(self.fps);
-        let mut latest_stats_ticker = clock.ticker(SECOND / self.fps as u32);
+        let mut latest_stats = RotateDiffWindowGroup::new(self.fps.into());
+        let mut latest_stats_ticker = clock.ticker(SECOND / self.fps.get() as u32);
 
-        let mut ui_ticker = tokio::time::interval(SECOND / self.fps as u32);
+        let mut ui_ticker = tokio::time::interval(SECOND / self.fps.get() as u32);
         ui_ticker.set_missed_tick_behavior(MissedTickBehavior::Burst);
 
         #[cfg(feature = "log")]

src/runner.rs

@@ -1,9 +1,10 @@
 //! This module defines traits for stateful and stateless benchmark suites.
-use anyhow::{anyhow, Result};
+use anyhow::Result;
 use async_trait::async_trait;
 use governor::{Quota, RateLimiter};
-use nonzero_ext::{nonzero, NonZero};
+use nonzero_ext::nonzero;
 use std::{
+    num::NonZeroU32,
     sync::{
         atomic::{AtomicU64, Ordering},
         Arc,
@@ -39,7 +40,7 @@ pub struct BenchOpts {
     pub duration: Option<Duration>,
 
     /// Rate limit for benchmarking, in iterations per second (ips).
-    pub rate: Option<u32>,
+    pub rate: Option<NonZeroU32>,
 }
 
 /// A trait for benchmark suites.
@@ -156,22 +157,16 @@ where
         let concurrency = self.opts.concurrency;
         let iterations = self.opts.iterations;
 
-        let rli = match self.opts.rate {
-            Some(r) => {
-                let quota = Quota::per_second(
-                    NonZero::new(r).ok_or_else(|| anyhow!("rate limit must be greater than 0, got {}", r))?,
-                )
-                .allow_burst(nonzero!(1u32));
-                let clock = &self.opts.clock;
-                Some(Arc::new(RateLimiter::direct_with_clock(quota, clock)))
-            }
-            None => None,
-        };
+        let buckets = self.opts.rate.map(|r| {
+            let quota = Quota::per_second(r).allow_burst(nonzero!(1u32));
+            let clock = &self.opts.clock;
+            Arc::new(RateLimiter::direct_with_clock(quota, clock))
+        });
 
         let mut set: JoinSet<Result<()>> = JoinSet::new();
         for worker in 0..concurrency {
             let mut b = self.clone();
-            let rli = rli.clone();
+            let buckets = buckets.clone();
             set.spawn(async move {
                 let mut state = b.suite.state(worker).await?;
                 let mut info = IterInfo::new(worker);
@@ -186,14 +181,16 @@ where
                         }
                     }
 
-                    if let Some(rli) = &rli {
+                    if let Some(buckets) = &buckets {
                         select! {
+                            biased;
                             _ = cancel.cancelled() => break,
-                            _ = rli.until_ready() => (),
+                            _ = buckets.until_ready() => (),
                         }
                     }
 
                     select! {
+                        biased;
                         _ = cancel.cancelled() => break,
                         _ = b.iteration(&mut state, &info) => (),
                     }
@@ -207,6 +204,7 @@ where
 
         if let Some(t) = self.opts.duration {
             select! {
+                biased;
                 _ = self.cancel.cancelled() => (),
                 _ = self.opts.clock.sleep(t) => self.cancel.cancel(),
                 _ = join_all(&mut set) => (),

src/stats/window.rs

@@ -1,5 +1,6 @@
-use std::collections::VecDeque;
+use std::{collections::VecDeque, num::NonZeroUsize};
 
+use nonzero_ext::nonzero;
 use tokio::time::Duration;
 
 use crate::report::IterReport;
@@ -8,13 +9,12 @@ use super::IterStats;
 
 pub struct RotateWindow {
     buckets: VecDeque<IterStats>,
-    size: usize,
+    size: NonZeroUsize,
 }
 
 impl RotateWindow {
-    fn new(size: usize) -> Self {
-        assert!(size > 0);
-        let mut win = Self { buckets: VecDeque::with_capacity(size), size };
+    fn new(size: NonZeroUsize) -> Self {
+        let mut win = Self { buckets: VecDeque::with_capacity(size.get()), size };
         win.rotate(IterStats::new());
         win
     }
@@ -25,7 +25,7 @@ impl RotateWindow {
     }
 
     fn rotate(&mut self, bucket: IterStats) {
-        if self.buckets.len() == self.size {
+        if self.buckets.len() == self.size.get() {
             self.buckets.pop_back();
         }
         self.buckets.push_front(bucket);
@@ -59,7 +59,7 @@ pub struct RotateWindowGroup {
 }
 
 impl RotateWindowGroup {
-    pub fn new(buckets: usize) -> Self {
+    pub fn new(buckets: NonZeroUsize) -> Self {
         Self {
             counter: 0,
             stats_by_sec: RotateWindow::new(buckets),
@@ -108,15 +108,14 @@ impl RotateDiffWindowGroup {
             &mut self.stats_last_10min,
         ]
     }
-    pub fn new(fps: u8) -> Self {
-        let fps = fps as usize;
-        let interval = Duration::from_secs_f64(1.0 / fps as f64);
+    pub fn new(fps: NonZeroUsize) -> Self {
+        let interval = Duration::from_secs_f64(1.0 / fps.get() as f64);
         let mut group = Self {
             interval,
-            stats_last_sec: RotateWindow::new(fps + 1),
-            stats_last_10sec: RotateWindow::new(fps * 10 + 1),
-            stats_last_min: RotateWindow::new(fps * 60 + 1),
-            stats_last_10min: RotateWindow::new(fps * 600 + 1),
+            stats_last_sec: RotateWindow::new(fps.saturating_add(1)),
+            stats_last_10sec: RotateWindow::new(fps.saturating_mul(nonzero!(10usize)).saturating_add(1)),
+            stats_last_min: RotateWindow::new(fps.saturating_mul(nonzero!(60usize)).saturating_add(1)),
+            stats_last_10min: RotateWindow::new(fps.saturating_mul(nonzero!(600usize)).saturating_add(1)),
         };
         group.rotate(&IterStats::new());
         group