From 7465aa429cf24e36235f162dc34c64d53926cb6d Mon Sep 17 00:00:00 2001 From: Robin Arenson Date: Thu, 3 Oct 2024 14:33:21 -0400 Subject: [PATCH 1/2] Retry logic and tests --- src/qos_net/src/proxy_stream.rs | 83 ++++++++++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 1 deletion(-) diff --git a/src/qos_net/src/proxy_stream.rs b/src/qos_net/src/proxy_stream.rs index 81b7fee4..ec923028 100644 --- a/src/qos_net/src/proxy_stream.rs +++ b/src/qos_net/src/proxy_stream.rs @@ -2,6 +2,9 @@ //! traits with `ProxyMsg`s. use std::io::{ErrorKind, Read, Write}; +use std::thread::sleep; +use std::time::Duration; + use borsh::BorshDeserialize; use qos_core::io::{SocketAddress, Stream, TimeVal}; @@ -294,10 +297,33 @@ impl Write for ProxyStream { } } +pub fn retry_with_backoff( + mut operation: F, + retry_count: u32, +) -> Result +where + F: FnMut() -> Result, +{ + let mut attempts = 0; + + loop { + match operation() { + Ok(result) => return Ok(result), + Err(_) if attempts < retry_count => { + attempts += 1; + let backoff_duration = + Duration::from_millis(2_u64.pow(attempts) * 100); // Exponential backoff + sleep(backoff_duration); + } + Err(e) => return Err(e), + } + } +} + #[cfg(test)] mod test { - use std::{io::ErrorKind, sync::Arc}; + use std::{io::ErrorKind, sync::Arc, cell::RefCell}; use chunked_transfer::Decoder; use httparse::Response; @@ -398,6 +424,61 @@ mod test { assert!(json_content["keys"].is_array()); } + #[test] + fn test_retry_with_backoff_success_after_retries() { + // This mock will fail the first 2 attempts, and succeed on the 3rd attempt. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + if *attempts <= 2 { + Err("fail") + } else { + Ok("success") + } + }; + + // Retry 3 times + let result: Result<&str, &str> = retry_with_backoff(operation, 3); + + assert_eq!(result, Ok("success")); + assert_eq!(*attempt_counter.borrow(), 3); + } + + #[test] + fn test_retry_with_backoff_failure_after_max_retries() { + // This mock will always fail. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + Err("fail") + }; + + // Retry 3 times + let result: Result<&str, &str> = retry_with_backoff(operation, 3); + + assert_eq!(result, Err("fail")); + assert_eq!(*attempt_counter.borrow(), 4); // 1 initial try + 3 retries + } + + #[test] + fn test_retry_with_backoff_no_retries() { + // This mock will fail the first time and there will be no retries. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + Err("fail") + }; + + // Retry 0 times + let result: Result<&str, &str> = retry_with_backoff(operation, 0); + + assert_eq!(result, Err("fail")); + assert_eq!(*attempt_counter.borrow(), 1); // Only 1 attempt + } + /// Struct representing a stream, with direct access to the proxy. /// Useful in tests! :) struct LocalStream { From 28bb019d7cf3fa1177fc09be2a88b665b6f2bad1 Mon Sep 17 00:00:00 2001 From: Robin Arenson Date: Thu, 3 Oct 2024 14:37:22 -0400 Subject: [PATCH 2/2] Moving to its own file --- src/qos_net/src/lib.rs | 1 + src/qos_net/src/proxy_stream.rs | 83 +------------------------------ src/qos_net/src/retry.rs | 88 +++++++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 82 deletions(-) create mode 100644 src/qos_net/src/retry.rs diff --git a/src/qos_net/src/lib.rs b/src/qos_net/src/lib.rs index f903c091..19e7c9fa 100644 --- a/src/qos_net/src/lib.rs +++ b/src/qos_net/src/lib.rs @@ -14,3 +14,4 @@ pub mod proxy; pub mod proxy_connection; pub mod proxy_msg; pub mod proxy_stream; +pub mod retry; diff --git a/src/qos_net/src/proxy_stream.rs b/src/qos_net/src/proxy_stream.rs index ec923028..81b7fee4 100644 --- a/src/qos_net/src/proxy_stream.rs +++ b/src/qos_net/src/proxy_stream.rs @@ -2,9 +2,6 @@ //! traits with `ProxyMsg`s. use std::io::{ErrorKind, Read, Write}; -use std::thread::sleep; -use std::time::Duration; - use borsh::BorshDeserialize; use qos_core::io::{SocketAddress, Stream, TimeVal}; @@ -297,33 +294,10 @@ impl Write for ProxyStream { } } -pub fn retry_with_backoff( - mut operation: F, - retry_count: u32, -) -> Result -where - F: FnMut() -> Result, -{ - let mut attempts = 0; - - loop { - match operation() { - Ok(result) => return Ok(result), - Err(_) if attempts < retry_count => { - attempts += 1; - let backoff_duration = - Duration::from_millis(2_u64.pow(attempts) * 100); // Exponential backoff - sleep(backoff_duration); - } - Err(e) => return Err(e), - } - } -} - #[cfg(test)] mod test { - use std::{io::ErrorKind, sync::Arc, cell::RefCell}; + use std::{io::ErrorKind, sync::Arc}; use chunked_transfer::Decoder; use httparse::Response; @@ -424,61 +398,6 @@ mod test { assert!(json_content["keys"].is_array()); } - #[test] - fn test_retry_with_backoff_success_after_retries() { - // This mock will fail the first 2 attempts, and succeed on the 3rd attempt. - let attempt_counter = RefCell::new(0); - let operation = || { - let mut attempts = attempt_counter.borrow_mut(); - *attempts += 1; - if *attempts <= 2 { - Err("fail") - } else { - Ok("success") - } - }; - - // Retry 3 times - let result: Result<&str, &str> = retry_with_backoff(operation, 3); - - assert_eq!(result, Ok("success")); - assert_eq!(*attempt_counter.borrow(), 3); - } - - #[test] - fn test_retry_with_backoff_failure_after_max_retries() { - // This mock will always fail. - let attempt_counter = RefCell::new(0); - let operation = || { - let mut attempts = attempt_counter.borrow_mut(); - *attempts += 1; - Err("fail") - }; - - // Retry 3 times - let result: Result<&str, &str> = retry_with_backoff(operation, 3); - - assert_eq!(result, Err("fail")); - assert_eq!(*attempt_counter.borrow(), 4); // 1 initial try + 3 retries - } - - #[test] - fn test_retry_with_backoff_no_retries() { - // This mock will fail the first time and there will be no retries. - let attempt_counter = RefCell::new(0); - let operation = || { - let mut attempts = attempt_counter.borrow_mut(); - *attempts += 1; - Err("fail") - }; - - // Retry 0 times - let result: Result<&str, &str> = retry_with_backoff(operation, 0); - - assert_eq!(result, Err("fail")); - assert_eq!(*attempt_counter.borrow(), 1); // Only 1 attempt - } - /// Struct representing a stream, with direct access to the proxy. /// Useful in tests! :) struct LocalStream { diff --git a/src/qos_net/src/retry.rs b/src/qos_net/src/retry.rs new file mode 100644 index 00000000..70100367 --- /dev/null +++ b/src/qos_net/src/retry.rs @@ -0,0 +1,88 @@ +use std::thread::sleep; +use std::time::Duration; + +pub fn retry_with_backoff( + mut operation: F, + retry_count: u32, +) -> Result +where + F: FnMut() -> Result, +{ + let mut attempts = 0; + + loop { + match operation() { + Ok(result) => return Ok(result), + Err(_) if attempts < retry_count => { + attempts += 1; + let backoff_duration = + Duration::from_millis(2_u64.pow(attempts) * 100); // Exponential backoff + sleep(backoff_duration); + } + Err(e) => return Err(e), + } + } +} + +#[cfg(test)] +mod test { + + use std::cell::RefCell; + + use super::*; + + #[test] + fn test_retry_with_backoff_success_after_retries() { + // This mock will fail the first 2 attempts, and succeed on the 3rd attempt. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + if *attempts <= 2 { + Err("fail") + } else { + Ok("success") + } + }; + + // Retry 3 times + let result: Result<&str, &str> = retry_with_backoff(operation, 3); + + assert_eq!(result, Ok("success")); + assert_eq!(*attempt_counter.borrow(), 3); + } + + #[test] + fn test_retry_with_backoff_failure_after_max_retries() { + // This mock will always fail. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + Err("fail") + }; + + // Retry 3 times + let result: Result<&str, &str> = retry_with_backoff(operation, 3); + + assert_eq!(result, Err("fail")); + assert_eq!(*attempt_counter.borrow(), 4); // 1 initial try + 3 retries + } + + #[test] + fn test_retry_with_backoff_no_retries() { + // This mock will fail the first time and there will be no retries. + let attempt_counter = RefCell::new(0); + let operation = || { + let mut attempts = attempt_counter.borrow_mut(); + *attempts += 1; + Err("fail") + }; + + // Retry 0 times + let result: Result<&str, &str> = retry_with_backoff(operation, 0); + + assert_eq!(result, Err("fail")); + assert_eq!(*attempt_counter.borrow(), 1); // Only 1 attempt + } +} \ No newline at end of file