atlassian_cli_api/
ratelimit.rs

1use chrono::{DateTime, Utc};
2use reqwest::Response;
3use std::sync::Arc;
4use std::time::Duration;
5use tokio::sync::Mutex;
6use tokio::time::timeout;
7use tracing::{debug, warn};
8
9#[derive(Clone)]
10pub struct RateLimiter {
11    state: Arc<Mutex<RateLimitState>>,
12}
13
14struct RateLimitState {
15    remaining: Option<u32>,
16    reset_at: Option<DateTime<Utc>>,
17    limit: Option<u32>,
18}
19
20impl RateLimiter {
21    pub fn new() -> Self {
22        Self {
23            state: Arc::new(Mutex::new(RateLimitState {
24                remaining: None,
25                reset_at: None,
26                limit: None,
27            })),
28        }
29    }
30
31    pub async fn update_from_response(&self, response: &Response) {
32        let Ok(mut state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
33            warn!("Rate limiter lock timeout, skipping update");
34            return;
35        };
36
37        if let Some(limit) = response.headers().get("x-ratelimit-limit") {
38            if let Ok(s) = limit.to_str() {
39                if let Ok(val) = s.parse() {
40                    state.limit = Some(val);
41                }
42            }
43        }
44
45        if let Some(remaining) = response.headers().get("x-ratelimit-remaining") {
46            if let Ok(s) = remaining.to_str() {
47                if let Ok(val) = s.parse() {
48                    state.remaining = Some(val);
49                }
50            }
51        }
52
53        if let Some(reset) = response.headers().get("x-ratelimit-reset") {
54            if let Ok(s) = reset.to_str() {
55                if let Ok(timestamp) = s.parse::<i64>() {
56                    if let Some(dt) = DateTime::from_timestamp(timestamp, 0) {
57                        state.reset_at = Some(dt);
58                    }
59                }
60            }
61        }
62
63        if let (Some(remaining), Some(limit)) = (state.remaining, state.limit) {
64            let usage_pct = ((limit - remaining) as f64 / limit as f64) * 100.0;
65            if usage_pct > 80.0 {
66                warn!(
67                    remaining,
68                    limit,
69                    usage_pct = format!("{:.1}%", usage_pct),
70                    "Rate limit usage high"
71                );
72            }
73        }
74    }
75
76    pub async fn check_limit(&self) -> Option<u64> {
77        let Ok(state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
78            warn!("Rate limiter lock timeout, skipping check");
79            return None;
80        };
81
82        if let Some(remaining) = state.remaining {
83            if remaining == 0 {
84                if let Some(reset_at) = state.reset_at {
85                    let now = Utc::now();
86                    if reset_at > now {
87                        let wait_secs = (reset_at - now).num_seconds() as u64;
88                        debug!(wait_secs, "Rate limit exceeded, need to wait");
89                        return Some(wait_secs);
90                    }
91                }
92            }
93        }
94
95        None
96    }
97
98    pub async fn get_info(&self) -> RateLimitInfo {
99        let Ok(state) = timeout(Duration::from_secs(5), self.state.lock()).await else {
100            warn!("Rate limiter lock timeout, returning empty info");
101            return RateLimitInfo {
102                limit: None,
103                remaining: None,
104                reset_at: None,
105            };
106        };
107        RateLimitInfo {
108            limit: state.limit,
109            remaining: state.remaining,
110            reset_at: state.reset_at,
111        }
112    }
113}
114
115impl Default for RateLimiter {
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121#[derive(Debug, Clone)]
122pub struct RateLimitInfo {
123    pub limit: Option<u32>,
124    pub remaining: Option<u32>,
125    pub reset_at: Option<DateTime<Utc>>,
126}
atlassian_cli_api/ratelimit.rs

atlassian_cli_api/
ratelimit.rs
