From 211d8681b2bd2182bbb39ff74a27daedf93b0562 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@harmonic.ai>
Date: Sat, 7 May 2022 23:11:42 -0400
Subject: [PATCH] Sleep for ratelimit depletion also.

---
 src/githubctl/github_endpoint_watcher.rs | 53 +++++++++++++++++++-----
 1 file changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/githubctl/github_endpoint_watcher.rs b/src/githubctl/github_endpoint_watcher.rs
index 0369719..04abb38 100644
--- a/src/githubctl/github_endpoint_watcher.rs
+++ b/src/githubctl/github_endpoint_watcher.rs
@@ -9,6 +9,8 @@ pub struct GithubEndpointWatcher {
     http_client: reqwest::Client,
     etag: Option<HeaderValue>,
     next_poll_allowed: u64,
+    ratelimit_remaining: Option<u64>,
+    ratelimit_reset: Option<u64>,
 }
 
 impl GithubEndpointWatcher {
@@ -25,6 +27,8 @@ impl GithubEndpointWatcher {
             http_client,
             etag: None,
             next_poll_allowed: 0,
+            ratelimit_remaining: None,
+            ratelimit_reset: None,
         })
     }
 
@@ -39,29 +43,27 @@ impl GithubEndpointWatcher {
             request = request.header(reqwest::header::ETAG, etag);
         }
         self.sleep_until_next_poll().await?;
+        self.sleep_until_ratelimit().await?;
         let request_started_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
         println!("Hitting url {}", self.url);
         let response = request.send().await?;
 
         let headers = response.headers();
         let etag = headers.get(reqwest::header::ETAG).map(|x| x.to_owned());
-        let poll_interval = headers.get("x-poll-interval").map(|x| x.to_owned());
+        let poll_interval = number_header(headers.get("x-poll-interval"))?;
         // let ratelimit_limit = headers.get("x-ratelimit-limit").map(|x| x.to_owned());
-        // let ratelimit_remaining = headers.get("x-ratelimit-remaining").map(|x| x.to_owned());
-        // let ratelimit_reset = headers.get("x-ratelimit-reset").map(|x| x.to_owned());
+        self.ratelimit_remaining = number_header(headers.get("x-ratelimit-remaining"))?;
+        self.ratelimit_reset = number_header(headers.get("x-ratelimit-reset"))?;
         // let ratelimit_used = headers.get("x-ratelimit-used").map(|x| x.to_owned());
         // let ratelimit_resource = headers.get("x-ratelimit-resource").map(|x| x.to_owned());
 
         let body = response.json::<serde_json::Value>().await?;
         println!("{}", serde_json::to_string(&body)?);
         self.etag = etag;
-        let poll_interval_parsed: u64 = match poll_interval {
-            Some(header) => header.to_str()?.parse()?,
-            None => {
-                println!("No poll interval returned, defaulting to 5 minute polling.");
-                300
-            }
-        };
+        let poll_interval_parsed: u64 = poll_interval.unwrap_or_else(|| {
+            println!("No poll interval returned, defaulting to 5 minute polling.");
+            300
+        });
         self.next_poll_allowed = request_started_at + poll_interval_parsed;
         Ok(body)
     }
@@ -76,4 +78,35 @@ impl GithubEndpointWatcher {
         }
         Ok(())
     }
+
+    async fn sleep_until_ratelimit(&self) -> Result<(), Box<dyn std::error::Error>> {
+        match self.ratelimit_remaining {
+            Some(remaining) if remaining <= 0 => {}
+            _ => {
+                // If the number of remaining requests is either
+                // unknown (None) or greater than 0, we do not need to
+                // sleep.
+                return Ok(());
+            }
+        };
+        let reset = self
+            .ratelimit_reset
+            .expect("The above match statement proves this is a Some()");
+        loop {
+            let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
+            if now >= reset {
+                break;
+            }
+            sleep(Duration::from_secs(reset - now)).await;
+        }
+        Ok(())
+    }
+}
+
+fn number_header(header: Option<&HeaderValue>) -> Result<Option<u64>, Box<dyn std::error::Error>> {
+    Ok(header
+        .map(|hdr| hdr.to_str())
+        .transpose()?
+        .map(|hdr| hdr.parse())
+        .transpose()?)
 }