From 3a1f34933ba856150de33f7815cdf66cd11cf74a Mon Sep 17 00:00:00 2001
From: Nelson J Morais <njmorais@gmail.com>
Date: Sat, 21 Nov 2020 02:08:13 +0000
Subject: [PATCH 1/2] GHA runners collector

---
 src/collectors/github_runners.rs | 174 +++++++++++++++++++++++++++++++
 src/collectors/mod.rs            |   9 ++
 src/config.rs                    |  11 ++
 src/lib.rs                       |   3 +-
 4 files changed, 196 insertions(+), 1 deletion(-)
 create mode 100644 src/collectors/github_runners.rs

diff --git a/src/collectors/github_runners.rs b/src/collectors/github_runners.rs
new file mode 100644
index 0000000..20076d5
--- /dev/null
+++ b/src/collectors/github_runners.rs
@@ -0,0 +1,174 @@
+use crate::Config;
+use anyhow::Error;
+use log::{debug, error};
+use prometheus::core::Desc;
+use prometheus::proto::MetricFamily;
+use prometheus::{core::Collector, IntGauge, Opts};
+use reqwest::header::{ACCEPT, AUTHORIZATION, USER_AGENT};
+use reqwest::Method;
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+use tokio::time::Duration;
+
+const GH_RUNNERS_ENDPOINT: &str = "https://api.github.com/repos/{owner_repo}/actions/runners";
+
+#[derive(Debug, serde::Deserialize)]
+struct ApiResponse {
+    total_count: usize,
+    runners: Vec<Runner>,
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct Runner {
+    id: usize,
+    name: String,
+    os: String,
+    status: String,
+    busy: bool,
+    labels: Vec<Label>,
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct Label {
+    id: usize,
+    name: String,
+    #[serde(rename = "type")]
+    the_type: String,
+}
+
+#[derive(Clone)]
+pub struct GithubRunners {
+    //api token to use
+    token: String,
+    // repos to track gha runners
+    repos: Vec<String>,
+    // metric namespace
+    ns: String,
+    // actual metrics
+    metrics: Arc<RwLock<Vec<IntGauge>>>,
+    // default metric description
+    desc: Desc,
+}
+
+impl GithubRunners {
+    pub async fn new(config: &Config) -> Result<Self, Error> {
+        let token = config.rust_runners_token.to_string();
+        let repos: Vec<String> = config
+            .gha_runners_repos
+            .split(',')
+            .map(|v| v.trim().to_string())
+            .collect();
+
+        let ns = String::from("gha_runner");
+        let rv = Self {
+            token,
+            repos,
+            ns: ns.clone(),
+            metrics: Arc::new(RwLock::new(Vec::new())),
+            desc: Desc::new(
+                ns,
+                "GHA runner's status".to_string(),
+                Vec::new(),
+                HashMap::new(),
+            )
+            .unwrap(),
+        };
+
+        let refresh_rate = config.gha_runners_cache_refresh;
+        let mut rv2 = rv.clone();
+        tokio::spawn(async move {
+            loop {
+                if let Err(e) = rv2.update_stats().await {
+                    error!("{:#?}", e);
+                }
+
+                tokio::time::delay_for(Duration::from_secs(refresh_rate)).await;
+            }
+        });
+
+        Ok(rv)
+    }
+
+    async fn update_stats(&mut self) -> Result<(), Error> {
+        let mut gauges = Vec::new();
+        let client = reqwest::Client::new();
+
+        for repo in self.repos.iter() {
+            let url = String::from(GH_RUNNERS_ENDPOINT).replace("{owner_repo}", repo);
+
+            debug!("Querying gha runner's status at: {}", url);
+            let req = client
+                .request(Method::GET, &url)
+                .header(
+                    USER_AGENT,
+                    "https://github.com/rust-lang/monitorbot (infra@rust-lang.org)",
+                )
+                .header(AUTHORIZATION, format!("{} {}", "token", self.token))
+                .header(ACCEPT, "application/vnd.github.v3+json")
+                .build()?;
+
+            let resp = client.execute(req).await?.json::<ApiResponse>().await?;
+
+            //debug!("ApiResponse: {:#?}", resp);
+
+            // convert to metrics
+            for runner in resp.runners.iter() {
+                let status = &runner.status.clone();
+                let value_busy = if runner.busy { 1 } else { 0 };
+                let label_repo = repo.clone();
+                let label_runner = runner.name.clone();
+
+                // online
+                let online = IntGauge::with_opts(
+                    Opts::new("online", "runner is online.")
+                        .namespace(self.ns.clone())
+                        .const_label("repo", label_repo.clone())
+                        .const_label("runner", label_runner.clone()),
+                )
+                .unwrap();
+
+                online.set(if status == "online" { 1 } else { 0 });
+                gauges.push(online);
+
+                // busy
+                let busy = IntGauge::with_opts(
+                    Opts::new("busy", "runner is busy.")
+                        .namespace(self.ns.clone())
+                        .const_label("repo", label_repo)
+                        .const_label("runner", label_runner),
+                )
+                .unwrap();
+
+                busy.set(value_busy);
+                gauges.push(busy);
+            }
+        }
+
+        // lock and replace old data
+        let mut guard = self.metrics.write().unwrap();
+        *guard = gauges;
+
+        Ok(())
+    }
+}
+
+impl Collector for GithubRunners {
+    fn desc(&self) -> Vec<&Desc> {
+        vec![&self.desc]
+    }
+
+    fn collect(&self) -> Vec<MetricFamily> {
+        self.metrics.read().map_or_else(
+            |e| {
+                error!("Unable to collect: {:#?}", e);
+                Vec::new()
+            },
+            |guard| {
+                guard.iter().fold(Vec::new(), |mut acc, item| {
+                    acc.extend(item.collect());
+                    acc
+                })
+            },
+        )
+    }
+}
diff --git a/src/collectors/mod.rs b/src/collectors/mod.rs
index f6ba43c..1b45196 100644
--- a/src/collectors/mod.rs
+++ b/src/collectors/mod.rs
@@ -1,6 +1,8 @@
 mod github_rate_limit;
+mod github_runners;
 
 pub use crate::collectors::github_rate_limit::GitHubRateLimit;
+pub use crate::collectors::github_runners::GithubRunners;
 
 use crate::MetricProvider;
 use anyhow::{Error, Result};
@@ -14,5 +16,12 @@ pub async fn register_collectors(p: &MetricProvider) -> Result<(), Error> {
             info!("Registering GitHubRateLimit collector");
             p.register_collector(rl)
         })
+        .await?;
+
+    GithubRunners::new(&p.config)
+        .and_then(|gr| async {
+            info!("Registering GitHubActionsRunners collector");
+            p.register_collector(gr)
+        })
         .await
 }
diff --git a/src/config.rs b/src/config.rs
index fed220e..32c3924 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -8,6 +8,14 @@ const ENVIRONMENT_VARIABLE_PREFIX: &str = "MONITORBOT_";
 pub struct Config {
     // authorization secret (token) to be able to scrape the metrics endpoint
     pub secret: String,
+    // github api token to be used when querying for gha runner's status
+    // note: token must have (repo scope) authorization
+    pub rust_runners_token: String,
+    // gh runner's repos to track they status. multiple repos are allowed
+    // ex. "rust,cargo,docs.rs"
+    pub gha_runners_repos: String,
+    // gha runner's status refresh rate frequency (in seconds)
+    pub gha_runners_cache_refresh: u64,
     // http server port to bind to
     pub port: u16,
     // github api tokens to collect rate limit statistics
@@ -20,6 +28,9 @@ impl Config {
     pub fn from_env() -> Result<Self, Error> {
         Ok(Self {
             secret: require_env("SECRET")?,
+            rust_runners_token: require_env("RUST_ORG_TOKEN")?,
+            gha_runners_repos: require_env("RUNNERS_REPOS")?,
+            gha_runners_cache_refresh: default_env("GHA_RUNNERS_REFRESH", 120)?,
             port: default_env("PORT", 3001)?,
             gh_rate_limit_tokens: require_env("RATE_LIMIT_TOKENS")?,
             gh_rate_limit_stats_cache_refresh: default_env("GH_RATE_LIMIT_STATS_REFRESH", 120)?,
diff --git a/src/lib.rs b/src/lib.rs
index a0860be..208ee87 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -25,7 +25,8 @@ pub struct MetricProvider {
 
 impl MetricProvider {
     pub fn new(config: Config) -> Self {
-        let register = Registry::new_custom(None, None).expect("Unable to build Registry");
+        let register = Registry::new_custom(Some("monitorbot".to_string()), None)
+            .expect("Unable to build Registry");
         Self { register, config }
     }
 

From e9d5f17be0b0661daadc9b3ea483e90a571d58ee Mon Sep 17 00:00:00 2001
From: Nelson J Morais <njmorais@gmail.com>
Date: Sun, 29 Nov 2020 01:16:23 +0000
Subject: [PATCH 2/2] rebased and refactor out to http module some modules

---
 src/collectors/github_rate_limit.rs | 233 ++++++++++++++++++----------
 src/collectors/github_runners.rs    |  42 ++---
 src/http.rs                         |  37 +++++
 src/lib.rs                          |   1 +
 4 files changed, 217 insertions(+), 96 deletions(-)
 create mode 100644 src/http.rs

diff --git a/src/collectors/github_rate_limit.rs b/src/collectors/github_rate_limit.rs
index cdca957..1cdd20b 100644
--- a/src/collectors/github_rate_limit.rs
+++ b/src/collectors/github_rate_limit.rs
@@ -1,11 +1,11 @@
 use prometheus::{core::Collector, IntGauge, Opts};
 
-use crate::Config;
-use anyhow::{Context, Error, Result};
-use log::{debug, error};
-use reqwest::header::{ACCEPT, AUTHORIZATION, USER_AGENT};
-use reqwest::{Client, Method, Request};
+use crate::{http, Config};
+use anyhow::{bail, Context, Error, Result};
+use log::{debug, error, warn};
+use prometheus::core::Desc;
 use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
 use tokio::time::Duration;
 
 #[derive(Clone)]
@@ -18,34 +18,17 @@ struct User {
 }
 
 const GH_API_USER_ENDPOINT: &str = "https://api.github.com/user";
-const GH_API_RATE_LIMIT_ENDPOINT: &str = "https://api.github.com/rate_limit";
-
-enum GithubReqBuilder {
-    User,
-    RateLimit,
-}
-
-impl GithubReqBuilder {
-    fn build_request(&self, client: &Client, token: &str) -> Result<Request, Error> {
-        let rb = match self {
-            Self::User => client.request(Method::GET, GH_API_USER_ENDPOINT),
-            Self::RateLimit => client.request(Method::GET, GH_API_RATE_LIMIT_ENDPOINT),
-        };
-
-        rb.header(
-            USER_AGENT,
-            "https://github.com/rust-lang/monitorbot (infra@rust-lang.org)",
-        )
-        .header(AUTHORIZATION, format!("{} {}", "token", token))
-        .header(ACCEPT, "application/vnd.github.v3+json")
-        .build()
-        .map_err(Error::from)
-    }
-}
 
 #[derive(Clone)]
 pub struct GitHubRateLimit {
-    users: Vec<User>,
+    // uninitialized tokens
+    uninit_tokens: Vec<String>,
+    // metric namespace
+    ns: String,
+    // metrics collection
+    metrics: Arc<RwLock<Vec<User>>>,
+    // default metric description
+    desc: Desc,
 }
 
 impl GitHubRateLimit {
@@ -56,11 +39,23 @@ impl GitHubRateLimit {
             .map(|v| v.trim().to_string())
             .collect();
 
-        let users = Self::get_users_for_tokens(tokens)
+        let ns = String::from("github_rate_limit");
+        let users = Self::get_users_for_tokens(ns.clone(), tokens)
             .await
             .context("Unable to get usernames for rate limit stats")?;
 
-        let rv = Self { users };
+        let rv = Self {
+            uninit_tokens: users.1,
+            ns: ns.clone(),
+            metrics: Arc::new(RwLock::new(users.0)),
+            desc: Desc::new(
+                ns,
+                "GH rate limit stats".to_string(),
+                Vec::new(),
+                HashMap::new(),
+            )
+            .unwrap(),
+        };
 
         let refresh_rate = config.gh_rate_limit_stats_cache_refresh;
         let mut rv2 = rv.clone();
@@ -74,15 +69,67 @@ impl GitHubRateLimit {
             }
         });
 
+        // start task to catch up with uninit tokens
+        if !rv.uninit_tokens.is_empty() {
+            let refresh_rate = config.gh_rate_limit_stats_cache_refresh + 10;
+            let mut rv2 = rv.clone();
+            tokio::spawn(async move {
+                loop {
+                    tokio::time::delay_for(Duration::from_secs(refresh_rate)).await;
+                    match Self::get_users_for_tokens(rv2.ns.clone(), rv2.uninit_tokens.clone())
+                        .await
+                    {
+                        Err(e) => {
+                            error!("{:#?}", e);
+                        }
+                        Ok(data) => {
+                            if data.0.is_empty() {
+                                continue;
+                            }
+
+                            if let Ok(mut guard) = rv2.metrics.try_write() {
+                                for u in data.0 {
+                                    guard.push(u)
+                                }
+                            }
+
+                            // drop the task, all tokens are under monitoring
+                            if data.1.is_empty() {
+                                rv2.uninit_tokens.clear();
+                                return;
+                            }
+
+                            // update data for next iteration
+                            rv2.uninit_tokens = data.1;
+                        }
+                    };
+                }
+            });
+        }
+
         Ok(rv)
     }
 
-    async fn get_users_for_tokens(tokens: Vec<String>) -> Result<Vec<User>, Error> {
-        let ns = String::from("monitorbot_github_rate_limit");
-        let mut rv: Vec<User> = Vec::new();
+    async fn get_users_for_tokens(
+        ns: String,
+        tokens: Vec<String>,
+    ) -> Result<(Vec<User>, Vec<String>), Error> {
+        let mut uninit_tokens: Vec<String> = Vec::with_capacity(tokens.len());
+        let mut metrics: Vec<User> = Vec::new();
         for token in tokens.into_iter() {
+            let username = match GitHubRateLimit::get_github_api_username(&token).await {
+                Ok(login) => login,
+                Err(e) => {
+                    warn!(
+                        "unable to get token's '{}' username at this time\n{:?}",
+                        token, e
+                    );
+                    uninit_tokens.push(token);
+                    continue;
+                }
+            };
+
             let ns2 = ns.clone();
-            let username = GitHubRateLimit::get_github_api_username(&token).await?;
             let user_future = tokio::task::spawn_blocking(move || {
                 let rate_limit = IntGauge::with_opts(
                     Opts::new("limit", "Rate limit.")
@@ -100,7 +147,7 @@ impl GitHubRateLimit {
 
                 let rate_reset = IntGauge::with_opts(
                     Opts::new("reset", "Rate reset.")
-                        .namespace(ns2.clone())
+                        .namespace(ns2)
                         .const_label("username", username.clone()),
                 )
                 .unwrap();
@@ -119,10 +166,10 @@ impl GitHubRateLimit {
                 _ => panic!("We need to decide if we wanna panic or keep going"),
             };
 
-            rv.push(user);
+            metrics.push(user);
         }
 
-        Ok(rv)
+        Ok((metrics, uninit_tokens))
     }
 
     async fn get_github_api_username(token: &str) -> Result<String, Error> {
@@ -131,44 +178,69 @@ impl GitHubRateLimit {
             pub login: String,
         }
 
-        let client = reqwest::Client::new();
-        let req = GithubReqBuilder::User.build_request(&client, &token)?;
-        let u = client.execute(req).await?.json::<GithubUser>().await?;
-
-        Ok(u.login)
+        match http::is_token_flagged(token).await {
+            Err(e) => bail!("checking if token is flagged: {:?}", e),
+            Ok(v) => {
+                if v {
+                    bail!("Token is flagged. unable to get username")
+                } else {
+                    Ok(http::get(&token, GH_API_USER_ENDPOINT)
+                        .send()
+                        .await?
+                        .json::<GithubUser>()
+                        .await?
+                        .login)
+                }
+            }
+        }
     }
 
     async fn update_stats(&mut self) -> Result<(), Error> {
         debug!("Updating rate limit stats");
 
-        #[derive(Debug, serde::Deserialize)]
-        struct GithubRateLimit {
-            pub rate: HashMap<String, usize>,
-        }
-
-        let client = reqwest::Client::new();
-        for u in self.users.iter_mut() {
-            let req = GithubReqBuilder::RateLimit
-                .build_request(&client, &u.token)
-                .context("Unable to build request to update stats")?;
-
-            let response = client
-                .execute(req)
-                .await
-                .context("Unable to execute request to update stats")?;
+        // lock and read what to query for
+        let tokens = {
+            self.metrics.read().map_or_else(
+                |_| Vec::new(),
+                |lock| {
+                    lock.iter().fold(Vec::new(), |mut acc, item| {
+                        acc.push(item.token.clone());
+                        acc
+                    })
+                },
+            )
+        };
 
-            let mut data = response
-                .json::<GithubRateLimit>()
-                .await
-                .context("Unable to deserialize rate limit stats")?;
+        // querying new updated data
+        let mut new_stats = Vec::<(String, i64, i64, i64)>::with_capacity(tokens.len());
+        for token in tokens.into_iter() {
+            let mut data = http::get_token_rate_limit_stats(&token).await?;
 
             let remaining = data.rate.remove("remaining").unwrap_or(0);
             let limit = data.rate.remove("limit").unwrap_or(0);
             let reset = data.rate.remove("reset").unwrap_or(0);
 
-            u.remaining.set(remaining as i64);
-            u.reset.set(reset as i64);
-            u.limit.set(limit as i64);
+            new_stats.push((
+                token.to_owned(),
+                remaining as i64,
+                limit as i64,
+                reset as i64,
+            ));
+        }
+
+        // lock and write
+        if let Ok(guard) = self.metrics.try_write() {
+            for new_data in new_stats.into_iter() {
+                guard
+                    .iter()
+                    .find(|item| item.token == new_data.0)
+                    .map(|item| {
+                        item.remaining.set(new_data.1);
+                        item.limit.set(new_data.2);
+                        item.reset.set(new_data.3);
+                    })
+                    .unwrap();
+            }
         }
 
         Ok(())
@@ -176,20 +248,25 @@ impl GitHubRateLimit {
 }
 
 impl Collector for GitHubRateLimit {
-    fn desc(&self) -> std::vec::Vec<&prometheus::core::Desc> {
-        // descriptions are being defined in the initialization of the metrics options
-        Vec::default()
+    fn desc(&self) -> Vec<&Desc> {
+        vec![&self.desc]
     }
 
     fn collect(&self) -> std::vec::Vec<prometheus::proto::MetricFamily> {
-        // collect MetricFamilys.
-        let mut mfs = Vec::new();
-        for user in self.users.iter() {
-            mfs.extend(user.limit.collect());
-            mfs.extend(user.remaining.collect());
-            mfs.extend(user.reset.collect());
-        }
-
-        mfs
+        self.metrics.read().map_or_else(
+            |e| {
+                error!("Unable to collect: {:#?}", e);
+                Vec::new()
+            },
+            |guard| {
+                guard.iter().fold(Vec::new(), |mut acc, item| {
+                    acc.extend(item.limit.collect());
+                    acc.extend(item.remaining.collect());
+                    acc.extend(item.reset.collect());
+
+                    acc
+                })
+            },
+        )
     }
 }
diff --git a/src/collectors/github_runners.rs b/src/collectors/github_runners.rs
index 20076d5..c3b7bd5 100644
--- a/src/collectors/github_runners.rs
+++ b/src/collectors/github_runners.rs
@@ -1,11 +1,9 @@
-use crate::Config;
+use crate::{http, http::is_token_flagged, Config};
 use anyhow::Error;
-use log::{debug, error};
+use log::{debug, error, warn};
 use prometheus::core::Desc;
 use prometheus::proto::MetricFamily;
 use prometheus::{core::Collector, IntGauge, Opts};
-use reqwest::header::{ACCEPT, AUTHORIZATION, USER_AGENT};
-use reqwest::Method;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 use tokio::time::Duration;
@@ -91,25 +89,33 @@ impl GithubRunners {
 
     async fn update_stats(&mut self) -> Result<(), Error> {
         let mut gauges = Vec::new();
-        let client = reqwest::Client::new();
-
         for repo in self.repos.iter() {
             let url = String::from(GH_RUNNERS_ENDPOINT).replace("{owner_repo}", repo);
 
-            debug!("Querying gha runner's status at: {}", url);
-            let req = client
-                .request(Method::GET, &url)
-                .header(
-                    USER_AGENT,
-                    "https://github.com/rust-lang/monitorbot (infra@rust-lang.org)",
-                )
-                .header(AUTHORIZATION, format!("{} {}", "token", self.token))
-                .header(ACCEPT, "application/vnd.github.v3+json")
-                .build()?;
+            // does this token still able to query github api? (rate limit wise)
+            match is_token_flagged(&self.token).await {
+                Err(e) => {
+                    error!("checking if token is flagged: {:?}", e);
+                    continue;
+                }
+                Ok(true) => {
+                    warn!(
+                        "token: '{}' is currently flagged. skipping data gathering",
+                        &self.token
+                    );
+                    continue;
+                }
+                Ok(false) => {}
+            }
 
-            let resp = client.execute(req).await?.json::<ApiResponse>().await?;
+            debug!("Querying gha runner's status at: {}", url);
+            let resp = http::get(&self.token, &url)
+                .send()
+                .await?
+                .json::<ApiResponse>()
+                .await?;
 
-            //debug!("ApiResponse: {:#?}", resp);
+            debug!("ApiResponse: {:#?}", resp);
 
             // convert to metrics
             for runner in resp.runners.iter() {
diff --git a/src/http.rs b/src/http.rs
new file mode 100644
index 0000000..4f6b141
--- /dev/null
+++ b/src/http.rs
@@ -0,0 +1,37 @@
+use anyhow::{Context, Error, Result};
+use reqwest::header::{ACCEPT, AUTHORIZATION, USER_AGENT};
+use reqwest::{Client, Method, RequestBuilder};
+use std::collections::HashMap;
+
+pub(crate) const GH_API_RATE_LIMIT_ENDPOINT: &str = "https://api.github.com/rate_limit";
+
+pub(crate) fn get(token: &str, url: &str) -> RequestBuilder {
+    Client::new()
+        .request(Method::GET, url)
+        .header(
+            USER_AGENT,
+            "https://github.com/rust-lang/monitorbot (infra@rust-lang.org)",
+        )
+        .header(AUTHORIZATION, format!("{} {}", "token", token))
+        .header(ACCEPT, "application/vnd.github.v3+json")
+}
+
+#[derive(Debug, serde::Deserialize)]
+pub(crate) struct GithubRateLimit {
+    pub rate: HashMap<String, usize>,
+}
+
+pub(crate) async fn is_token_flagged(token: &str) -> Result<bool, Error> {
+    get_token_rate_limit_stats(token)
+        .await
+        .map(|mut r| Ok(r.rate.remove("remaining").unwrap_or(0) == 0))?
+}
+
+pub(crate) async fn get_token_rate_limit_stats(token: &str) -> Result<GithubRateLimit, Error> {
+    get(token, GH_API_RATE_LIMIT_ENDPOINT)
+        .send()
+        .await?
+        .json::<GithubRateLimit>()
+        .await
+        .context("Unable to get rate limit stats")
+}
diff --git a/src/lib.rs b/src/lib.rs
index 208ee87..fcc1c04 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@
 
 pub mod collectors;
 mod config;
+pub(crate) mod http;
 
 pub use config::Config;