From 7d0a5956091a3ba27c1656bbfbe56c3368e09593 Mon Sep 17 00:00:00 2001
From: "Heinz N. Gies" <heinz@licenser.net>
Date: Wed, 29 Jun 2022 14:13:14 +0200
Subject: [PATCH] Add parameter for timeout on a per test level

Signed-off-by: Heinz N. Gies <heinz@licenser.net>
---
 serial_test/src/code_lock.rs          |  31 +------
 serial_test/src/lib.rs                |   1 -
 serial_test/src/parallel_code_lock.rs |  22 +++--
 serial_test/src/serial_code_lock.rs   |  33 ++++---
 serial_test/tests/tests.rs            |   2 +-
 serial_test_derive/src/lib.rs         | 125 ++++++++++++++++----------
 6 files changed, 117 insertions(+), 97 deletions(-)

diff --git a/serial_test/src/code_lock.rs b/serial_test/src/code_lock.rs
index 1b5dffb..d6b86cf 100644
--- a/serial_test/src/code_lock.rs
+++ b/serial_test/src/code_lock.rs
@@ -1,7 +1,7 @@
 use crate::rwlock::{Locks, MutexGuardWrapper};
 use lazy_static::lazy_static;
 #[cfg(feature = "logging")]
-use log::{debug, warn};
+use log::debug;
 use parking_lot::RwLock;
 use std::{
     collections::HashMap,
@@ -48,11 +48,6 @@ lazy_static! {
     static ref MUTEX_ID: Arc<AtomicU32> = Arc::new(AtomicU32::new(1));
 }
 
-#[cfg(feature = "timeout")]
-lazy_static! {
-    static ref MAX_WAIT: Arc<RwLock<Duration>> = Arc::new(RwLock::new(Duration::from_secs(60)));
-}
-
 impl Default for UniqueReentrantMutex {
     fn default() -> Self {
         Self {
@@ -62,25 +57,7 @@ impl Default for UniqueReentrantMutex {
     }
 }
 
-/// Sets the maximum amount of time the serial locks will wait to unlock.
-/// By default, this is set to 60 seconds, which is almost always much longer than is needed.
-/// This is deliberately set high to try and avoid situations where we accidentally hit the limits
-/// but is set at all so we can timeout rather than hanging forever.
-///
-/// However, sometimes if you've got a *lot* of serial tests it might theoretically not be enough,
-/// hence this method.
-///
-/// This function is only available when the `timeout` feature is enabled.
-#[cfg(feature = "timeout")]
-pub fn set_max_wait(max_wait: Duration) {
-    *MAX_WAIT.write() = max_wait;
-}
-
-pub(crate) fn wait_duration() -> Duration {
-    *MAX_WAIT.read()
-}
-
-pub(crate) fn check_new_key(name: &str) {
+pub(crate) fn check_new_key(name: &str, max_wait: Option<Duration>) {
     let start = Instant::now();
     loop {
         #[cfg(all(feature = "logging", feature = "timeout"))]
@@ -111,9 +88,9 @@ pub(crate) fn check_new_key(name: &str) {
         // Odds are another test was also locking on the write and has now written the key
 
         #[cfg(feature = "timeout")]
-        {
+        if let Some(max_wait) = max_wait {
             let duration = start.elapsed();
-            if duration > wait_duration() {
+            if duration > max_wait {
                 panic!("Timeout waiting for '{}' {:?}", name, duration);
             }
         }
diff --git a/serial_test/src/lib.rs b/serial_test/src/lib.rs
index 00032fa..01268a1 100644
--- a/serial_test/src/lib.rs
+++ b/serial_test/src/lib.rs
@@ -63,7 +63,6 @@ mod parallel_file_lock;
 #[cfg(feature = "file_locks")]
 mod serial_file_lock;
 
-pub use code_lock::set_max_wait;
 pub use parallel_code_lock::{
     local_async_parallel_core, local_async_parallel_core_with_return, local_parallel_core,
     local_parallel_core_with_return,
diff --git a/serial_test/src/parallel_code_lock.rs b/serial_test/src/parallel_code_lock.rs
index e854cec..e8af569 100644
--- a/serial_test/src/parallel_code_lock.rs
+++ b/serial_test/src/parallel_code_lock.rs
@@ -2,14 +2,15 @@
 
 use crate::code_lock::{check_new_key, LOCKS};
 use futures::FutureExt;
-use std::{ops::Deref, panic};
+use std::{ops::Deref, panic, time::Duration};
 
 #[doc(hidden)]
 pub fn local_parallel_core_with_return<E>(
     name: &str,
+    max_wait: Option<Duration>,
     function: fn() -> Result<(), E>,
 ) -> Result<(), E> {
-    check_new_key(name);
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     unlock.deref()[name].start_parallel();
@@ -24,8 +25,8 @@ pub fn local_parallel_core_with_return<E>(
 }
 
 #[doc(hidden)]
-pub fn local_parallel_core(name: &str, function: fn()) {
-    check_new_key(name);
+pub fn local_parallel_core(name: &str, max_wait: Option<Duration>, function: fn()) {
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     unlock.deref()[name].start_parallel();
@@ -41,9 +42,10 @@ pub fn local_parallel_core(name: &str, function: fn()) {
 #[doc(hidden)]
 pub async fn local_async_parallel_core_with_return<E>(
     name: &str,
+    max_wait: Option<Duration>,
     fut: impl std::future::Future<Output = Result<(), E>> + panic::UnwindSafe,
 ) -> Result<(), E> {
-    check_new_key(name);
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     unlock.deref()[name].start_parallel();
@@ -60,9 +62,10 @@ pub async fn local_async_parallel_core_with_return<E>(
 #[doc(hidden)]
 pub async fn local_async_parallel_core(
     name: &str,
+    max_wait: Option<Duration>,
     fut: impl std::future::Future<Output = ()> + panic::UnwindSafe,
 ) {
-    check_new_key(name);
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     unlock.deref()[name].start_parallel();
@@ -84,7 +87,7 @@ mod tests {
     #[test]
     fn unlock_on_assert_sync_without_return() {
         let _ = panic::catch_unwind(|| {
-            local_parallel_core("unlock_on_assert_sync_without_return", || {
+            local_parallel_core("unlock_on_assert_sync_without_return", None, || {
                 assert!(false);
             })
         });
@@ -100,6 +103,7 @@ mod tests {
         let _ = panic::catch_unwind(|| {
             local_parallel_core_with_return(
                 "unlock_on_assert_sync_with_return",
+                None,
                 || -> Result<(), Error> {
                     assert!(false);
                     Ok(())
@@ -119,7 +123,8 @@ mod tests {
             assert!(false);
         }
         async fn call_serial_test_fn() {
-            local_async_parallel_core("unlock_on_assert_async_without_return", demo_assert()).await
+            local_async_parallel_core("unlock_on_assert_async_without_return", None, demo_assert())
+                .await
         }
         // as per https://stackoverflow.com/a/66529014/320546
         let _ = panic::catch_unwind(|| {
@@ -145,6 +150,7 @@ mod tests {
         async fn call_serial_test_fn() {
             local_async_parallel_core_with_return(
                 "unlock_on_assert_async_with_return",
+                None,
                 demo_assert(),
             )
             .await;
diff --git a/serial_test/src/serial_code_lock.rs b/serial_test/src/serial_code_lock.rs
index 6c80687..3384023 100644
--- a/serial_test/src/serial_code_lock.rs
+++ b/serial_test/src/serial_code_lock.rs
@@ -1,14 +1,15 @@
 #![allow(clippy::await_holding_lock)]
 
 use crate::code_lock::{check_new_key, LOCKS};
-use std::ops::Deref;
+use std::{ops::Deref, time::Duration};
 
 #[doc(hidden)]
 pub fn local_serial_core_with_return<E>(
     name: &str,
+    max_wait: Option<Duration>,
     function: fn() -> Result<(), E>,
 ) -> Result<(), E> {
-    check_new_key(name);
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     // _guard needs to be named to avoid being instant dropped
@@ -17,8 +18,8 @@ pub fn local_serial_core_with_return<E>(
 }
 
 #[doc(hidden)]
-pub fn local_serial_core(name: &str, function: fn()) {
-    check_new_key(name);
+pub fn local_serial_core(name: &str, max_wait: Option<Duration>, function: fn()) {
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     // _guard needs to be named to avoid being instant dropped
@@ -29,9 +30,10 @@ pub fn local_serial_core(name: &str, function: fn()) {
 #[doc(hidden)]
 pub async fn local_async_serial_core_with_return<E>(
     name: &str,
+    max_wait: Option<Duration>,
     fut: impl std::future::Future<Output = Result<(), E>>,
 ) -> Result<(), E> {
-    check_new_key(name);
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     // _guard needs to be named to avoid being instant dropped
@@ -40,8 +42,12 @@ pub async fn local_async_serial_core_with_return<E>(
 }
 
 #[doc(hidden)]
-pub async fn local_async_serial_core(name: &str, fut: impl std::future::Future<Output = ()>) {
-    check_new_key(name);
+pub async fn local_async_serial_core(
+    name: &str,
+    max_wait: Option<Duration>,
+    fut: impl std::future::Future<Output = ()>,
+) {
+    check_new_key(name, max_wait);
 
     let unlock = LOCKS.read_recursive();
     // _guard needs to be named to avoid being instant dropped
@@ -53,13 +59,14 @@ pub async fn local_async_serial_core(name: &str, fut: impl std::future::Future<O
 #[allow(clippy::print_stdout)]
 mod tests {
     use super::local_serial_core;
-    use crate::code_lock::{check_new_key, wait_duration, LOCKS};
+    use crate::code_lock::{check_new_key, LOCKS};
     use itertools::Itertools;
     use parking_lot::RwLock;
     use std::{
         ops::Deref,
         sync::{Arc, Barrier},
         thread,
+        time::Duration,
     };
 
     #[test]
@@ -76,15 +83,15 @@ mod tests {
             let c = barrier.clone();
             threads.push(thread::spawn(move || {
                 c.wait();
-                check_new_key("foo");
+                check_new_key("foo", None);
                 {
                     let unlock = local_locks
-                        .try_read_recursive_for(wait_duration())
+                        .try_read_recursive_for(Duration::from_secs(1))
                         .expect("read lock didn't work");
                     let mutex = unlock.deref().get("foo").unwrap();
 
                     let mut ptr_guard = local_ptrs
-                        .try_write_for(wait_duration())
+                        .try_write_for(Duration::from_secs(1))
                         .expect("write lock didn't work");
                     ptr_guard.push(mutex.id);
                 }
@@ -96,7 +103,7 @@ mod tests {
             thread.join().expect("thread join worked");
         }
         let ptrs_read_lock = ptrs
-            .try_read_recursive_for(wait_duration())
+            .try_read_recursive_for(Duration::from_secs(1))
             .expect("ptrs read work");
         assert_eq!(ptrs_read_lock.len(), count);
         println!("{:?}", ptrs_read_lock);
@@ -106,7 +113,7 @@ mod tests {
     #[test]
     fn unlock_on_assert() {
         let _ = std::panic::catch_unwind(|| {
-            local_serial_core("assert", || {
+            local_serial_core("assert", None, || {
                 assert!(false);
             })
         });
diff --git a/serial_test/tests/tests.rs b/serial_test/tests/tests.rs
index add8b10..63213f8 100644
--- a/serial_test/tests/tests.rs
+++ b/serial_test/tests/tests.rs
@@ -2,7 +2,7 @@ use serial_test::local_serial_core;
 
 #[test]
 fn test_empty_serial_call() {
-    local_serial_core("beta", || {
+    local_serial_core("beta", None, || {
         println!("Bar");
     });
 }
diff --git a/serial_test_derive/src/lib.rs b/serial_test_derive/src/lib.rs
index 56bb310..57ab08a 100644
--- a/serial_test_derive/src/lib.rs
+++ b/serial_test_derive/src/lib.rs
@@ -201,13 +201,37 @@ impl<T: ToTokens> ToTokens for QuoteOption<T> {
     }
 }
 
-fn get_raw_args(attr: proc_macro2::TokenStream) -> Vec<String> {
+enum Arg {
+    Name(String),
+    Timeout(u64),
+}
+
+fn get_raw_args(attr: proc_macro2::TokenStream) -> Vec<Arg> {
     let mut attrs = attr.into_iter().collect::<Vec<TokenTree>>();
-    let mut raw_args: Vec<String> = Vec::new();
+    let mut raw_args: Vec<Arg> = Vec::new();
     while !attrs.is_empty() {
         match attrs.remove(0) {
             TokenTree::Ident(id) => {
-                raw_args.push(id.to_string());
+                let name = id.to_string();
+                if name == "timeout" {
+                    match attrs.first() {
+                        Some(TokenTree::Punct(p)) if p.as_char() == '=' && !attrs.is_empty() => {
+                            attrs.remove(0);
+                            if let TokenTree::Literal(lit) = attrs.remove(0) {
+                                let millis = lit
+                                    .to_string()
+                                    .parse::<u64>()
+                                    .expect("Not a valid duration for Timeout");
+                                raw_args.push(Arg::Timeout(millis));
+                            } else {
+                                panic!("Timeout argument must be a literal duration");
+                            }
+                        }
+                        _ => raw_args.push(Arg::Name(name)),
+                    }
+                } else {
+                    raw_args.push(Arg::Name(name));
+                }
             }
             TokenTree::Literal(literal) => {
                 let string_literal = literal.to_string();
@@ -215,7 +239,9 @@ fn get_raw_args(attr: proc_macro2::TokenStream) -> Vec<String> {
                     panic!("Expected a string literal, got '{}'", string_literal);
                 }
                 // Hacky way of getting a string without the enclosing quotes
-                raw_args.push(string_literal[1..string_literal.len() - 1].to_string());
+                raw_args.push(Arg::Name(
+                    string_literal[1..string_literal.len() - 1].to_string(),
+                ));
             }
             x => {
                 panic!("Expected either strings or literals as args, not {}", x);
@@ -233,60 +259,65 @@ fn get_raw_args(attr: proc_macro2::TokenStream) -> Vec<String> {
     raw_args
 }
 
-fn get_core_key(attr: proc_macro2::TokenStream) -> String {
-    let mut raw_args = get_raw_args(attr);
-    match raw_args.len() {
-        0 => "".to_string(),
-        1 => raw_args.pop().unwrap(),
-        n => {
-            panic!(
-                "Expected either 0 or 1 arguments, got {}: {:?}",
-                n, raw_args
-            );
+#[derive(Default, Debug)]
+struct Config {
+    name: String,
+    timeout: Option<u64>,
+    path: Option<String>,
+}
+
+fn get_core_key(attr: proc_macro2::TokenStream) -> Config {
+    let raw_args = get_raw_args(attr);
+    let mut c = Config::default();
+    let mut name_found = false;
+    for a in raw_args {
+        match a {
+            Arg::Name(name) if !name_found => {
+                c.name = name;
+                name_found = true;
+            }
+            Arg::Name(name) => {
+                c.path = Some(name);
+            }
+            Arg::Timeout(timeout) => {
+                c.timeout = Some(timeout);
+            }
         }
     }
+    c
 }
 
 fn local_serial_core(
     attr: proc_macro2::TokenStream,
     input: proc_macro2::TokenStream,
 ) -> proc_macro2::TokenStream {
-    let key = get_core_key(attr);
-    serial_setup(input, vec![Box::new(key)], "local")
+    let config = get_core_key(attr);
+    let args: Vec<Box<dyn ToTokens>> =
+        vec![Box::new(config.name), Box::new(QuoteOption(config.timeout))];
+    serial_setup(input, args, "local")
 }
 
 fn local_parallel_core(
     attr: proc_macro2::TokenStream,
     input: proc_macro2::TokenStream,
 ) -> proc_macro2::TokenStream {
-    let key = get_core_key(attr);
-    parallel_setup(input, vec![Box::new(key)], "local")
+    let config = get_core_key(attr);
+    let args: Vec<Box<dyn ToTokens>> =
+        vec![Box::new(config.name), Box::new(QuoteOption(config.timeout))];
+    parallel_setup(input, args, "local")
 }
 
 fn fs_args(attr: proc_macro2::TokenStream) -> Vec<Box<dyn ToTokens>> {
-    let none_ident = Box::new(format_ident!("None"));
-    let mut args: Vec<Box<dyn quote::ToTokens>> = Vec::new();
-    let mut raw_args = get_raw_args(attr);
-    match raw_args.len() {
-        0 => {
-            args.push(Box::new("".to_string()));
-            args.push(none_ident);
-        }
-        1 => {
-            args.push(Box::new(raw_args.pop().unwrap()));
-            args.push(none_ident);
-        }
-        2 => {
-            let key = raw_args.remove(0);
-            let path = raw_args.remove(0);
-            args.push(Box::new(key));
-            args.push(Box::new(QuoteOption(Some(path))));
-        }
-        n => {
-            panic!("Expected 0-2 arguments, got {}: {:?}", n, raw_args);
-        }
-    }
-    args
+    let config = get_core_key(attr);
+    vec![
+        Box::new(config.name),
+        Box::new(QuoteOption(config.timeout)),
+        if let Some(path) = config.path {
+            Box::new(QuoteOption(Some(path)))
+        } else {
+            Box::new(format_ident!("None"))
+        },
+    ]
 }
 
 fn fs_serial_core(
@@ -435,7 +466,7 @@ mod tests {
         let compare = quote! {
             #[test]
             fn foo () {
-                serial_test::local_serial_core("", || {} );
+                serial_test::local_serial_core("",  :: std :: option :: Option :: None, || {} );
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));
@@ -457,7 +488,7 @@ mod tests {
             #[test]
             #[something_else]
             fn foo () {
-                serial_test::local_serial_core("", || {} );
+                serial_test::local_serial_core("",  :: std :: option :: Option :: None, || {} );
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));
@@ -472,7 +503,7 @@ mod tests {
         let stream = local_serial_core(attrs.into(), input);
         let compare = quote! {
             async fn foo () {
-                serial_test::local_async_serial_core("", || async {} ).await;
+                serial_test::local_async_serial_core("",  :: std :: option :: Option :: None, || async {} ).await;
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));
@@ -487,7 +518,7 @@ mod tests {
         let stream = local_serial_core(attrs.into(), input);
         let compare = quote! {
             async fn foo () -> Result<(), ()> {
-                serial_test::local_async_serial_core_with_return("", || async { Ok(()) } ).await;
+                serial_test::local_async_serial_core_with_return("", :: std :: option :: Option :: None, || async { Ok(()) } ).await;
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));
@@ -515,7 +546,7 @@ mod tests {
         let compare = quote! {
             #[test]
             fn foo () {
-                serial_test::fs_serial_core("foo", None, || {} );
+                serial_test::fs_serial_core("foo",  :: std :: option :: Option :: None, None, || {} );
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));
@@ -539,7 +570,7 @@ mod tests {
         let compare = quote! {
             #[test]
             fn foo () {
-                serial_test::fs_serial_core("foo", ::std::option::Option::Some("bar_path"), || {} );
+                serial_test::fs_serial_core("foo",  :: std :: option :: Option :: None, ::std::option::Option::Some("bar_path"), || {} );
             }
         };
         assert_eq!(format!("{}", compare), format!("{}", stream));