use std::path::Path;
use std::sync::{Mutex, PoisonError};
use anyhow::{bail, Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const WARN_PERCENT: f64 = 80.0;
const RESOURCE_KEYS: &[&str] = &["graphql", "core", "search"];
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct RateLimitResource {
pub used: u64,
pub limit: u64,
pub remaining: u64,
pub percent: f64,
pub reset: i64,
}
impl RateLimitResource {
fn from_value(v: &Value) -> Option<Self> {
let used = v.get("used").and_then(Value::as_u64)?;
let limit = v.get("limit").and_then(Value::as_u64)?;
let remaining = v
.get("remaining")
.and_then(Value::as_u64)
.unwrap_or_else(|| limit.saturating_sub(used));
let reset = v.get("reset").and_then(Value::as_i64).unwrap_or(0);
Some(Self {
used,
limit,
remaining,
percent: percent_of(used, limit),
reset,
})
}
#[must_use]
pub fn over_warn(&self) -> bool {
self.percent >= WARN_PERCENT
}
#[must_use]
pub fn new(used: u64, limit: u64, remaining: u64, reset: i64) -> Self {
Self {
used,
limit,
remaining,
percent: percent_of(used, limit),
reset,
}
}
}
fn percent_of(used: u64, limit: u64) -> f64 {
if limit == 0 {
return 0.0;
}
let raw = used as f64 / limit as f64 * 100.0;
(raw * 10.0).round() / 10.0
}
#[derive(Debug, Clone, Copy, PartialEq, Default, Serialize, Deserialize)]
pub struct RateLimitSnapshot {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub graphql: Option<RateLimitResource>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub core: Option<RateLimitResource>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub search: Option<RateLimitResource>,
}
impl RateLimitSnapshot {
fn resources(&self) -> Vec<(&'static str, RateLimitResource)> {
[
("graphql", self.graphql),
("core", self.core),
("search", self.search),
]
.into_iter()
.filter_map(|(name, res)| res.map(|r| (name, r)))
.collect()
}
#[must_use]
pub fn max_percent(&self) -> f64 {
self.resources()
.iter()
.map(|(_, r)| r.percent)
.fold(0.0_f64, f64::max)
}
#[must_use]
pub fn over_warn(&self) -> bool {
self.resources().iter().any(|(_, r)| r.over_warn())
}
#[must_use]
pub fn summary_line(&self) -> String {
let body = self
.resources()
.iter()
.map(|(name, r)| format_resource(name, r))
.collect::<Vec<_>>()
.join(" · ");
if body.is_empty() {
return body;
}
if self.over_warn() {
format!("⚠ {body}")
} else {
body
}
}
#[must_use]
pub fn tray_label(&self) -> String {
let resources = self.resources();
if resources.is_empty() {
return String::new();
}
let body = resources
.iter()
.map(|(name, r)| format!("{name} {}%", trim_percent(r.percent)))
.collect::<Vec<_>>()
.join(" · ");
if self.over_warn() {
format!("github: {body} ⚠")
} else {
format!("github: {body}")
}
}
}
fn format_resource(name: &str, r: &RateLimitResource) -> String {
format!(
"{name} {}% ({}/{}, resets {})",
trim_percent(r.percent),
r.used,
r.limit,
format_reset_utc(r.reset)
)
}
fn trim_percent(percent: f64) -> String {
if (percent.fract()).abs() < f64::EPSILON {
format!("{}", percent as i64)
} else {
format!("{percent}")
}
}
#[must_use]
pub fn format_reset_utc(epoch: i64) -> String {
match DateTime::<Utc>::from_timestamp(epoch, 0) {
Some(dt) if epoch > 0 => dt.format("%H:%MZ").to_string(),
_ => "??:??Z".to_string(),
}
}
fn parse_rate_limit(body: &Value) -> RateLimitSnapshot {
let resources = body.get("resources");
let read = |key: &str| -> Option<RateLimitResource> {
resources
.and_then(|r| r.get(key))
.and_then(RateLimitResource::from_value)
};
debug_assert!(RESOURCE_KEYS.contains(&"graphql"));
RateLimitSnapshot {
graphql: read("graphql"),
core: read("core"),
search: read("search"),
}
}
fn run_gh_rate_limit(bin: &Path) -> Result<Value> {
let output = crate::github_metrics::run_gh(bin, ["api", "rate_limit"], "api rate_limit", None)
.with_context(|| {
format!(
"failed to run {} (is the GitHub CLI installed?)",
bin.display()
)
})?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!("gh api rate_limit failed: {}", stderr.trim());
}
serde_json::from_slice(&output.stdout).context("gh api rate_limit returned invalid JSON")
}
pub fn resolve_rate_limit_with(bin: &Path) -> Result<RateLimitSnapshot> {
let body = run_gh_rate_limit(bin)?;
Ok(parse_rate_limit(&body))
}
#[derive(Debug, Default)]
pub struct RateLimitCache {
snapshot: Mutex<Option<RateLimitSnapshot>>,
}
impl RateLimitCache {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn get(&self) -> Option<RateLimitSnapshot> {
*self.lock()
}
pub fn replace(&self, next: RateLimitSnapshot) -> bool {
let mut guard = self.lock();
let changed = *guard != Some(next);
*guard = Some(next);
changed
}
pub fn observe_graphql(&self, graphql: RateLimitResource) -> bool {
let mut guard = self.lock();
let mut snap = (*guard).unwrap_or_default();
let changed = snap.graphql != Some(graphql);
snap.graphql = Some(graphql);
*guard = Some(snap);
changed
}
fn lock(&self) -> std::sync::MutexGuard<'_, Option<RateLimitSnapshot>> {
self.snapshot.lock().unwrap_or_else(PoisonError::into_inner)
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::float_cmp)]
mod tests {
use super::*;
use crate::test_support::shim::{retry_on_etxtbsy, shim_lock, write_exec_script};
use serde_json::json;
use std::path::PathBuf;
use std::sync::MutexGuard;
fn sample_body() -> Value {
json!({
"resources": {
"core": {"limit": 5000, "used": 27, "remaining": 4973, "reset": 1_700_000_000_i64},
"graphql": {"limit": 5000, "used": 4100, "remaining": 900, "reset": 1_700_003_000_i64},
"search": {"limit": 30, "used": 3, "remaining": 27, "reset": 1_700_000_060_i64},
},
"rate": {"limit": 5000, "used": 27, "remaining": 4973, "reset": 1_700_000_000_i64}
})
}
#[test]
fn parse_reads_every_resource_and_computes_percent() {
let snap = parse_rate_limit(&sample_body());
let graphql = snap.graphql.expect("graphql present");
assert_eq!(graphql.used, 4100);
assert_eq!(graphql.limit, 5000);
assert_eq!(graphql.remaining, 900);
assert_eq!(graphql.percent, 82.0);
assert_eq!(graphql.reset, 1_700_003_000);
let core = snap.core.expect("core present");
assert_eq!(core.percent, 0.5);
assert!(snap.search.is_some());
}
#[test]
fn parse_tolerates_a_missing_search_resource() {
let body = json!({"resources": {
"core": {"limit": 5000, "used": 27, "remaining": 4973, "reset": 1},
"graphql": {"limit": 5000, "used": 10, "remaining": 4990, "reset": 1},
}});
let snap = parse_rate_limit(&body);
assert!(snap.graphql.is_some());
assert!(snap.core.is_some());
assert!(snap.search.is_none());
}
#[test]
fn parse_tolerates_a_missing_resources_block() {
let snap = parse_rate_limit(&json!({}));
assert_eq!(snap, RateLimitSnapshot::default());
}
#[test]
fn parse_derives_remaining_when_absent() {
let body = json!({"resources": {"core": {"limit": 100, "used": 40}}});
let core = parse_rate_limit(&body).core.expect("core present");
assert_eq!(core.remaining, 60);
assert_eq!(core.percent, 40.0);
}
#[test]
fn percent_guards_a_zero_limit() {
assert_eq!(percent_of(0, 0), 0.0);
assert_eq!(percent_of(5, 0), 0.0);
assert_eq!(percent_of(1, 3), 33.3);
}
#[test]
fn over_warn_fires_only_at_or_above_the_threshold() {
let res = |used: u64| RateLimitResource {
used,
limit: 100,
remaining: 100 - used,
percent: percent_of(used, 100),
reset: 0,
};
assert!(!res(79).over_warn());
assert!(res(80).over_warn());
assert!(res(95).over_warn());
let snap = RateLimitSnapshot {
core: Some(res(10)),
graphql: Some(res(85)),
search: None,
};
assert!(snap.over_warn());
assert_eq!(snap.max_percent(), 85.0);
}
#[test]
fn format_reset_utc_renders_clock_time_or_a_placeholder() {
assert_eq!(format_reset_utc(1_700_000_000), "22:13Z");
assert_eq!(format_reset_utc(0), "??:??Z");
assert_eq!(format_reset_utc(-5), "??:??Z");
}
#[test]
fn summary_line_lists_resources_and_marks_the_warn_case() {
let snap = parse_rate_limit(&sample_body());
let line = snap.summary_line();
assert!(line.contains("graphql 82% (4100/5000, resets"), "{line}");
assert!(line.contains("core 0.5% (27/5000"), "{line}");
assert!(line.starts_with("⚠ "), "{line}");
}
#[test]
fn summary_line_omits_the_marker_below_threshold() {
let body = json!({"resources": {
"graphql": {"limit": 5000, "used": 10, "remaining": 4990, "reset": 1_700_000_000_i64},
"core": {"limit": 5000, "used": 27, "remaining": 4973, "reset": 1_700_000_000_i64},
}});
let line = parse_rate_limit(&body).summary_line();
assert!(!line.starts_with('⚠'), "{line}");
assert!(line.starts_with("graphql 0.2%"), "{line}");
}
#[test]
fn summary_line_is_empty_without_resources() {
assert!(RateLimitSnapshot::default().summary_line().is_empty());
assert!(RateLimitSnapshot::default().tray_label().is_empty());
}
#[test]
fn tray_label_is_compact_and_marks_the_warn_case() {
let snap = parse_rate_limit(&sample_body());
let label = snap.tray_label();
assert!(label.starts_with("github: graphql 82%"), "{label}");
assert!(label.ends_with('⚠'), "{label}");
}
#[test]
fn tray_label_omits_the_marker_below_threshold() {
let body = json!({"resources": {
"graphql": {"limit": 5000, "used": 10, "remaining": 4990, "reset": 1},
"core": {"limit": 5000, "used": 27, "remaining": 4973, "reset": 1},
}});
let label = parse_rate_limit(&body).tray_label();
assert_eq!(label, "github: graphql 0.2% · core 0.5%", "{label}");
assert!(!label.contains('⚠'), "{label}");
}
#[test]
fn cache_get_and_replace_report_changes() {
let cache = RateLimitCache::new();
assert!(cache.get().is_none());
let a = parse_rate_limit(&sample_body());
assert!(cache.replace(a));
assert_eq!(cache.get(), Some(a));
assert!(!cache.replace(a));
let b = parse_rate_limit(&json!({"resources": {
"graphql": {"limit": 5000, "used": 4200, "remaining": 800, "reset": 1}
}}));
assert!(cache.replace(b));
}
#[test]
fn observe_graphql_updates_only_graphql_and_preserves_core_search() {
let cache = RateLimitCache::new();
cache.replace(RateLimitSnapshot {
graphql: Some(RateLimitResource::new(10, 5000, 4990, 0)),
core: Some(RateLimitResource::new(3, 5000, 4997, 0)),
search: Some(RateLimitResource::new(0, 30, 30, 0)),
});
assert!(cache.observe_graphql(RateLimitResource::new(45, 5000, 4955, 0)));
let snap = cache.get().unwrap();
assert_eq!(snap.graphql.unwrap().used, 45, "graphql updated");
assert_eq!(snap.core.unwrap().used, 3, "core preserved");
assert_eq!(snap.search.unwrap().used, 0, "search preserved");
assert!(!cache.observe_graphql(RateLimitResource::new(45, 5000, 4955, 0)));
let empty = RateLimitCache::new();
assert!(empty.observe_graphql(RateLimitResource::new(1, 5000, 4999, 0)));
assert!(empty.get().unwrap().core.is_none());
}
fn fake_gh(dir: &Path, stdout: &str, code: i32) -> (PathBuf, MutexGuard<'static, ()>) {
let guard = shim_lock();
let path = dir.join("fake-gh");
write_exec_script(
&path,
&format!("#!/bin/sh\ncat <<'JSON'\n{stdout}\nJSON\nexit {code}\n"),
);
(path, guard)
}
#[test]
fn resolve_errors_when_gh_is_missing() {
let err = resolve_rate_limit_with(Path::new("/no/such/gh/xyzzy")).unwrap_err();
let msg = format!("{err:#}");
assert!(msg.contains("failed to run"), "{msg}");
assert!(msg.contains("GitHub CLI"), "{msg}");
}
#[test]
fn resolve_errors_on_a_nonzero_exit() {
let dir = tempfile::tempdir().unwrap();
let (bin, _shim) = fake_gh(dir.path(), "", 1);
let err = retry_on_etxtbsy(|| resolve_rate_limit_with(&bin)).unwrap_err();
assert!(
format!("{err:#}").contains("gh api rate_limit failed"),
"{err:#}"
);
}
#[test]
fn resolve_errors_on_unparseable_output() {
let dir = tempfile::tempdir().unwrap();
let (bin, _shim) = fake_gh(dir.path(), "not json at all", 0);
let err = retry_on_etxtbsy(|| resolve_rate_limit_with(&bin)).unwrap_err();
assert!(format!("{err:#}").contains("invalid JSON"), "{err:#}");
}
#[test]
fn resolve_reads_a_real_reply_end_to_end() {
let dir = tempfile::tempdir().unwrap();
let (bin, _shim) = fake_gh(dir.path(), &sample_body().to_string(), 0);
let snap = retry_on_etxtbsy(|| resolve_rate_limit_with(&bin)).unwrap();
assert_eq!(snap.graphql.unwrap().percent, 82.0);
assert_eq!(snap.core.unwrap().used, 27);
}
}