atlassian_cli_api/
ratelimit.rs

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

atlassian_cli_api/
ratelimit.rs
