safe point

Signed-off-by: TXXT <[email protected]>

Author: hunterlxt

src/lib.rs

@@ -230,7 +230,7 @@ use std::env::VarError;
 use std::fmt::Debug;
 use std::str::FromStr;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, Condvar, Mutex, MutexGuard, RwLock, TryLockError};
+use std::sync::{Arc, Condvar, Mutex, RwLock, TryLockError};
 use std::time::{Duration, Instant};
 use std::{env, thread};
 
@@ -421,7 +421,7 @@ impl FromStr for Action {
 }
 
 #[cfg_attr(feature = "cargo-clippy", allow(clippy::mutex_atomic))]
-#[derive(Debug)]
+#[derive(Debug, Default)]
 struct FailPoint {
     pause: Mutex<bool>,
     pause_notifier: Condvar,
@@ -509,98 +509,115 @@ impl FailPoint {
 /// Registry with failpoints configuration.
 type Registry = HashMap<String, Arc<FailPoint>>;
 
-#[derive(Debug, Default)]
-struct FailPointRegistry {
+/// Fail point registry. Each thread should be bound to exact one registry.
+/// Threads bound to the same registry share the same failpoints configuration.
+#[derive(Debug, Default, Clone)]
+pub struct FailPointRegistry {
     // TODO: remove rwlock or store *mut FailPoint
-    registry: RwLock<Registry>,
+    registry: Arc<RwLock<Registry>>,
 }
 
-lazy_static::lazy_static! {
-    static ref REGISTRY: FailPointRegistry = FailPointRegistry::default();
-    static ref SCENARIO: Mutex<&'static FailPointRegistry> = Mutex::new(&REGISTRY);
+/// Generate a new failpoint registry. The new registry will inherit the
+/// global failpoints configuration.
+pub fn new_fail_group() -> FailPointRegistry {
+    let registry = REGISTRY_GLOBAL.registry.read().unwrap().clone();
+    FailPointRegistry {
+        registry: Arc::new(RwLock::new(registry)),
+    }
 }
 
-/// Test scenario with configured fail points.
-#[derive(Debug)]
-pub struct FailScenario<'a> {
-    scenario_guard: MutexGuard<'a, &'static FailPointRegistry>,
-}
+impl FailPointRegistry {
+    /// Register the current thread to this failpoints registry.
+    pub fn register_current(&self) {
+        let id = thread::current().id();
+        REGISTRY_GROUP
+            .write()
+            .unwrap()
+            .insert(id, self.registry.clone());
+    }
 
-impl<'a> FailScenario<'a> {
-    /// Set up the system for a fail points scenario.
-    ///
-    /// Configures all fail points specified in the `FAILPOINTS` environment variable.
-    /// It does not otherwise change any existing fail point configuration.
-    ///
-    /// The format of `FAILPOINTS` is `failpoint=actions;...`, where
-    /// `failpoint` is the name of the fail point. For more information
-    /// about fail point actions see the [`cfg`](fn.cfg.html) function and
-    /// the [`fail_point`](macro.fail_point.html) macro.
-    ///
-    /// `FAILPOINTS` may configure fail points that are not actually defined. In
-    /// this case the configuration has no effect.
-    ///
-    /// This function should generally be called prior to running a test with fail
-    /// points, and afterward paired with [`teardown`](#method.teardown).
-    ///
-    /// # Panics
-    ///
-    /// Panics if an action is not formatted correctly.
-    pub fn setup() -> Self {
-        // Cleanup first, in case of previous failed/panic'ed test scenarios.
-        let scenario_guard = SCENARIO.lock().unwrap_or_else(|e| e.into_inner());
-        let mut registry = scenario_guard.registry.write().unwrap();
-        Self::cleanup(&mut registry);
-
-        let failpoints = match env::var("FAILPOINTS") {
-            Ok(s) => s,
-            Err(VarError::NotPresent) => return Self { scenario_guard },
-            Err(e) => panic!("invalid failpoints: {:?}", e),
-        };
-        for mut cfg in failpoints.trim().split(';') {
-            cfg = cfg.trim();
-            if cfg.is_empty() {
-                continue;
-            }
-            let (name, order) = partition(cfg, '=');
-            match order {
-                None => panic!("invalid failpoint: {:?}", cfg),
-                Some(order) => {
-                    if let Err(e) = set(&mut registry, name.to_owned(), order) {
-                        panic!("unable to configure failpoint \"{}\": {}", name, e);
-                    }
-                }
-            }
-        }
-        Self { scenario_guard }
+    /// Deregister the current thread to this failpoints registry.
+    pub fn deregister_current(&self) {
+        let id = thread::current().id();
+        REGISTRY_GROUP.write().unwrap().remove(&id);
     }
 
-    /// Tear down the fail point system.
-    ///
-    /// Clears the configuration of all fail points. Any paused fail
-    /// points will be notified before they are deactivated.
-    ///
-    /// This function should generally be called after running a test with fail points.
-    /// Calling `teardown` without previously calling `setup` results in a no-op.
-    pub fn teardown(self) {
-        drop(self)
+    /// Clean registered fail points in this registry.
+    pub fn cleanup(&self) {
+        let mut registry = self.registry.write().unwrap();
+        cleanup(&mut registry);
     }
+}
 
-    /// Clean all registered fail points.
-    fn cleanup(registry: &mut std::sync::RwLockWriteGuard<'a, Registry>) {
-        for p in registry.values() {
-            // wake up all pause failpoint.
-            p.set_actions("", vec![]);
+lazy_static::lazy_static! {
+    static ref REGISTRY_GROUP: RwLock<HashMap<thread::ThreadId, Arc<RwLock<Registry>>>> = Default::default();
+    static ref REGISTRY_GLOBAL: FailPointRegistry = Default::default();
+}
+
+/// Set up the system for a fail points scenario.
+///
+/// Configures all fail points specified in the `FAILPOINTS` environment variable.
+/// It does not otherwise change any existing fail point configuration.
+///
+/// The format of `FAILPOINTS` is `failpoint=actions;...`, where
+/// `failpoint` is the name of the fail point. For more information
+/// about fail point actions see the [`cfg`](fn.cfg.html) function and
+/// the [`fail_point`](macro.fail_point.html) macro.
+///
+/// `FAILPOINTS` may configure fail points that are not actually defined. In
+/// this case the configuration has no effect.
+///
+/// This function should generally be called prior to running a test with fail
+/// points, and afterward paired with [`teardown`](#method.teardown).
+///
+/// # Panics
+///
+/// Panics if an action is not formatted correctly.
+pub fn setup() {
+    let mut registry = REGISTRY_GLOBAL.registry.write().unwrap();
+    cleanup(&mut registry);
+
+    let failpoints = match env::var("FAILPOINTS") {
+        Ok(s) => s,
+        Err(VarError::NotPresent) => return,
+        Err(e) => panic!("invalid failpoints: {:?}", e),
+    };
+    for mut cfg in failpoints.trim().split(';') {
+        cfg = cfg.trim();
+        if cfg.is_empty() {
+            continue;
+        }
+        let (name, order) = partition(cfg, '=');
+        match order {
+            None => panic!("invalid failpoint: {:?}", cfg),
+            Some(order) => {
+                if let Err(e) = set(&mut registry, name.to_owned(), order) {
+                    panic!("unable to configure failpoint \"{}\": {}", name, e);
+                }
+            }
         }
-        registry.clear();
     }
 }
 
-impl<'a> Drop for FailScenario<'a> {
-    fn drop(&mut self) {
-        let mut registry = self.scenario_guard.registry.write().unwrap();
-        Self::cleanup(&mut registry)
+/// Tear down the fail point system.
+///
+/// Clears the configuration of all fail points. Any paused fail
+/// points will be notified before they are deactivated.
+///
+/// This function should generally be called after running a test with fail points.
+/// Calling `teardown` without previously calling `setup` results in a no-op.
+pub fn teardown() {
+    let mut registry = REGISTRY_GLOBAL.registry.write().unwrap();
+    cleanup(&mut registry);
+}
+
+/// Clean all registered fail points.
+fn cleanup(registry: &mut std::sync::RwLockWriteGuard<Registry>) {
+    for p in registry.values() {
+        // wake up all pause failpoint.
+        p.set_actions("", vec![]);
     }
+    registry.clear();
 }
 
 /// Returns whether code generation for failpoints is enabled.
@@ -616,7 +633,15 @@ pub const fn has_failpoints() -> bool {
 ///
 /// Return a vector of `(name, actions)` pairs.
 pub fn list() -> Vec<(String, String)> {
-    let registry = REGISTRY.registry.read().unwrap();
+    let id = thread::current().id();
+    let group = REGISTRY_GROUP.read().unwrap();
+
+    let registry = group
+        .get(&id)
+        .unwrap_or(&REGISTRY_GLOBAL.registry)
+        .read()
+        .unwrap();
+
     registry
         .iter()
         .map(|(name, fp)| (name.to_string(), fp.actions_str.read().unwrap().clone()))
@@ -625,8 +650,16 @@ pub fn list() -> Vec<(String, String)> {
 
 #[doc(hidden)]
 pub fn eval<R, F: FnOnce(Option<String>) -> R>(name: &str, f: F) -> Option<R> {
+    let id = thread::current().id();
+    let group = REGISTRY_GROUP.read().unwrap();
+
+    let registry = group
+        .get(&id)
+        .unwrap_or(&REGISTRY_GLOBAL.registry)
+        .read()
+        .unwrap();
+
     let p = {
-        let registry = REGISTRY.registry.read().unwrap();
         match registry.get(name) {
             None => return None,
             Some(p) => p.clone(),
@@ -666,7 +699,14 @@ pub fn eval<R, F: FnOnce(Option<String>) -> R>(name: &str, f: F) -> Option<R> {
 /// A call to `cfg` with a particular fail point name overwrites any existing actions for
 /// that fail point, including those set via the `FAILPOINTS` environment variable.
 pub fn cfg<S: Into<String>>(name: S, actions: &str) -> Result<(), String> {
-    let mut registry = REGISTRY.registry.write().unwrap();
+    let id = thread::current().id();
+    let group = REGISTRY_GROUP.read().unwrap();
+    let mut registry = group
+        .get(&id)
+        .unwrap_or(&REGISTRY_GLOBAL.registry)
+        .write()
+        .unwrap();
+
     set(&mut registry, name.into(), actions)
 }
 
@@ -679,7 +719,14 @@ where
     S: Into<String>,
     F: Fn() + Send + Sync + 'static,
 {
-    let mut registry = REGISTRY.registry.write().unwrap();
+    let id = thread::current().id();
+    let group = REGISTRY_GROUP.read().unwrap();
+    let mut registry = group
+        .get(&id)
+        .unwrap_or(&REGISTRY_GLOBAL.registry)
+        .write()
+        .unwrap();
+
     let p = registry
         .entry(name.into())
         .or_insert_with(|| Arc::new(FailPoint::new()));
@@ -691,9 +738,17 @@ where
 
 /// Remove a fail point.
 ///
-/// If the fail point doesn't exist, nothing will happen.
+/// If the fail point doesn't exist, it will try to delete the corresponding
+/// action in the `REGISTRY_GLOBAL`.
 pub fn remove<S: AsRef<str>>(name: S) {
-    let mut registry = REGISTRY.registry.write().unwrap();
+    let id = thread::current().id();
+    let group = REGISTRY_GROUP.read().unwrap();
+    let mut registry = group
+        .get(&id)
+        .unwrap_or(&REGISTRY_GLOBAL.registry)
+        .write()
+        .unwrap();
+
     if let Some(p) = registry.remove(name.as_ref()) {
         // wake up all pause failpoint.
         p.set_actions("", vec![]);
@@ -990,10 +1045,7 @@ mod tests {
         }
     }
 
-    // This case should be tested as integration case, but when calling `teardown` other cases
-    // like `test_pause` maybe also affected, so it's better keep it here.
     #[test]
-    #[cfg_attr(not(feature = "failpoints"), ignore)]
     fn test_setup_and_teardown() {
         let f1 = || {
             fail_point!("setup_and_teardown1", |_| 1);
@@ -1007,7 +1059,7 @@ mod tests {
             "FAILPOINTS",
             "setup_and_teardown1=return;setup_and_teardown2=pause;",
         );
-        let scenario = FailScenario::setup();
+        setup();
         assert_eq!(f1(), 1);
 
         let (tx, rx) = mpsc::channel();
@@ -1016,7 +1068,7 @@ mod tests {
         });
         assert!(rx.recv_timeout(Duration::from_millis(500)).is_err());
 
-        scenario.teardown();
+        teardown();
         assert_eq!(rx.recv_timeout(Duration::from_millis(500)).unwrap(), 0);
         assert_eq!(f1(), 0);
     }