feat: Proactive Rate Limitting in Github Client (#75)

* implemented a rate limitter

* adding threshold to config

* fixing the update rate limit

* minor change

* dropping locka guard and reading the states only

* minor commit

* added jitter

* fixed the jitter issue

* fixed error handling for headers

* refactor: removed commented function for cleanup

* docs: update installation instructions to include prerequisites and database setup (#76)

* fix:removed double call to update_rate_limit

* refactor:fixed formatting and changed some field names

* template for testing the rate limit funciton

* test:adding test for rate_limmitting

* minor change

* jitter tests

* fix: add threshold as an input parameter

* fix: better error handling message for failing to parse header

* chore: formating

* chore: add docs for functions

* fix: cargo fmt

* chore: add docs to field

---------

Co-authored-by: rituraj-tripathy_pinegit <rituraj.tripathy@pinelabs.com>
Co-authored-by: isSerge <isserge@protonmail.com>

Author: 3 people

Cargo.lock

@@ -24,3 +24,4 @@ thiserror = "2.0.12"
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter", "fmt"] }
 futures = "0.3.31"
+rand = "0.9.1"

Cargo.toml

@@ -18,10 +18,11 @@ pub async fn handle(ctx: Context<'_>, label_name: &str, label_page: usize) -> Bo
 
     let (repo_id, from_page) = match dialogue_state {
         Some(CommandState::ViewingRepoLabels { repo_id, from_page }) => (repo_id, from_page),
-        _ =>
+        _ => {
             return Err(BotHandlerError::InvalidInput(
                 "Invalid state: expected ViewingRepoLabels".to_string(),
-            )),
+            ));
+        }
     };
 
     let repo =

src/bot_handler/callbacks/toggle_label.rs

@@ -6,6 +6,7 @@ const DEFAULT_POLL_INTERVAL: u64 = 10;
 const DEFAULT_REPOS_PER_USER: usize = 20;
 const DEFAULT_LABELS_PER_REPO: usize = 10;
 const DEFAULT_MAX_CONCURRENCY: usize = 10;
+const DEFAULT_RATE_LIMIT_THRESHOLD: u64 = 10;
 
 /// Represents the application configuration.
 #[derive(Debug)]
@@ -26,6 +27,8 @@ pub struct Config {
     pub max_labels_per_repo: usize,
     /// The maximum number of concurrent requests to make to the GitHub API.
     pub max_concurrency: usize,
+    /// The threshold before the bot should pause operations.
+    pub rate_limit_threshold: u64,
 }
 
 impl Config {
@@ -54,6 +57,10 @@ impl Config {
                 .ok()
                 .and_then(|v| v.parse().ok())
                 .unwrap_or(DEFAULT_MAX_CONCURRENCY),
+            rate_limit_threshold: env::var("RATE_LIMIT_THRESHOLD")
+                .ok()
+                .and_then(|v| v.parse().ok())
+                .unwrap_or(DEFAULT_RATE_LIMIT_THRESHOLD),
         })
     }
 }

src/config.rs

@@ -2,17 +2,29 @@
 #[cfg(test)]
 mod tests;
 
-use std::{collections::HashSet, time::Duration};
+use std::{
+    collections::HashSet,
+    sync::Arc,
+    time::{Duration, Instant},
+};
 
 use async_trait::async_trait;
 use backoff::{Error as BackoffError, ExponentialBackoff, future::retry};
 use graphql_client::{GraphQLQuery, Response};
 use mockall::automock;
+use rand::{Rng, rng};
 use reqwest::{
     Client,
     header::{AUTHORIZATION, HeaderMap, HeaderValue, USER_AGENT},
 };
 use thiserror::Error;
+use tokio::sync::Mutex;
+
+#[derive(Debug)]
+struct RateLimitState {
+    remaining: u32,
+    reset_at: Instant,
+}
 
 /// Represents errors that can occur when interacting with the GitHub API.
 #[derive(Debug, Error)]
@@ -43,6 +55,10 @@ pub enum GithubError {
     /// An error indicating that the request was not authorized.
     #[error("GitHub authentication failed")]
     Unauthorized,
+
+    /// An error indicating that a required header could not be parsed.
+    #[error("Failed to parse header: {0}")]
+    HeaderError(String),
 }
 
 // Helper function to check if a GraphQL error is retryable
@@ -117,22 +133,33 @@ pub struct Labels;
 pub struct DefaultGithubClient {
     client: Client,
     graphql_url: String,
+    rate_limit: Arc<Mutex<RateLimitState>>,
+    rate_limit_threshold: u64,
 }
 
 impl DefaultGithubClient {
     /// Creates a new `DefaultGithubClient`.
-    pub fn new(github_token: &str, graphql_url: &str) -> Result<Self, GithubError> {
+    pub fn new(
+        github_token: &str,
+        graphql_url: &str,
+        rate_limit_threshold: u64,
+    ) -> Result<Self, GithubError> {
         // Build the HTTP client with the GitHub token.
         let mut headers = HeaderMap::new();
 
         headers.insert(AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {github_token}"))?);
         headers.insert(USER_AGENT, HeaderValue::from_static("github-activity-rs"));
 
         let client = reqwest::Client::builder().default_headers(headers).build()?;
-
+        let initial_state = RateLimitState { remaining: u32::MAX, reset_at: Instant::now() };
         tracing::debug!("HTTP client built successfully.");
 
-        Ok(Self { client, graphql_url: graphql_url.to_string() })
+        Ok(Self {
+            client,
+            graphql_url: graphql_url.to_string(),
+            rate_limit: Arc::new(Mutex::new(initial_state)),
+            rate_limit_threshold,
+        })
     }
 
     /// Re-usable configuration for exponential backoff.
@@ -158,6 +185,9 @@ impl DefaultGithubClient {
     {
         // closure that Backoff expects
         let operation = || async {
+            // 0. Rate limit guard
+            self.rate_limit_guard().await;
+
             // 1. Build the request
             let request_body = Q::build_query(variables.clone());
 
@@ -170,7 +200,13 @@ impl DefaultGithubClient {
                     },
                 )?;
 
-            // 3. HTTP-status check
+            //3 Update rate limit state from headers
+            if let Err(e) = self.update_rate_limit_from_headers(resp.headers()).await {
+                // Option A: warn and continue
+                tracing::warn!("Could not update rate-limit info: {}", e);
+            }
+
+            // 4. HTTP-status check
             if !resp.status().is_success() {
                 let status = resp.status();
                 let text = resp.text().await.unwrap_or_else(|e| {
@@ -214,15 +250,15 @@ impl DefaultGithubClient {
                 return Err(be);
             }
 
-            // 4. Parse JSON
+            // 5. Parse JSON
             let body: Response<Q::ResponseData> = resp.json().await.map_err(|e| {
                 tracing::warn!("Failed to parse JSON: {e}. Retrying...");
                 BackoffError::transient(GithubError::GraphQLApiError(format!(
                     "JSON parse error: {e}"
                 )))
             })?;
 
-            // 5. GraphQL errors?
+            // 6. GraphQL errors?
             if let Some(errors) = &body.errors {
                 let is_rate_limit_error = errors.iter().any(|e| {
                     e.message.to_lowercase().contains("rate limit") || is_retryable_graphql_error(e)
@@ -239,7 +275,7 @@ impl DefaultGithubClient {
                 }
             }
 
-            // 6. Unwrap the data or permanent-fail
+            // 7. Unwrap the data or permanent-fail
             body.data.ok_or_else(|| {
                 tracing::error!("GraphQL response had no data field; permanent failure");
                 BackoffError::permanent(GithubError::GraphQLApiError(
@@ -251,6 +287,80 @@ impl DefaultGithubClient {
         // kick off the retry loop
         retry(Self::backoff_config(), operation).await
     }
+
+    /// Rate limit guard that sleeps until the rate limit resets if we're close
+    /// to the threshold.
+    async fn rate_limit_guard(&self) {
+        let (remaining, reset_at) = {
+            let state = self.rate_limit.lock().await;
+            (state.remaining, state.reset_at)
+        };
+
+        // define a safety threshold
+        let threshold = self.rate_limit_threshold as u32;
+        if remaining <= threshold {
+            let now = Instant::now();
+            if now < reset_at {
+                let wait = reset_at - now;
+                tracing::info!(
+                    "Approaching rate limit ({} left). Sleeping {:?} until reset...",
+                    remaining,
+                    wait
+                );
+
+                // Sleep until the rate limit resets
+                //added a jitter to avoid thundering herd problem
+                let max_jitter = wait.as_millis() as u64 / 10;
+                let jitter_ms = rng().random_range(0..=max_jitter);
+                tokio::time::sleep(wait + Duration::from_millis(jitter_ms)).await;
+            }
+        }
+    }
+
+    /// Update the rate limit state from the response headers.
+    async fn update_rate_limit_from_headers(&self, headers: &HeaderMap) -> Result<(), GithubError> {
+        // Names are case-insensitive in HeaderMap
+        let rem_val = headers.get("X-RateLimit-Remaining").ok_or_else(|| {
+            let msg = "Missing X-RateLimit-Remaining header".to_string();
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+        let rem_str = rem_val.to_str().map_err(|e| {
+            let msg = format!("Invalid X-RateLimit-Remaining value: {e}");
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+        let remaining = rem_str.parse::<u32>().map_err(|e| {
+            let msg = format!("Cannot parse remaining as u32: {e}");
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+
+        let reset_val = headers.get("X-RateLimit-Reset").ok_or_else(|| {
+            let msg = "Missing X-RateLimit-Reset header".to_string();
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+        let reset_str = reset_val.to_str().map_err(|e| {
+            let msg = format!("Invalid X-RateLimit-Reset value: {e}");
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+        let reset_unix = reset_str.parse::<u64>().map_err(|e| {
+            let msg = format!("Cannot parse reset timestamp as u64: {e}");
+            tracing::error!("{}", msg);
+            GithubError::HeaderError(msg)
+        })?;
+
+        // All good — update the shared state
+        let mut state = self.rate_limit.lock().await;
+        state.remaining = remaining;
+        let reset_in = reset_unix.saturating_sub(chrono::Utc::now().timestamp() as u64);
+        state.reset_at = Instant::now() + Duration::from_secs(reset_in);
+
+        tracing::debug!("Rate limit updated: {} remaining, resets in {}s", remaining, reset_in);
+        Ok(())
+    }
 }
 
 #[async_trait]