use std::path::PathBuf;
use std::sync::{Arc, Mutex, PoisonError, TryLockError};
use std::time::{Duration, Instant};
use anyhow::{bail, Result};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde_json::{json, Value};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use crate::daemon::service::{DaemonService, MenuItem, MenuSnapshot, ServiceStatus};
use crate::github_metrics::{GhCounts, IncrementalCounts};
use crate::request_log;
const STARTUP_DELAY: Duration = Duration::from_secs(5);
const PERIODIC_INTERVAL: Duration = Duration::from_secs(10 * 60);
const SHUTDOWN_SCAN_BUDGET: Duration = Duration::from_secs(2);
type SharedTally = Arc<Mutex<IncrementalCounts>>;
pub struct GithubCountersService {
started_at: DateTime<Utc>,
tally: Option<SharedTally>,
stopping: CancellationToken,
logger: Mutex<Option<JoinHandle<()>>>,
}
impl Default for GithubCountersService {
fn default() -> Self {
Self::new()
}
}
impl GithubCountersService {
#[must_use]
pub fn new() -> Self {
Self::with_log_path(request_log::log_file_path())
}
fn with_log_path(path: Option<PathBuf>) -> Self {
let (started_at, tally) = match path {
Some(path) => {
let tally = IncrementalCounts::start(path);
(tally.since(), Some(Arc::new(Mutex::new(tally))))
}
None => (Utc::now(), None),
};
Self {
started_at,
tally,
stopping: CancellationToken::new(),
logger: Mutex::new(None),
}
}
pub fn start_counter_logger(&self) {
self.start_counter_logger_with(STARTUP_DELAY, PERIODIC_INTERVAL);
}
fn start_counter_logger_with(&self, startup_delay: Duration, interval: Duration) {
if tokio::runtime::Handle::try_current().is_err() {
tracing::debug!("no tokio runtime; github counter logger not started");
return;
}
let mut guard = self.logger.lock().unwrap_or_else(PoisonError::into_inner);
if guard.is_some() {
return;
}
let stopping = self.stopping.clone();
let tally = self.tally.clone();
*guard = Some(tokio::spawn(async move {
tokio::select! {
() = stopping.cancelled() => return,
() = tokio::time::sleep(startup_delay) => {}
}
emit_background(&tally, "startup", &stopping).await;
loop {
tokio::select! {
() = stopping.cancelled() => break,
() = tokio::time::sleep(interval) => {
emit_background(&tally, "periodic", &stopping).await;
}
}
}
}));
}
}
fn cancelled_by(token: &CancellationToken) -> impl FnMut() -> bool + Send + 'static {
let token = token.clone();
move || token.is_cancelled()
}
async fn counts_since_boot(
tally: &Option<SharedTally>,
stop: impl FnMut() -> bool + Send + 'static,
) -> (GhCounts, bool) {
let Some(tally) = tally else {
return (GhCounts::default(), true);
};
let reader = Arc::clone(tally);
let read = tokio::task::spawn_blocking(move || {
let mut tally = reader.lock().unwrap_or_else(PoisonError::into_inner);
let current = tally.refresh(stop);
(tally.counts().clone(), current)
})
.await;
match read {
Ok(counts) => counts,
Err(e) => {
tracing::warn!("github counters: the log read did not complete: {e}");
(last_known(tally), false)
}
}
}
fn last_known(tally: &SharedTally) -> GhCounts {
match tally.try_lock() {
Ok(tally) => tally.counts().clone(),
Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner().counts().clone(),
Err(TryLockError::WouldBlock) => GhCounts::default(),
}
}
fn log_summary(phase: &str, counts: &GhCounts, current: bool) {
let mut summary = counts.summary_line();
if !current {
summary.push_str(" (incomplete: the log read was cut short)");
}
tracing::info!("github api calls ({phase}): {summary}");
}
async fn emit_background(tally: &Option<SharedTally>, phase: &str, stopping: &CancellationToken) {
let (counts, current) = counts_since_boot(tally, cancelled_by(stopping)).await;
if !stopping.is_cancelled() {
log_summary(phase, &counts, current);
}
}
#[async_trait]
impl DaemonService for GithubCountersService {
fn name(&self) -> &'static str {
"github"
}
async fn handle(&self, op: &str, _payload: Value) -> Result<Value> {
match op {
"summary" => {
let (counts, current) =
counts_since_boot(&self.tally, cancelled_by(&self.stopping)).await;
let mut value = counts.to_json();
if let Value::Object(map) = &mut value {
map.insert("since".to_string(), json!(self.started_at.to_rfc3339()));
if !current {
map.insert("incomplete".to_string(), json!(true));
}
}
Ok(value)
}
other => bail!("unknown github op: {other}"),
}
}
fn menu(&self) -> MenuSnapshot {
MenuSnapshot {
title: "GitHub API".to_string(),
items: vec![MenuItem::Label(
"call counters logged to daemon.log".to_string(),
)],
}
}
async fn menu_action(&self, _action_id: &str) -> Result<()> {
Ok(())
}
async fn status(&self) -> ServiceStatus {
let (counts, current) = counts_since_boot(&self.tally, cancelled_by(&self.stopping)).await;
let mut summary = format!("{} GitHub API call(s) since start", counts.api_total());
if !current {
summary.push_str(" (incomplete)");
}
ServiceStatus {
name: self.name().to_string(),
healthy: true,
summary,
detail: counts.to_json(),
}
}
async fn shutdown(&self) {
self.stopping.cancel();
let logger = self
.logger
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(logger) = logger {
let _ = logger.await;
}
let deadline = Instant::now() + SHUTDOWN_SCAN_BUDGET;
let (counts, current) =
counts_since_boot(&self.tally, move || Instant::now() >= deadline).await;
log_summary("shutdown", &counts, current);
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use crate::github_metrics::aggregate;
use chrono::{Duration as ChronoDuration, SecondsFormat};
use std::io::Write;
use std::path::Path;
fn gh_line(since: DateTime<Utc>, secs: i64, command: &[&str]) -> String {
format!(
"{{\"kind\":\"gh\",\"timestamp\":\"{}\",\"command\":{},\"source\":\"daemon\"}}\n",
(since + ChronoDuration::seconds(secs)).to_rfc3339_opts(SecondsFormat::Millis, true),
serde_json::to_string(command).unwrap(),
)
}
fn append(path: &Path, text: &str) {
std::fs::OpenOptions::new()
.append(true)
.create(true)
.open(path)
.unwrap()
.write_all(text.as_bytes())
.unwrap();
}
fn tallied(svc: &GithubCountersService) -> u64 {
svc.tally.as_ref().unwrap().lock().unwrap().counts().total()
}
#[tokio::test]
async fn service_reports_status_and_handles_summary() {
let dir = tempfile::tempdir().unwrap();
let svc = GithubCountersService::with_log_path(Some(dir.path().join("log.jsonl")));
assert_eq!(svc.name(), "github");
let status = svc.status().await;
assert_eq!(status.name, "github");
assert!(status.healthy);
assert!(status.detail.get("api_total").is_some());
let summary = svc.handle("summary", Value::Null).await.unwrap();
assert!(summary.get("since").is_some());
assert!(summary.get("by_source").is_some());
assert!(summary.get("incomplete").is_none());
assert!(svc.handle("bogus", Value::Null).await.is_err());
assert_eq!(svc.menu().title, "GitHub API");
svc.menu_action("x").await.unwrap();
svc.shutdown().await;
}
#[tokio::test]
async fn without_a_log_path_there_is_nothing_to_count() {
let svc = GithubCountersService::with_log_path(None);
let status = svc.status().await;
assert!(status.healthy);
assert_eq!(status.summary, "0 GitHub API call(s) since start");
svc.shutdown().await;
}
#[tokio::test]
async fn counts_only_the_records_appended_after_start() {
let dir = tempfile::tempdir().unwrap();
let log = dir.path().join("log.jsonl");
let future = Utc::now() + ChronoDuration::hours(1);
let pre: String = (0..50)
.map(|_| gh_line(future, 0, &["pr", "list"]))
.collect();
append(&log, &pre);
let svc = GithubCountersService::with_log_path(Some(log.clone()));
assert_eq!(
aggregate(&log, Some(svc.started_at), None, None).total(),
50
);
assert_eq!(
svc.status().await.summary,
"0 GitHub API call(s) since start"
);
append(&log, &gh_line(svc.started_at, 1, &["api", "graphql"]));
append(&log, &gh_line(svc.started_at, 2, &["pr", "view"]));
let status = svc.status().await;
assert_eq!(status.summary, "2 GitHub API call(s) since start");
let summary = svc.handle("summary", Value::Null).await.unwrap();
assert_eq!(summary["api_total"], 2);
assert_eq!(summary["by_subcommand"]["pr view"], 1);
assert_eq!(summary["since"], svc.started_at.to_rfc3339());
}
#[tokio::test]
async fn a_scan_is_abandoned_once_shutdown_begins() {
let dir = tempfile::tempdir().unwrap();
let log = dir.path().join("log.jsonl");
let svc = GithubCountersService::with_log_path(Some(log.clone()));
for n in 1..=3 {
append(&log, &gh_line(svc.started_at, n, &["pr", "list"]));
}
svc.stopping.cancel();
let status = svc.status().await;
assert_eq!(
status.summary,
"0 GitHub API call(s) since start (incomplete)"
);
let summary = svc.handle("summary", Value::Null).await.unwrap();
assert_eq!(summary["incomplete"], true);
assert_eq!(tallied(&svc), 0);
svc.shutdown().await;
assert_eq!(tallied(&svc), 3);
}
#[tokio::test]
async fn a_read_that_fails_outright_reports_what_the_tally_last_held() {
let dir = tempfile::tempdir().unwrap();
let log = dir.path().join("log.jsonl");
let svc = GithubCountersService::with_log_path(Some(log.clone()));
append(&log, &gh_line(svc.started_at, 1, &["pr", "list"]));
append(&log, &gh_line(svc.started_at, 2, &["pr", "view"]));
assert_eq!(
svc.status().await.summary,
"2 GitHub API call(s) since start"
);
let (counts, current) = counts_since_boot(&svc.tally, || panic!("induced")).await;
assert!(!current);
assert_eq!(counts.total(), 2);
append(&log, &gh_line(svc.started_at, 3, &["pr", "list"]));
assert_eq!(
svc.status().await.summary,
"3 GitHub API call(s) since start"
);
}
#[tokio::test]
async fn the_logger_keeps_the_tally_current_and_stops_on_shutdown() {
let dir = tempfile::tempdir().unwrap();
let log = dir.path().join("log.jsonl");
let svc = GithubCountersService::with_log_path(Some(log.clone()));
svc.start_counter_logger_with(Duration::from_millis(1), Duration::from_millis(5));
svc.start_counter_logger_with(Duration::from_millis(1), Duration::from_millis(5));
append(&log, &gh_line(svc.started_at, 1, &["pr", "list"]));
let deadline = Instant::now() + Duration::from_secs(10);
while tallied(&svc) == 0 {
assert!(
Instant::now() < deadline,
"the logger never read the record"
);
tokio::time::sleep(Duration::from_millis(10)).await;
}
svc.shutdown().await;
assert!(svc.logger.lock().unwrap().is_none());
assert!(svc.stopping.is_cancelled());
}
}