diff --git a/src/util.rs b/src/util.rs index 9d2327b..2fcfaca 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,25 +1,90 @@ use core::time::Duration; use std::thread::sleep; -const RETRY_INTERVAL_MSEC: u64 = 1000; - +/// Retries the supplied function until it returns `Ok` or the supplied maximum +/// retry limit is reached. +/// +/// # Examples +/// +/// ```rust +/// use std::{fs, io}; +/// use ue_rs::retry_loop; +/// use std::sync::atomic::{AtomicUsize, Ordering}; +/// use std::path::Path; +/// +/// fn read_possibly_extant_file>(path: P) -> io::Result { +/// let file_path = Path::new(path.as_ref()); +/// +/// if file_path.exists() { +/// fs::read_to_string(file_path) +/// } else { +/// Err(io::Error::new(io::ErrorKind::NotFound, io::Error::last_os_error())) +/// } +/// } +/// +/// let result = retry_loop(|| read_possibly_extant_file("might_exist.txt"), 3); +/// ``` pub fn retry_loop(mut func: F, max_tries: u32) -> Result where F: FnMut() -> Result, { - let mut tries = 0; + const RETRY_INTERVAL: Duration = Duration::from_millis(1000); - loop { + for _ in 0..max_tries - 1 { match func() { - ok @ Ok(_) => return ok, - err @ Err(_) => { - tries += 1; - - if tries >= max_tries { - return err; - } - sleep(Duration::from_millis(RETRY_INTERVAL_MSEC)); - } + Ok(ret) => return Ok(ret), + Err(_) => sleep(RETRY_INTERVAL), } } + + func() +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + use super::*; + + #[test] + fn test_success_first_try() { + let res: Result<(), ()> = retry_loop(|| Ok(()), 3); + assert_eq!(res, Ok(())); + } + + #[test] + fn test_fail_first_try() { + let res: Result = retry_loop(|| Err(()), 3); + assert_eq!(res, Err(())); + } + + #[test] + fn test_success_after_some_retries() { + let attempts = Cell::new(0); + + let mut func = || { + let current: u8 = attempts.get(); + attempts.set(current.wrapping_add(1)); + if current < 2 { + Err(()) + } else { + Ok(()) + } + }; + + let res = retry_loop(&mut func, 5); + assert_eq!(res, Ok(())); + assert_eq!(attempts.get(), 3); + } + + #[test] + fn test_fail_after_max_tries() { + let attempts = Cell::new(0); + let mut func = || { + attempts.set(attempts.get() + 1); + Err(()) + }; + let res: Result<(), ()> = retry_loop(&mut func, 5); + assert_eq!(res, Err(())); + assert_eq!(attempts.get(), 5); + } }