Feature/Proactive Rate Limitting in Github Client

Cargo.toml

@@ -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"

readme.md

@@ -73,7 +73,15 @@ src
    cd good-first-bot-rs
    ```
 
-2. **Set up environment variables:**
+2. **Prerequisites:**
+
+   - **Rust:** Ensure you have [Rust](https://www.rust-lang.org/tools/install) installed.
+   - **sqlx-cli:** Install the SQLx command-line tool for database migrations.
+     ```bash
+     cargo install sqlx-cli
+     ```
+
+3. **Set up environment variables:**
 
 Create a .env file in the project root with the following keys (values are
 examples):
@@ -103,10 +111,18 @@ MAX_LABELS_PER_REPO=5
 - MAX_LABELS_PER_REPO: (Optional) Maximum number of labels per repository a user
   can track. Default is 10
 
-3. **Install Dependencies:**
+4. **Database Setup:**
+
+Run the following command to set up the database and apply migrations:
+
+```bash
+mkdir data
+sqlx database setup
+```
+
+5. **Install Dependencies:**
 
-Ensure you have [Rust](https://www.rust-lang.org/tools/install) installed. Then,
-in the project directory run:
+In the project directory run:
 
 ```bash
 cargo build

src/bot_handler/callbacks/toggle_label.rs

@@ -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/config.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;
 
 #[derive(Debug)]
 pub struct Config {
@@ -17,6 +18,7 @@ pub struct Config {
     pub max_repos_per_user: usize,
     pub max_labels_per_repo: usize,
     pub max_concurrency: usize,
+    pub rate_limit_threshold: u64,
 }
 
 impl Config {
@@ -44,6 +46,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/github/mod.rs

@@ -1,6 +1,7 @@
 #[cfg(test)]
 mod tests;
 
+use rand::{Rng, rng};
 use std::{collections::HashSet, time::Duration};
 
 use async_trait::async_trait;
@@ -11,7 +12,15 @@ use reqwest::{
     Client,
     header::{AUTHORIZATION, HeaderMap, HeaderValue, USER_AGENT},
 };
+use std::{sync::Arc, time::Instant};
 use thiserror::Error;
+use tokio::sync::Mutex;
+
+#[derive(Debug)]
+struct RateLimitState {
+    remaining: u32,
+    reset_at: Instant,
+}
 
 #[derive(Debug, Error)]
 pub enum GithubError {
@@ -33,6 +42,8 @@ pub enum GithubError {
     RateLimited,
     #[error("GitHub authentication failed")]
     Unauthorized,
+    #[error("Failed to parse header: {0}")]
+    HeaderError(String),
 }
 
 // Helper function to check if a GraphQL error is retryable
@@ -102,21 +113,32 @@ pub struct Labels;
 pub struct DefaultGithubClient {
     client: Client,
     graphql_url: String,
+    rate_limit: Arc<Mutex<RateLimitState>>,
+    rate_limit_threshold: u64,
 }
 
 impl 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.
@@ -142,6 +164,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());
 
@@ -154,7 +179,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| {
@@ -183,30 +214,33 @@ impl DefaultGithubClient {
                             ))
                         }
                     }
-                    reqwest::StatusCode::NOT_FOUND =>
-                        GithubError::GraphQLApiError(format!("HTTP Not Found ({status}): {text}")),
+                    reqwest::StatusCode::NOT_FOUND => {
+                        GithubError::GraphQLApiError(format!("HTTP Not Found ({status}): {text}"))
+                    }
                     _ => GithubError::GraphQLApiError(format!("HTTP Error ({status}): {text}")),
                 };
 
                 let be = match github_err {
                     GithubError::RateLimited => BackoffError::transient(github_err),
                     _ if status.is_server_error()
                         || status == reqwest::StatusCode::TOO_MANY_REQUESTS =>
-                        BackoffError::transient(github_err),
+                    {
+                        BackoffError::transient(github_err)
+                    }
                     _ => BackoffError::permanent(github_err),
                 };
                 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)
@@ -223,7 +257,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(
@@ -235,6 +269,78 @@ 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; // acquire lock
+            (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;
+            }
+        }
+    }
+
+    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]

src/github/tests.rs

@@ -1,10 +1,8 @@
-use std::collections::HashMap;
-
 use super::*;
-
+use std::collections::HashMap;
 #[test]
 fn test_new_github_client() {
-    let client = DefaultGithubClient::new("test_token", "https://api.github.com/graphql");
+    let client = DefaultGithubClient::new("test_token", "https://api.github.com/graphql", 10);
     assert!(client.is_ok());
 }
 
@@ -37,3 +35,23 @@ fn test_is_not_retryable_graphql_error() {
 
     assert!(!is_retryable_graphql_error(&error));
 }
+
+#[tokio::test]
+async fn test_update_rate_limit_from_headers() {
+    let client =
+        DefaultGithubClient::new("fake", "https://api.github.com/graphql", 5).expect("client init");
+
+    // Build fake headers with remaining=3, reset in 60s
+    let mut headers = HeaderMap::new();
+    headers.insert("X-RateLimit-Remaining", HeaderValue::from_static("3"));
+    let reset_ts = (chrono::Utc::now().timestamp() as u64) + 60;
+    headers.insert("X-RateLimit-Reset", HeaderValue::from_str(&reset_ts.to_string()).unwrap());
+
+    client.update_rate_limit_from_headers(&headers).await;
+
+    let state = client.rate_limit.lock().await;
+    assert_eq!(state.remaining, 3);
+    // We expect reset_at ≈ now + 60s (within a small delta)
+    let diff = state.reset_at.checked_duration_since(Instant::now()).unwrap();
+    assert!(diff >= Duration::from_secs(59) && diff <= Duration::from_secs(61));
+}

src/main.rs

@@ -56,6 +56,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
     let github_client = Arc::new(github::DefaultGithubClient::new(
         &config.github_token,
         &config.github_graphql_url,
+        config.rate_limit_threshold,
     )?);
 
     let messaging_service = Arc::new(TelegramMessagingService::new(bot.clone()));

src/messaging/utils.rs

@@ -13,8 +13,9 @@ pub fn github_color_to_emoji(hex_color: &str) -> &str {
         "fef2c0" | "fbca04" | "e4e669" | "ffeb3b" | "f9e076" | "fadc73" => "🟡",
 
         // Greens
-        "0e8a16" | "006b75" | "5ab302" | "a2eeef" | "008672" | "c2e0c6" | "1aa34a" | "4caf50" =>
-            "🟢",
+        "0e8a16" | "006b75" | "5ab302" | "a2eeef" | "008672" | "c2e0c6" | "1aa34a" | "4caf50" => {
+            "🟢"
+        }
 
         // Blues / Teals
         "0052cc" | "c5def5" | "0075ca" | "1d76db" | "89d2fc" | "00bcd4" | "b3f4f4" => "🔵",

src/poller/mod.rs

@@ -195,6 +195,9 @@ impl GithubPoller {
                         );
                         return Err(PollerError::Github(github_error));
                     }
+                    GithubError::HeaderError(_) => {
+                        tracing::error!("Header error while polling issues for repository");
+                    }
                 },
                 unexpected_error => {
                     tracing::error!(