rustfs/crates/concurrency/src/timeout.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Timeout management for operations

use rustfs_io_core::{TimeoutConfig as CoreTimeoutConfig, TimeoutError, calculate_adaptive_timeout};
use std::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;

/// Facade policy for the concurrency-layer timeout manager.
#[derive(Debug, Clone, Copy)]
pub struct TimeoutManagerPolicy {
    /// Default timeout duration
    pub default_timeout: Duration,
    /// Maximum timeout duration
    pub max_timeout: Duration,
    /// Minimum timeout floor (prevents dynamic calculation from going too low).
    pub min_timeout: Duration,
    /// Enable dynamic timeout calculation
    pub enable_dynamic: bool,
}

impl Default for TimeoutManagerPolicy {
    fn default() -> Self {
        Self {
            default_timeout: Duration::from_secs(30),
            max_timeout: Duration::from_secs(300),
            min_timeout: Duration::from_secs(5),
            enable_dynamic: true,
        }
    }
}

impl TimeoutManagerPolicy {
    /// Convert the facade policy into the reusable io-core timeout configuration.
    ///
    /// This keeps the concurrency layer explicitly wired to the shared core
    /// timeout primitives without changing the facade's public behavior.
    pub fn to_core_config(&self) -> CoreTimeoutConfig {
        CoreTimeoutConfig {
            base_timeout: self.default_timeout,
            timeout_per_mb: Duration::ZERO,
            max_timeout: self.max_timeout,
            min_timeout: self.min_timeout,
            get_object_timeout: self.default_timeout,
            put_object_timeout: self.max_timeout,
            list_objects_timeout: self.default_timeout,
            enable_dynamic_timeout: self.enable_dynamic,
        }
    }
}

/// Timeout manager
pub struct TimeoutManager {
    config: TimeoutManagerPolicy,
    core_config: CoreTimeoutConfig,
}

impl TimeoutManager {
    /// Create a new timeout manager
    pub fn new(default_timeout: Duration, max_timeout: Duration, enable_dynamic: bool) -> Self {
        let min_timeout = default_timeout.min(max_timeout);
        Self::from_policy(TimeoutManagerPolicy {
            default_timeout,
            max_timeout,
            min_timeout,
            enable_dynamic,
        })
    }

    /// Create a new timeout manager from the facade policy type.
    pub fn from_policy(config: TimeoutManagerPolicy) -> Self {
        let config = TimeoutManagerPolicy {
            // Guard clamp(min, max) from panic when callers provide an
            // out-of-order policy (or very small max_timeout).
            min_timeout: config.min_timeout.min(config.max_timeout),
            ..config
        };
        let core_config = config.to_core_config();
        Self { config, core_config }
    }

    /// Get the configuration
    pub fn config(&self) -> &TimeoutManagerPolicy {
        &self.config
    }

    /// Get the derived io-core timeout configuration.
    pub fn core_config(&self) -> &CoreTimeoutConfig {
        &self.core_config
    }

    /// Calculate timeout for a given size
    pub fn calculate_timeout(&self, size: u64, _history: &[Duration]) -> Duration {
        if !self.config.enable_dynamic {
            return self.config.default_timeout;
        }

        calculate_adaptive_timeout(self.core_config.base_timeout, None, 0, size)
            .clamp(self.core_config.min_timeout, self.core_config.max_timeout)
    }

    /// Wrap an operation with timeout control
    pub async fn wrap_operation<F, T, E>(&self, operation: F, timeout: Option<Duration>) -> Result<T, TimeoutError>
    where
        F: std::future::Future<Output = Result<T, E>>,
        E: Into<TimeoutError>,
    {
        let timeout = timeout.unwrap_or(self.config.default_timeout);

        match tokio::time::timeout(timeout, operation).await {
            Ok(Ok(result)) => Ok(result),
            Ok(Err(e)) => Err(e.into()),
            Err(_) => Err(TimeoutError::TimedOut(timeout)),
        }
    }

    /// Create a timeout guard for manual timeout control
    pub fn create_guard(&self, timeout: Option<Duration>) -> TimeoutGuard {
        TimeoutGuard::new(timeout.unwrap_or(self.core_config.base_timeout))
    }
}

/// Timeout guard for manual timeout control
pub struct TimeoutGuard {
    timeout: Duration,
    start: Instant,
    cancel_token: CancellationToken,
}

impl TimeoutGuard {
    fn new(timeout: Duration) -> Self {
        Self {
            timeout,
            start: Instant::now(),
            cancel_token: CancellationToken::new(),
        }
    }

    /// Check if timeout has elapsed
    pub fn is_timed_out(&self) -> bool {
        self.start.elapsed() > self.timeout
    }

    /// Get remaining time
    pub fn remaining(&self) -> Duration {
        self.timeout.saturating_sub(self.start.elapsed())
    }

    /// Get the cancellation token
    pub fn cancel_token(&self) -> CancellationToken {
        self.cancel_token.clone()
    }

    /// Cancel the operation
    pub fn cancel(&self) {
        self.cancel_token.cancel();
    }
}

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

    #[test]
    fn test_timeout_config() {
        let config = TimeoutManagerPolicy::default();
        assert!(config.default_timeout < config.max_timeout);
    }

    #[test]
    fn test_timeout_policy_to_core_config() {
        let policy = TimeoutManagerPolicy::default();
        let core = policy.to_core_config();
        assert_eq!(core.base_timeout, policy.default_timeout);
        assert_eq!(core.max_timeout, policy.max_timeout);
        assert_eq!(core.min_timeout, policy.min_timeout);
        assert_eq!(core.get_object_timeout, policy.default_timeout);
        assert!(core.enable_dynamic_timeout);
    }

    #[test]
    fn test_timeout_manager_new_sanitizes_min_timeout_with_small_max_timeout() {
        let manager = TimeoutManager::new(Duration::from_secs(1), Duration::from_secs(1), true);
        let timeout = manager.calculate_timeout(1024, &[]);
        assert_eq!(timeout, Duration::from_secs(1));
    }

    #[test]
    fn test_timeout_manager_from_policy_sanitizes_min_timeout() {
        let manager = TimeoutManager::from_policy(TimeoutManagerPolicy {
            default_timeout: Duration::from_secs(30),
            max_timeout: Duration::from_secs(1),
            min_timeout: Duration::from_secs(5),
            enable_dynamic: true,
        });

        assert_eq!(manager.config().min_timeout, Duration::from_secs(1));
        assert_eq!(manager.core_config().min_timeout, Duration::from_secs(1));
    }

    #[tokio::test]
    async fn test_wrap_operation_success() {
        let manager = TimeoutManager::new(Duration::from_secs(5), Duration::from_secs(10), true);

        let result = manager.wrap_operation(async { Ok::<_, TimeoutError>(42) }, None).await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 42);
    }
}