bot-forge 1.0.2

//! Process-wide cancellation routing for one active application run.
//!
//! The signal handler records cancellation even when no run is active. [`CancellationSession`]
//! binds subsequent signals to one shared [`CancellationToken`]; dropping the session detaches the
//! token without uninstalling the process handler.

use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::{Duration, Instant};

use crate::error::ForgeError;

static ACTIVE_CANCELLATION: OnceLock<Mutex<Weak<AtomicU8>>> = OnceLock::new();
static PENDING_CANCELLATION: AtomicU8 = AtomicU8::new(0);

/// Thread-safe monotonic cancellation state shared by schedulers and side-effect boundaries.
///
/// Level one requests graceful cancellation; level two and above request forced termination.
#[derive(Clone, Default)]
pub(crate) struct CancellationToken(Arc<AtomicU8>);

impl CancellationToken {
    /// Request graceful cancellation without downgrading a forced request.
    #[cfg(test)]
    pub(crate) fn cancel(&self) {
        self.0.fetch_max(1, Ordering::SeqCst);
    }

    /// Return whether graceful or forced cancellation has been requested.
    pub(crate) fn is_cancelled(&self) -> bool {
        self.0.load(Ordering::SeqCst) > 0
    }

    /// Return whether forced cancellation has been requested.
    pub(crate) fn is_forced(&self) -> bool {
        self.0.load(Ordering::SeqCst) > 1
    }
}

/// Install the process-wide signal handler used by cancellation sessions.
///
/// # Errors
///
/// Returns [`ctrlc::Error`] when the platform handler cannot be registered or this process has
/// already installed one.
pub(crate) fn install_handler() -> Result<(), ctrlc::Error> {
    ctrlc::set_handler(|| {
        if let Some(state) = ACTIVE_CANCELLATION
            .get()
            .and_then(|active| active.lock().ok())
            .and_then(|active| active.upgrade())
        {
            state.fetch_add(1, Ordering::SeqCst);
        } else {
            PENDING_CANCELLATION.fetch_add(1, Ordering::SeqCst);
        }
    })
}

/// Return whether the active or pending process state contains a cancellation request.
pub(crate) fn requested() -> bool {
    active_level() > 0
}

/// Return whether the active or pending process state contains a forced request.
pub(crate) fn forced() -> bool {
    active_level() > 1
}

/// Wait for `duration`, polling process cancellation at bounded intervals.
///
/// # Errors
///
/// Returns [`ForgeError::Command`] when cancellation is observed before the delay completes.
pub(crate) fn delay(duration: Duration, label: &str) -> Result<(), ForgeError> {
    delay_while(duration, label, requested)
}

fn delay_while(
    duration: Duration,
    label: &str,
    cancelled: impl Fn() -> bool,
) -> Result<(), ForgeError> {
    let started = Instant::now();
    while started.elapsed() < duration {
        if cancelled() {
            return Err(ForgeError::Command(format!(
                "cancelled while waiting for {label}"
            )));
        }
        std::thread::sleep(
            duration
                .saturating_sub(started.elapsed())
                .min(Duration::from_millis(50)),
        );
    }
    Ok(())
}

fn active_level() -> u8 {
    let active = ACTIVE_CANCELLATION
        .get()
        .and_then(|active| active.lock().ok())
        .and_then(|active| active.upgrade())
        .map_or(0, |state| state.load(Ordering::SeqCst));
    active.max(PENDING_CANCELLATION.load(Ordering::SeqCst))
}

/// RAII scope that binds process cancellation signals to one run token.
pub(crate) struct CancellationSession {
    token: CancellationToken,
}

impl CancellationSession {
    /// Begin a fresh scope and clear signals pending from an earlier scope.
    pub(crate) fn begin() -> Self {
        PENDING_CANCELLATION.store(0, Ordering::SeqCst);
        let token = CancellationToken::default();
        let active = ACTIVE_CANCELLATION.get_or_init(|| Mutex::new(Weak::new()));
        if let Ok(mut current) = active.lock() {
            *current = Arc::downgrade(&token.0);
        }
        Self { token }
    }

    /// Clone the token bound to this session.
    pub(crate) fn token(&self) -> CancellationToken {
        self.token.clone()
    }
}

impl Drop for CancellationSession {
    fn drop(&mut self) {
        if let Some(active) = ACTIVE_CANCELLATION.get()
            && let Ok(mut current) = active.lock()
        {
            *current = Weak::new();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::{Duration, Instant};

    use crate::cancellation::delay_while;

    #[test]
    fn cancellable_delay_stops_promptly() {
        let polls = AtomicUsize::new(0);
        let started = Instant::now();
        let error = delay_while(Duration::from_secs(1), "retry backoff", || {
            polls.fetch_add(1, Ordering::SeqCst) >= 1
        })
        .unwrap_err();
        assert!(error.to_string().contains("cancelled"));
        assert!(started.elapsed() < Duration::from_millis(150));
    }
}