dataforge/multithreading/

safety.rs

//! 线程安全机制模块

use std::sync::{Arc, RwLock, Mutex};
use std::collections::HashMap;
use serde_json::Value;
use crate::error::{DataForgeError, Result};

/// 线程安全的数据缓存
pub struct ThreadSafeCache<K, V> {
    data: Arc<RwLock<HashMap<K, V>>>,
    max_size: usize,
}

impl<K, V> ThreadSafeCache<K, V>
where
    K: std::hash::Hash + Eq + Clone,
    V: Clone,
{
    /// 创建新的线程安全缓存
    pub fn new(max_size: usize) -> Self {
        Self {
            data: Arc::new(RwLock::new(HashMap::new())),
            max_size,
        }
    }

    /// 获取值
    pub fn get(&self, key: &K) -> Option<V> {
        self.data.read().unwrap().get(key).cloned()
    }

    /// 插入值
    pub fn insert(&self, key: K, value: V) -> Result<()> {
        let mut data = self.data.write().unwrap();
        
        if data.len() >= self.max_size && !data.contains_key(&key) {
            return Err(DataForgeError::validation("Cache is full"));
        }
        
        data.insert(key, value);
        Ok(())
    }

    /// 移除值
    pub fn remove(&self, key: &K) -> Option<V> {
        self.data.write().unwrap().remove(key)
    }

    /// 清空缓存
    pub fn clear(&self) {
        self.data.write().unwrap().clear();
    }

    /// 获取缓存大小
    pub fn len(&self) -> usize {
        self.data.read().unwrap().len()
    }

    /// 检查是否为空
    pub fn is_empty(&self) -> bool {
        self.data.read().unwrap().is_empty()
    }
}

/// 线程安全的计数器
pub struct ThreadSafeCounter {
    value: Arc<Mutex<usize>>,
}

impl ThreadSafeCounter {
    /// 创建新的计数器
    pub fn new(initial: usize) -> Self {
        Self {
            value: Arc::new(Mutex::new(initial)),
        }
    }

    /// 增加计数
    pub fn increment(&self) -> usize {
        let mut value = self.value.lock().unwrap();
        *value += 1;
        *value
    }

    /// 减少计数
    pub fn decrement(&self) -> usize {
        let mut value = self.value.lock().unwrap();
        if *value > 0 {
            *value -= 1;
        }
        *value
    }

    /// 获取当前值
    pub fn get(&self) -> usize {
        *self.value.lock().unwrap()
    }

    /// 重置计数器
    pub fn reset(&self) {
        *self.value.lock().unwrap() = 0;
    }
}

/// 线程安全的数据生成器注册表
pub struct ThreadSafeGeneratorRegistry {
    generators: Arc<RwLock<HashMap<String, Arc<dyn Fn() -> Value + Send + Sync>>>>,
}

impl ThreadSafeGeneratorRegistry {
    /// 创建新的注册表
    pub fn new() -> Self {
        Self {
            generators: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// 注册生成器
    pub fn register<F>(&self, name: String, generator: F)
    where
        F: Fn() -> Value + Send + Sync + 'static,
    {
        self.generators.write().unwrap().insert(name, Arc::new(generator));
    }

    /// 获取生成器
    pub fn get(&self, name: &str) -> Option<Arc<dyn Fn() -> Value + Send + Sync>> {
        self.generators.read().unwrap().get(name).cloned()
    }

    /// 移除生成器
    pub fn remove(&self, name: &str) -> bool {
        self.generators.write().unwrap().remove(name).is_some()
    }

    /// 列出所有生成器名称
    pub fn list_generators(&self) -> Vec<String> {
        self.generators.read().unwrap().keys().cloned().collect()
    }

    /// 获取生成器数量
    pub fn count(&self) -> usize {
        self.generators.read().unwrap().len()
    }
}

impl Default for ThreadSafeGeneratorRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for ThreadSafeGeneratorRegistry {
    fn clone(&self) -> Self {
        Self {
            generators: Arc::clone(&self.generators),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;


    #[test]
    fn test_thread_safe_cache() {
        let cache = ThreadSafeCache::new(10);
        
        assert!(cache.insert("key1".to_string(), "value1".to_string()).is_ok());
        assert_eq!(cache.get(&"key1".to_string()), Some("value1".to_string()));
        assert_eq!(cache.len(), 1);
        
        cache.clear();
        assert!(cache.is_empty());
    }

    #[test]
    fn test_thread_safe_counter() {
        let counter = ThreadSafeCounter::new(0);
        
        assert_eq!(counter.increment(), 1);
        assert_eq!(counter.increment(), 2);
        assert_eq!(counter.get(), 2);
        
        assert_eq!(counter.decrement(), 1);
        assert_eq!(counter.get(), 1);
        
        counter.reset();
        assert_eq!(counter.get(), 0);
    }

    #[test]
    fn test_thread_safe_generator_registry() {
        let registry = ThreadSafeGeneratorRegistry::new();
        
        registry.register("test".to_string(), || Value::String("test".to_string()));
        
        assert_eq!(registry.count(), 1);
        assert!(registry.get("test").is_some());
        
        let generator = registry.get("test").unwrap();
        let result = generator();
        assert_eq!(result, Value::String("test".to_string()));
        
        assert!(registry.remove("test"));
        assert_eq!(registry.count(), 0);
    }

    #[test]
    fn test_concurrent_access() {
        let counter = Arc::new(ThreadSafeCounter::new(0));
        let mut handles = vec![];

        for _ in 0..10 {
            let counter = Arc::clone(&counter);
            let handle = thread::spawn(move || {
                for _ in 0..100 {
                    counter.increment();
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().unwrap();
        }

        assert_eq!(counter.get(), 1000);
    }
}