github_bot_sdk/client/

rate_limit.rs

//! Rate limit tracking for GitHub API operations.
//!
//! GitHub enforces rate limits on API requests. This module provides types and functions
//! for tracking rate limits from response headers and checking them before making requests.

use chrono::{DateTime, Utc};
use reqwest::header::HeaderMap;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use crate::auth::InstallationId;

/// Authentication context for rate limit tracking.
///
/// GitHub enforces separate rate limits for app-level and installation-level operations.
/// This enum distinguishes between the two contexts.
///
/// # Examples
///
/// ```
/// use github_bot_sdk::client::RateLimitContext;
/// use github_bot_sdk::auth::InstallationId;
///
/// // App-level context (using JWT)
/// let app_context = RateLimitContext::App;
///
/// // Installation-level context (using installation token)
/// let install_context = RateLimitContext::Installation(InstallationId::new(12345));
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RateLimitContext {
    /// App-level operations (authenticated with JWT)
    App,
    /// Installation-level operations (authenticated with installation token)
    Installation(InstallationId),
}

/// Rate limit information from GitHub API response headers.
///
/// GitHub includes rate limit information in HTTP response headers:
/// - `X-RateLimit-Limit`: Maximum requests allowed per hour
/// - `X-RateLimit-Remaining`: Requests remaining in current window
/// - `X-RateLimit-Reset`: Unix timestamp when the rate limit resets
///
/// # Examples
///
/// ```
/// use github_bot_sdk::client::RateLimit;
/// use chrono::{Utc, Duration};
///
/// let reset_time = Utc::now() + Duration::hours(1);
/// let rate_limit = RateLimit::new(5000, 4500, reset_time, "core");
///
/// assert!(!rate_limit.is_exhausted());
/// assert!(rate_limit.remaining() > 1000);
/// ```
#[derive(Debug, Clone)]
pub struct RateLimit {
    /// Maximum requests allowed per hour
    limit: u32,
    /// Requests remaining in current window
    remaining: u32,
    /// When the rate limit resets
    reset_at: DateTime<Utc>,
    /// The resource this rate limit applies to (e.g., "core", "search")
    resource: String,
}

impl RateLimit {
    /// Create a new rate limit from GitHub API response.
    ///
    /// # Arguments
    ///
    /// * `limit` - Maximum requests allowed
    /// * `remaining` - Requests remaining
    /// * `reset_at` - When the limit resets
    /// * `resource` - The resource type (default "core")
    pub fn new(
        limit: u32,
        remaining: u32,
        reset_at: DateTime<Utc>,
        resource: impl Into<String>,
    ) -> Self {
        Self {
            limit,
            remaining,
            reset_at,
            resource: resource.into(),
        }
    }

    /// Get the maximum number of requests allowed.
    pub fn limit(&self) -> u32 {
        self.limit
    }

    /// Get the number of requests remaining.
    pub fn remaining(&self) -> u32 {
        self.remaining
    }

    /// Get when the rate limit resets.
    pub fn reset_at(&self) -> DateTime<Utc> {
        self.reset_at
    }

    /// Get the resource this rate limit applies to.
    pub fn resource(&self) -> &str {
        &self.resource
    }

    /// Check if the rate limit is exhausted (no requests remaining).
    pub fn is_exhausted(&self) -> bool {
        self.remaining == 0
    }

    /// Check if we're close to exhausting the rate limit.
    ///
    /// # Arguments
    ///
    /// * `margin` - The safety margin as a fraction (0.0 to 1.0)
    ///
    /// Returns true if remaining requests are below the margin threshold.
    ///
    /// # Examples
    ///
    /// ```
    /// # use github_bot_sdk::client::RateLimit;
    /// # use chrono::Utc;
    /// let rate_limit = RateLimit::new(5000, 400, Utc::now(), "core");
    ///
    /// // Check if we're below 10% remaining
    /// assert!(rate_limit.is_near_exhaustion(0.1));
    /// ```
    pub fn is_near_exhaustion(&self, margin: f64) -> bool {
        let threshold = (self.limit as f64 * margin) as u32;
        self.remaining <= threshold
    }

    /// Check if the rate limit has been reset.
    ///
    /// Returns true if the current time is past the reset time.
    pub fn has_reset(&self) -> bool {
        Utc::now() >= self.reset_at
    }
}

/// Parse rate limit information from HTTP response headers.
///
/// Extracts rate limit data from GitHub API response headers:
/// - `X-RateLimit-Limit`
/// - `X-RateLimit-Remaining`
/// - `X-RateLimit-Reset`
/// - `X-RateLimit-Resource` (optional, defaults to "core")
///
/// # Arguments
///
/// * `headers` - HTTP response headers from GitHub API
///
/// # Returns
///
/// `Some(RateLimit)` if all required headers are present and valid,
/// `None` if headers are missing or invalid.
///
/// # Examples
///
/// ```no_run
/// # use github_bot_sdk::client::parse_rate_limit_from_headers;
/// # use reqwest::header::HeaderMap;
/// # fn example(headers: &HeaderMap) {
/// if let Some(rate_limit) = parse_rate_limit_from_headers(headers) {
///     println!("Remaining: {}", rate_limit.remaining());
/// }
/// # }
/// ```
pub fn parse_rate_limit_from_headers(headers: &HeaderMap) -> Option<RateLimit> {
    // Extract required headers
    let limit_str = headers.get("x-ratelimit-limit")?.to_str().ok()?;
    let remaining_str = headers.get("x-ratelimit-remaining")?.to_str().ok()?;
    let reset_str = headers.get("x-ratelimit-reset")?.to_str().ok()?;

    // Parse numeric values
    let limit = limit_str.parse::<u32>().ok()?;
    let remaining = remaining_str.parse::<u32>().ok()?;
    let reset_timestamp = reset_str.parse::<i64>().ok()?;

    // Convert Unix timestamp to DateTime
    let reset_at = DateTime::from_timestamp(reset_timestamp, 0)?;

    // Get resource type (defaults to "core" if not present)
    let resource = headers
        .get("x-ratelimit-resource")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("core")
        .to_string();

    Some(RateLimit::new(limit, remaining, reset_at, resource))
}

/// Thread-safe rate limit tracker for GitHub API operations.
///
/// Tracks rate limits for different GitHub API resources and authentication contexts
/// (app-level vs installation-level) and provides methods to check rate limits
/// before making requests.
///
/// # Examples
///
/// ```
/// use github_bot_sdk::client::{RateLimiter, RateLimitContext};
/// use github_bot_sdk::auth::InstallationId;
///
/// let rate_limiter = RateLimiter::new(0.1); // 10% safety margin
///
/// // Check if we can make an app-level request
/// if rate_limiter.can_proceed(&RateLimitContext::App, "core") {
///     // Make API request
/// }
///
/// // Check if we can make an installation-level request
/// let install_id = InstallationId::new(12345);
/// if rate_limiter.can_proceed(&RateLimitContext::Installation(install_id), "core") {
///     // Make API request
/// }
/// ```
#[derive(Debug, Clone)]
pub struct RateLimiter {
    /// Rate limits by (context, resource) key
    limits: Arc<RwLock<HashMap<(RateLimitContext, String), RateLimit>>>,
    /// Safety margin (0.0 to 1.0) - buffer before hitting limits
    margin: f64,
}

impl RateLimiter {
    /// Create a new rate limiter with the specified safety margin.
    ///
    /// # Arguments
    ///
    /// * `margin` - Safety margin (0.0 to 1.0) to keep as a buffer
    ///
    /// # Examples
    ///
    /// ```
    /// use github_bot_sdk::client::RateLimiter;
    ///
    /// // Keep 10% buffer
    /// let limiter = RateLimiter::new(0.1);
    /// ```
    pub fn new(margin: f64) -> Self {
        Self {
            limits: Arc::new(RwLock::new(HashMap::new())),
            margin: margin.clamp(0.0, 1.0),
        }
    }

    /// Update rate limit information from response headers.
    ///
    /// # Arguments
    ///
    /// * `context` - The authentication context (app or installation)
    /// * `headers` - HTTP response headers containing rate limit info
    pub fn update_from_headers(&self, context: &RateLimitContext, headers: &HeaderMap) {
        if let Some(rate_limit) = parse_rate_limit_from_headers(headers) {
            let resource = rate_limit.resource().to_string();
            if let Ok(mut limits) = self.limits.write() {
                limits.insert((context.clone(), resource), rate_limit);
            }
        }
    }

    /// Check if we can proceed with a request for the given context and resource.
    ///
    /// # Arguments
    ///
    /// * `context` - The authentication context (app or installation)
    /// * `resource` - The resource type (e.g., "core", "search")
    ///
    /// # Returns
    ///
    /// `true` if we have sufficient rate limit remaining (considering safety margin),
    /// `false` if we're at or near the rate limit.
    pub fn can_proceed(&self, context: &RateLimitContext, resource: &str) -> bool {
        // If we don't have rate limit data yet, allow the request
        let limits = match self.limits.read() {
            Ok(limits) => limits,
            Err(_) => return true, // Lock poisoned, allow request
        };

        // If we don't have data for this context/resource, allow the request
        let key = (context.clone(), resource.to_string());
        let rate_limit = match limits.get(&key) {
            Some(limit) => limit,
            None => return true,
        };

        // If the rate limit has reset, allow the request
        if rate_limit.has_reset() {
            return true;
        }

        // Check if we're exhausted or near exhaustion
        !rate_limit.is_exhausted() && !rate_limit.is_near_exhaustion(self.margin)
    }

    /// Get the current rate limit for a context and resource.
    ///
    /// # Arguments
    ///
    /// * `context` - The authentication context (app or installation)
    /// * `resource` - The resource type
    ///
    /// # Returns
    ///
    /// `Some(RateLimit)` if we have rate limit data for this context/resource,
    /// `None` if we haven't received rate limit headers yet.
    pub fn get_limit(&self, context: &RateLimitContext, resource: &str) -> Option<RateLimit> {
        let key = (context.clone(), resource.to_string());
        self.limits.read().ok()?.get(&key).cloned()
    }
}

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

#[cfg(test)]
#[path = "rate_limit_tests.rs"]
mod tests;