ostium_rust_sdk/

retry.rs

use crate::error::OstiumError;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::{sleep, Instant};
use tracing::{debug, warn};

/// Configuration for retry behavior
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retry attempts
    pub max_attempts: u32,
    /// Initial delay between retries
    pub initial_delay: Duration,
    /// Maximum delay between retries
    pub max_delay: Duration,
    /// Multiplier for exponential backoff
    pub backoff_multiplier: f64,
    /// Maximum jitter factor (0.0 to 1.0)
    pub jitter_factor: f64,
    /// Timeout for individual operations
    pub operation_timeout: Duration,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(30),
            backoff_multiplier: 2.0,
            jitter_factor: 0.1,
            operation_timeout: Duration::from_secs(30),
        }
    }
}

impl RetryConfig {
    /// Create a config optimized for network operations
    pub fn network() -> Self {
        Self {
            max_attempts: 5,
            initial_delay: Duration::from_millis(200),
            max_delay: Duration::from_secs(10),
            backoff_multiplier: 1.5,
            jitter_factor: 0.2,
            operation_timeout: Duration::from_secs(30),
        }
    }

    /// Create a config optimized for contract interactions
    pub fn contract() -> Self {
        Self {
            max_attempts: 3,
            initial_delay: Duration::from_millis(500),
            max_delay: Duration::from_secs(20),
            backoff_multiplier: 2.0,
            jitter_factor: 0.1,
            operation_timeout: Duration::from_secs(60),
        }
    }

    /// Create a config optimized for GraphQL queries
    pub fn graphql() -> Self {
        Self {
            max_attempts: 4,
            initial_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(5),
            backoff_multiplier: 1.8,
            jitter_factor: 0.15,
            operation_timeout: Duration::from_secs(15),
        }
    }
}

/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CircuitState {
    /// Circuit is closed - normal operation
    Closed,
    /// Circuit is open - blocking requests due to failures
    Open,
    /// Circuit is half-open - allowing limited requests to test recovery
    HalfOpen,
}

/// Circuit breaker for preventing cascading failures
#[derive(Debug)]
pub struct CircuitBreaker {
    state: Arc<AtomicU64>, // Packed: state (8 bits) + failure_count (24 bits) + last_failure_time (32 bits)
    failure_threshold: u32,
    recovery_timeout: Duration,
    success_threshold: u32,
}

impl CircuitBreaker {
    /// Create a new circuit breaker with the specified failure threshold and recovery timeout
    pub fn new(failure_threshold: u32, recovery_timeout: Duration) -> Self {
        Self {
            state: Arc::new(AtomicU64::new(0)), // Initial state: Closed (0)
            failure_threshold,
            recovery_timeout,
            success_threshold: 3,
        }
    }

    /// Execute an operation through the circuit breaker
    pub fn call<F, Fut, T>(
        &self,
        operation: F,
    ) -> impl std::future::Future<Output = Result<T, OstiumError>>
    where
        F: FnOnce() -> Fut,
        Fut: std::future::Future<Output = Result<T, OstiumError>>,
    {
        let state = self.state.clone();
        let failure_threshold = self.failure_threshold;
        let recovery_timeout = self.recovery_timeout;
        let _success_threshold = self.success_threshold;

        async move {
            let current_state = Self::decode_state(state.load(Ordering::Acquire));

            match current_state.0 {
                CircuitState::Open => {
                    let time_since_failure = Instant::now().duration_since(
                        Instant::now() - Duration::from_secs(current_state.2 as u64),
                    );

                    if time_since_failure >= recovery_timeout {
                        // Transition to half-open
                        let new_packed = Self::encode_state(CircuitState::HalfOpen, 0, 0);
                        state.store(new_packed, Ordering::Release);
                        debug!("Circuit breaker transitioning to half-open");
                    } else {
                        return Err(OstiumError::Network(
                            "Circuit breaker is open - too many recent failures".to_string(),
                        ));
                    }
                }
                CircuitState::HalfOpen => {
                    // Allow limited requests through
                }
                CircuitState::Closed => {
                    // Normal operation
                }
            }

            match operation().await {
                Ok(result) => {
                    // Success - reset failure count or close circuit
                    match current_state.0 {
                        CircuitState::HalfOpen => {
                            let new_packed = Self::encode_state(CircuitState::Closed, 0, 0);
                            state.store(new_packed, Ordering::Release);
                            debug!("Circuit breaker closed after successful recovery");
                        }
                        _ => {
                            let new_packed = Self::encode_state(CircuitState::Closed, 0, 0);
                            state.store(new_packed, Ordering::Release);
                        }
                    }
                    Ok(result)
                }
                Err(error) => {
                    // Failure - increment count and possibly open circuit
                    let new_failure_count = current_state.1 + 1;
                    let current_time = Instant::now().elapsed().as_secs() as u32;

                    if new_failure_count >= failure_threshold {
                        let new_packed =
                            Self::encode_state(CircuitState::Open, new_failure_count, current_time);
                        state.store(new_packed, Ordering::Release);
                        warn!(
                            "Circuit breaker opened after {} failures",
                            new_failure_count
                        );
                    } else {
                        let new_packed =
                            Self::encode_state(current_state.0, new_failure_count, current_time);
                        state.store(new_packed, Ordering::Release);
                    }

                    Err(error)
                }
            }
        }
    }

    fn encode_state(state: CircuitState, failure_count: u32, last_failure_time: u32) -> u64 {
        let state_bits = match state {
            CircuitState::Closed => 0u64,
            CircuitState::Open => 1u64,
            CircuitState::HalfOpen => 2u64,
        };

        (state_bits << 56) | ((failure_count as u64 & 0xFFFFFF) << 32) | (last_failure_time as u64)
    }

    fn decode_state(packed: u64) -> (CircuitState, u32, u32) {
        let state = match (packed >> 56) & 0xFF {
            0 => CircuitState::Closed,
            1 => CircuitState::Open,
            2 => CircuitState::HalfOpen,
            _ => CircuitState::Closed,
        };
        let failure_count = ((packed >> 32) & 0xFFFFFF) as u32;
        let last_failure_time = (packed & 0xFFFFFFFF) as u32;

        (state, failure_count, last_failure_time)
    }

    /// Get the current state of the circuit breaker
    pub fn state(&self) -> CircuitState {
        Self::decode_state(self.state.load(Ordering::Acquire)).0
    }
}

/// Retry executor with exponential backoff and jitter
pub struct RetryExecutor {
    config: RetryConfig,
    circuit_breaker: Option<CircuitBreaker>,
}

impl RetryExecutor {
    /// Create a new retry executor with the specified configuration
    pub fn new(config: RetryConfig) -> Self {
        Self {
            config,
            circuit_breaker: None,
        }
    }

    /// Add a circuit breaker to the retry executor
    pub fn with_circuit_breaker(
        mut self,
        failure_threshold: u32,
        recovery_timeout: Duration,
    ) -> Self {
        self.circuit_breaker = Some(CircuitBreaker::new(failure_threshold, recovery_timeout));
        self
    }

    /// Execute an operation with retry logic
    pub async fn execute<F, Fut, T>(&self, operation: F) -> Result<T, OstiumError>
    where
        F: Fn() -> Fut,
        Fut: std::future::Future<Output = Result<T, OstiumError>>,
    {
        let mut attempt = 0;
        let mut delay = self.config.initial_delay;

        loop {
            attempt += 1;

            debug!(
                "Executing operation attempt {}/{}",
                attempt, self.config.max_attempts
            );

            // Use circuit breaker if configured
            let result = if let Some(ref circuit_breaker) = self.circuit_breaker {
                circuit_breaker.call(&operation).await
            } else {
                operation().await
            };

            match result {
                Ok(value) => {
                    if attempt > 1 {
                        debug!("Operation succeeded after {} attempts", attempt);
                    }
                    return Ok(value);
                }
                Err(error) => {
                    if !self.should_retry(&error) || attempt >= self.config.max_attempts {
                        warn!("Operation failed after {} attempts: {}", attempt, error);
                        return Err(error);
                    }

                    debug!(
                        "Operation failed on attempt {}, retrying after {:?}: {}",
                        attempt, delay, error
                    );

                    // Sleep with jitter
                    let jittered_delay = self.add_jitter(delay);
                    sleep(jittered_delay).await;

                    // Calculate next delay with exponential backoff
                    delay = std::cmp::min(
                        Duration::from_millis(
                            (delay.as_millis() as f64 * self.config.backoff_multiplier) as u64,
                        ),
                        self.config.max_delay,
                    );
                }
            }
        }
    }

    /// Determine if an error should trigger a retry
    fn should_retry(&self, error: &OstiumError) -> bool {
        match error {
            // Always retry network errors
            OstiumError::Network(_) => true,

            // Retry HTTP errors that might be transient
            OstiumError::Http(e) => e.is_timeout() || e.is_connect() || e.is_request(),

            // Retry specific contract errors
            OstiumError::Contract(msg) => {
                msg.contains("timeout")
                    || msg.contains("connection")
                    || msg.contains("temporarily unavailable")
                    || msg.contains("rate limit")
            }

            // Retry GraphQL errors that might be transient
            OstiumError::GraphQL(msg) => {
                msg.contains("timeout")
                    || msg.contains("server error")
                    || msg.contains("503")
                    || msg.contains("502")
                    || msg.contains("504")
            }

            // Retry provider errors that might be transient
            OstiumError::Provider(msg) => {
                msg.contains("timeout") || msg.contains("connection") || msg.contains("rate limit")
            }

            // Don't retry these errors as they're likely permanent
            OstiumError::Validation(_) => false,
            OstiumError::Wallet(_) => false,
            OstiumError::Config(_) => false,
            OstiumError::Json(_) => false,
            OstiumError::Decimal(_) => false,
            OstiumError::Other(_) => false,
        }
    }

    /// Add jitter to delay to prevent thundering herd
    fn add_jitter(&self, delay: Duration) -> Duration {
        if self.config.jitter_factor <= 0.0 {
            return delay;
        }

        let jitter_range = (delay.as_millis() as f64 * self.config.jitter_factor) as u64;
        let jitter = fastrand::u64(0..=jitter_range);

        Duration::from_millis(delay.as_millis() as u64 + jitter)
    }
}

/// Convenience macros for different operation types
/// Execute an operation with network-optimized retry settings
#[macro_export]
macro_rules! retry_network {
    ($operation:expr) => {
        $crate::retry::RetryExecutor::new($crate::retry::RetryConfig::network())
            .execute(|| async { $operation })
            .await
    };
}

/// Execute an operation with contract-optimized retry settings and circuit breaker
#[macro_export]
macro_rules! retry_contract {
    ($operation:expr) => {
        $crate::retry::RetryExecutor::new($crate::retry::RetryConfig::contract())
            .with_circuit_breaker(5, std::time::Duration::from_secs(60))
            .execute(|| async { $operation })
            .await
    };
}

/// Execute an operation with GraphQL-optimized retry settings
#[macro_export]
macro_rules! retry_graphql {
    ($operation:expr) => {
        $crate::retry::RetryExecutor::new($crate::retry::RetryConfig::graphql())
            .execute(|| async { $operation })
            .await
    };
}

// Re-export macros for convenience
pub use retry_contract;
pub use retry_graphql;
pub use retry_network;

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::sync::Arc;

    #[tokio::test]
    async fn test_retry_success_after_failures() {
        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        let config = RetryConfig {
            max_attempts: 3,
            initial_delay: Duration::from_millis(10),
            ..Default::default()
        };

        let executor = RetryExecutor::new(config);

        let result = executor
            .execute(|| {
                let counter = counter_clone.clone();
                async move {
                    let count = counter.fetch_add(1, Ordering::SeqCst);
                    if count < 2 {
                        Err(OstiumError::Network("Temporary failure".to_string()))
                    } else {
                        Ok("Success".to_string())
                    }
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "Success");
        assert_eq!(counter.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_retry_exhaustion() {
        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        let config = RetryConfig {
            max_attempts: 2,
            initial_delay: Duration::from_millis(10),
            ..Default::default()
        };

        let executor = RetryExecutor::new(config);

        let result: Result<String, OstiumError> = executor
            .execute(|| {
                let counter = counter_clone.clone();
                async move {
                    counter.fetch_add(1, Ordering::SeqCst);
                    Err(OstiumError::Network("Permanent failure".to_string()))
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(counter.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn test_circuit_breaker() {
        let circuit_breaker = CircuitBreaker::new(2, Duration::from_millis(100));

        // First failure
        let result1: Result<String, OstiumError> = circuit_breaker
            .call(|| async { Err(OstiumError::Network("Failure".to_string())) })
            .await;
        assert!(result1.is_err());
        assert_eq!(circuit_breaker.state(), CircuitState::Closed);

        // Second failure - should open circuit
        let result2: Result<String, OstiumError> = circuit_breaker
            .call(|| async { Err(OstiumError::Network("Failure".to_string())) })
            .await;
        assert!(result2.is_err());
        assert_eq!(circuit_breaker.state(), CircuitState::Open);

        // Third call should be rejected immediately
        let result3 = circuit_breaker
            .call(|| async { Ok("Should not execute".to_string()) })
            .await;
        assert!(result3.is_err());
        assert!(result3
            .unwrap_err()
            .to_string()
            .contains("Circuit breaker is open"));
    }
}