use std::io::{IsTerminal, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
static SPINNER_ACTIVE: AtomicBool = AtomicBool::new(false);
pub struct SpinnerSafeWriter;
impl std::io::Write for SpinnerSafeWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let mut err = std::io::stderr().lock();
if SPINNER_ACTIVE.load(Ordering::Relaxed) {
err.write_all(b"\r\x1b[K")?;
}
err.write_all(buf)?;
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
std::io::stderr().lock().flush()
}
}
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for SpinnerSafeWriter {
type Writer = SpinnerSafeWriter;
fn make_writer(&'a self) -> Self::Writer {
SpinnerSafeWriter
}
}
pub(crate) struct Spinner {
stop: Arc<AtomicBool>,
phase: Arc<Mutex<String>>,
handle: Option<std::thread::JoinHandle<()>>,
tty: bool,
}
impl Spinner {
pub(crate) fn start(label: &'static str) -> Option<Self> {
let is_tty = std::io::stderr().is_terminal();
if !is_tty && std::env::var_os("CARTOG_PROGRESS").is_none() {
return None;
}
let stop = Arc::new(AtomicBool::new(false));
let phase = Arc::new(Mutex::new(label.to_string()));
let stop_clone = Arc::clone(&stop);
let phase_clone = Arc::clone(&phase);
if is_tty {
SPINNER_ACTIVE.store(true, Ordering::Relaxed);
}
let handle = std::thread::spawn(move || {
if is_tty {
Self::run_tty(&stop_clone, &phase_clone);
} else {
Self::run_plain(&stop_clone, &phase_clone);
}
});
Some(Self {
stop,
phase,
handle: Some(handle),
tty: is_tty,
})
}
pub(crate) fn set_phase(&self, phase: impl Into<String>) {
if let Ok(mut p) = self.phase.lock() {
*p = phase.into();
}
}
fn run_tty(stop: &AtomicBool, phase: &Mutex<String>) {
const FRAMES: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
let mut i = 0usize;
let start = std::time::Instant::now();
while !stop.load(Ordering::Relaxed) {
let elapsed = start.elapsed().as_secs();
let label = phase.lock().map(|p| p.clone()).unwrap_or_default();
let mut err = std::io::stderr().lock();
let _ = write!(err, "\r\x1b[K{} {label} ({elapsed}s)", FRAMES[i]);
let _ = err.flush();
drop(err);
i = (i + 1) % FRAMES.len();
std::thread::sleep(Duration::from_millis(100));
}
let mut err = std::io::stderr().lock();
let _ = write!(err, "\r\x1b[K");
let _ = err.flush();
}
fn run_plain(stop: &AtomicBool, phase: &Mutex<String>) {
let start = std::time::Instant::now();
let mut last_label = String::new();
let mut last_emit = std::time::Instant::now();
while !stop.load(Ordering::Relaxed) {
let label = phase.lock().map(|p| p.clone()).unwrap_or_default();
let changed = label != last_label;
if changed || last_emit.elapsed() >= Duration::from_secs(5) {
let elapsed = start.elapsed().as_secs();
eprintln!(" {label}… ({elapsed}s)");
last_label = label;
last_emit = std::time::Instant::now();
}
std::thread::sleep(Duration::from_millis(200));
}
}
pub(crate) fn stop(mut self) {
self.stop.store(true, Ordering::Relaxed);
if let Some(h) = self.handle.take() {
let _ = h.join();
}
}
}
impl Drop for Spinner {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
if let Some(h) = self.handle.take() {
let _ = h.join();
}
if self.tty {
SPINNER_ACTIVE.store(false, Ordering::Relaxed);
}
}
}
pub(crate) fn install_cancel_probe() -> impl Fn() -> bool + Send + Sync + 'static {
let interrupted = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&interrupted);
let _ = ctrlc::set_handler(move || flag.store(true, Ordering::SeqCst));
move || interrupted.load(Ordering::SeqCst)
}
pub(crate) fn capitalize_phase(label: String) -> String {
let mut chars = label.chars();
match chars.next() {
Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
None => label,
}
}
pub(crate) fn spinner_callback<U>(
spinner: &Option<Arc<Spinner>>,
label_of: fn(&U) -> String,
) -> Option<impl Fn(U)> {
spinner.as_ref().map(|s| {
let s = Arc::clone(s);
move |u: U| s.set_phase(capitalize_phase(label_of(&u)))
})
}
pub(crate) fn stop_spinner(spinner: Option<Arc<Spinner>>) {
if let Some(s) = spinner.and_then(Arc::into_inner) {
s.stop();
}
}
#[cfg(test)]
mod tests {
use super::*;
use cartog_indexer as indexer;
use cartog_rag as rag;
#[test]
fn capitalized_index_phase_labels() {
use indexer::ProgressUpdate as U;
let cap = |u: &U| capitalize_phase(u.label());
assert_eq!(cap(&U::Walking), "Scanning files");
assert_eq!(cap(&U::Parsing { done: 0, total: 12 }), "Parsing 12 files");
assert_eq!(
cap(&U::Parsing { done: 4, total: 12 }),
"Parsing 4/12 files"
);
assert_eq!(cap(&U::Storing { done: 0, total: 5 }), "Storing 5 files");
assert_eq!(cap(&U::Storing { done: 3, total: 5 }), "Storing 3/5 files");
assert_eq!(
cap(&U::ResolvingLsp { done: 0, total: 9 }),
"Resolving 9 edges with LSP"
);
assert_eq!(
cap(&U::ResolvingLsp { done: 3, total: 9 }),
"Resolving 3/9 edges with LSP"
);
}
#[test]
fn capitalized_rag_phase_labels() {
use rag::indexer::ProgressUpdate as U;
let cap = |u: &U| capitalize_phase(u.label());
assert_eq!(cap(&U::Preparing), "Preparing");
assert_eq!(
cap(&U::Embedding {
processed: 64,
total: 256
}),
"Embedding 64/256"
);
assert_eq!(cap(&U::Storing), "Storing embeddings");
}
#[test]
fn capitalize_phase_handles_empty() {
assert_eq!(capitalize_phase(String::new()), "");
}
}