atlassian-cli-api 0.4.2

API client library for Atlassian Cloud products
Documentation
use chrono::{DateTime, Utc};
use reqwest::Response;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::timeout;
use tracing::{debug, warn};

#[derive(Clone)]
pub struct RateLimiter {
    state: Arc<Mutex<RateLimitState>>,
}

struct RateLimitState {
    remaining: Option<u32>,
    reset_at: Option<DateTime<Utc>>,
    limit: Option<u32>,
}

impl RateLimiter {
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(RateLimitState {
                remaining: None,
                reset_at: None,
                limit: None,
            })),
        }
    }

    pub async fn update_from_response(&self, response: &Response) {
        let Ok(mut state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
            warn!("Rate limiter lock timeout, skipping update");
            return;
        };

        if let Some(limit) = response.headers().get("x-ratelimit-limit") {
            if let Ok(s) = limit.to_str() {
                if let Ok(val) = s.parse() {
                    state.limit = Some(val);
                }
            }
        }

        if let Some(remaining) = response.headers().get("x-ratelimit-remaining") {
            if let Ok(s) = remaining.to_str() {
                if let Ok(val) = s.parse() {
                    state.remaining = Some(val);
                }
            }
        }

        if let Some(reset) = response.headers().get("x-ratelimit-reset") {
            if let Ok(s) = reset.to_str() {
                if let Ok(timestamp) = s.parse::<i64>() {
                    if let Some(dt) = DateTime::from_timestamp(timestamp, 0) {
                        state.reset_at = Some(dt);
                    }
                }
            }
        }

        if let (Some(remaining), Some(limit)) = (state.remaining, state.limit) {
            let usage_pct = ((limit - remaining) as f64 / limit as f64) * 100.0;
            if usage_pct > 80.0 {
                warn!(
                    remaining,
                    limit,
                    usage_pct = format!("{:.1}%", usage_pct),
                    "Rate limit usage high"
                );
            }
        }
    }

    pub async fn check_limit(&self) -> Option<u64> {
        let Ok(state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
            warn!("Rate limiter lock timeout, skipping check");
            return None;
        };

        if let Some(remaining) = state.remaining {
            if remaining == 0 {
                if let Some(reset_at) = state.reset_at {
                    let now = Utc::now();
                    if reset_at > now {
                        let wait_secs = (reset_at - now).num_seconds() as u64;
                        debug!(wait_secs, "Rate limit exceeded, need to wait");
                        return Some(wait_secs);
                    }
                }
            }
        }

        None
    }

    pub async fn get_info(&self) -> RateLimitInfo {
        let Ok(state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
            warn!("Rate limiter lock timeout, returning empty info");
            return RateLimitInfo {
                limit: None,
                remaining: None,
                reset_at: None,
            };
        };
        RateLimitInfo {
            limit: state.limit,
            remaining: state.remaining,
            reset_at: state.reset_at,
        }
    }
}

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

#[derive(Debug, Clone)]
pub struct RateLimitInfo {
    pub limit: Option<u32>,
    pub remaining: Option<u32>,
    pub reset_at: Option<DateTime<Utc>>,
}