revoke_resilience/

timeout.rs

use crate::error::{ResilienceError, Result};
use std::time::Duration;
use tokio::time::timeout as tokio_timeout;
use tracing::{debug, warn};

/// 超时处理器
#[derive(Debug, Clone)]
pub struct Timeout {
    duration: Duration,
    name: Option<String>,
}

impl Timeout {
    /// 创建新的超时处理器
    pub fn new(duration: Duration) -> Self {
        Self {
            duration,
            name: None,
        }
    }

    /// 设置名称
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// 执行带超时的操作
    pub async fn execute<F, T>(&self, future: F) -> Result<T>
    where
        F: std::future::Future<Output = Result<T>>,
    {
        let name = self.name.as_deref().unwrap_or("未命名");
        
        debug!("开始执行操作 '{}', 超时: {:?}", name, self.duration);
        
        match tokio_timeout(self.duration, future).await {
            Ok(result) => {
                debug!("操作 '{}' 在超时前完成", name);
                result
            }
            Err(_) => {
                warn!("操作 '{}' 超时 ({:?})", name, self.duration);
                Err(ResilienceError::Timeout(format!(
                    "操作 '{}' 在 {:?} 后超时",
                    name, self.duration
                )))
            }
        }
    }

    /// 执行带超时的操作，返回 Option
    pub async fn execute_optional<F, T>(&self, future: F) -> Option<T>
    where
        F: std::future::Future<Output = T>,
    {
        let name = self.name.as_deref().unwrap_or("未命名");
        
        debug!("开始执行操作 '{}', 超时: {:?}", name, self.duration);
        
        match tokio_timeout(self.duration, future).await {
            Ok(result) => {
                debug!("操作 '{}' 在超时前完成", name);
                Some(result)
            }
            Err(_) => {
                warn!("操作 '{}' 超时 ({:?})", name, self.duration);
                None
            }
        }
    }
}

/// 快速超时函数
pub async fn with_timeout<F, T>(duration: Duration, future: F) -> Result<T>
where
    F: std::future::Future<Output = Result<T>>,
{
    Timeout::new(duration).execute(future).await
}

/// 快速超时函数（返回 Option）
pub async fn with_timeout_optional<F, T>(duration: Duration, future: F) -> Option<T>
where
    F: std::future::Future<Output = T>,
{
    Timeout::new(duration).execute_optional(future).await
}

/// 动态超时管理器
#[derive(Debug)]
pub struct DynamicTimeout {
    base_duration: Duration,
    min_duration: Duration,
    max_duration: Duration,
    adjustment_factor: f64,
    current_duration: parking_lot::RwLock<Duration>,
    success_count: std::sync::atomic::AtomicU32,
    timeout_count: std::sync::atomic::AtomicU32,
}

impl DynamicTimeout {
    /// 创建新的动态超时管理器
    pub fn new(base_duration: Duration) -> Self {
        Self {
            base_duration,
            min_duration: Duration::from_millis(100),
            max_duration: Duration::from_secs(300),
            adjustment_factor: 1.5,
            current_duration: parking_lot::RwLock::new(base_duration),
            success_count: std::sync::atomic::AtomicU32::new(0),
            timeout_count: std::sync::atomic::AtomicU32::new(0),
        }
    }

    /// 设置最小超时
    pub fn with_min_duration(mut self, duration: Duration) -> Self {
        self.min_duration = duration;
        self
    }

    /// 设置最大超时
    pub fn with_max_duration(mut self, duration: Duration) -> Self {
        self.max_duration = duration;
        self
    }

    /// 设置调整因子
    pub fn with_adjustment_factor(mut self, factor: f64) -> Self {
        self.adjustment_factor = factor;
        self
    }

    /// 获取当前超时时间
    pub fn current_duration(&self) -> Duration {
        *self.current_duration.read()
    }

    /// 执行带动态超时的操作
    pub async fn execute<F, T>(&self, future: F) -> Result<T>
    where
        F: std::future::Future<Output = Result<T>>,
    {
        let current = self.current_duration();
        let timeout = Timeout::new(current);
        
        match timeout.execute(future).await {
            Ok(result) => {
                self.record_success();
                Ok(result)
            }
            Err(ResilienceError::Timeout(_)) => {
                self.record_timeout();
                Err(ResilienceError::Timeout(format!(
                    "动态超时: {:?}",
                    current
                )))
            }
            Err(e) => Err(e),
        }
    }

    /// 记录成功
    fn record_success(&self) {
        use std::sync::atomic::Ordering;
        
        let success_count = self.success_count.fetch_add(1, Ordering::Relaxed);
        let timeout_count = self.timeout_count.load(Ordering::Relaxed);
        
        // 每 10 次成功后调整
        if success_count % 10 == 0 && timeout_count == 0 {
            let mut current = self.current_duration.write();
            let new_duration = current.as_millis() as f64 / self.adjustment_factor;
            let new_duration = Duration::from_millis(new_duration as u64).max(self.min_duration);
            
            if new_duration != *current {
                debug!("减少超时时间: {:?} -> {:?}", current, new_duration);
                *current = new_duration;
                self.timeout_count.store(0, Ordering::Relaxed);
            }
        }
    }

    /// 记录超时
    fn record_timeout(&self) {
        use std::sync::atomic::Ordering;
        
        let timeout_count = self.timeout_count.fetch_add(1, Ordering::Relaxed);
        
        // 连续 3 次超时后调整
        if timeout_count >= 2 {
            let mut current = self.current_duration.write();
            let new_duration = current.as_millis() as f64 * self.adjustment_factor;
            let new_duration = Duration::from_millis(new_duration as u64).min(self.max_duration);
            
            if new_duration != *current {
                debug!("增加超时时间: {:?} -> {:?}", current, new_duration);
                *current = new_duration;
                self.timeout_count.store(0, Ordering::Relaxed);
                self.success_count.store(0, Ordering::Relaxed);
            }
        }
    }

    /// 重置到基础超时
    pub fn reset(&self) {
        use std::sync::atomic::Ordering;
        
        *self.current_duration.write() = self.base_duration;
        self.success_count.store(0, Ordering::Relaxed);
        self.timeout_count.store(0, Ordering::Relaxed);
        debug!("超时重置到基础值: {:?}", self.base_duration);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::time::sleep;

    #[tokio::test]
    async fn test_timeout_success() {
        let timeout = Timeout::new(Duration::from_millis(100));
        
        let result = timeout.execute(async {
            sleep(Duration::from_millis(50)).await;
            Ok(42)
        }).await;
        
        assert_eq!(result.unwrap(), 42);
    }

    #[tokio::test]
    async fn test_timeout_failure() {
        let timeout = Timeout::new(Duration::from_millis(50));
        
        let result = timeout.execute(async {
            sleep(Duration::from_millis(100)).await;
            Ok(42)
        }).await;
        
        assert!(matches!(result, Err(ResilienceError::Timeout(_))));
    }

    #[tokio::test]
    async fn test_with_timeout() {
        let result = with_timeout(Duration::from_millis(100), async {
            sleep(Duration::from_millis(50)).await;
            Ok("success")
        }).await;
        
        assert_eq!(result.unwrap(), "success");
    }

    #[tokio::test]
    async fn test_with_timeout_optional() {
        let result = with_timeout_optional(Duration::from_millis(100), async {
            sleep(Duration::from_millis(50)).await;
            42
        }).await;
        
        assert_eq!(result, Some(42));
        
        let result = with_timeout_optional(Duration::from_millis(50), async {
            sleep(Duration::from_millis(100)).await;
            42
        }).await;
        
        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_dynamic_timeout() {
        let dynamic = DynamicTimeout::new(Duration::from_millis(100))
            .with_min_duration(Duration::from_millis(50))
            .with_max_duration(Duration::from_millis(200))
            .with_adjustment_factor(2.0);
        
        // 初始超时应该是 100ms
        assert_eq!(dynamic.current_duration(), Duration::from_millis(100));
        
        // 记录几次超时
        for _ in 0..3 {
            let _ = dynamic.execute(async {
                sleep(Duration::from_millis(150)).await;
                Ok(())
            }).await;
        }
        
        // 超时时间应该增加
        assert!(dynamic.current_duration() > Duration::from_millis(100));
        
        // 重置
        dynamic.reset();
        assert_eq!(dynamic.current_duration(), Duration::from_millis(100));
    }
}