overbrainer 0.9.0

//! What differs between the front ends that run the pipeline and training flows:
//! where their events go, what interrupts them, and where their summary lines go.

use std::future::Future;
use std::io;
use std::pin::{Pin, pin};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context as TaskContext, Poll, Waker};
use std::time::Duration;

use anyhow::Context;
use tokio::signal::unix::{SignalKind, signal};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

use crate::events::{EventBus, Observer};

/// A front end of the flows in `cli::data` and `cli::train`.
pub(crate) enum Frontend {
    /// The command line: a new bus rendered to stderr and seen by the observer
    /// if any, Ctrl-C (SIGINT), stdout.
    Cli(Option<Arc<dyn Observer>>),
    /// The terminal UI: the task's own bus, the token that detaches it, the flag
    /// that abandons a Runpod provisioning, and a sink into the app.
    Tui {
        /// The task's bus, which the app forwards.
        bus: EventBus,
        /// Cancelled to interrupt the task's flow (a raced watch detaches).
        detach: CancellationToken,
        /// Set to make Runpod provisioning stop and delete its pod.
        abandon: Arc<AtomicBool>,
        /// Receives what the command line would print.
        report: Arc<dyn Fn(Report) + Send + Sync>,
    },
}

/// What a flow tells its front end besides events.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Report {
    /// A line the command line prints on stdout.
    Line(String),
    /// A training run was created, with this ID (the command line says nothing).
    RunCreated(String),
    /// The training run ended stopped, with its snapshot at this step (the
    /// command line says it in a line).
    RunStopped(u64),
}

impl Frontend {
    /// The bus a flow publishes on, with what renders it.
    pub(crate) fn open_bus(&self) -> BusGuard {
        match self {
            Self::Cli(observer) => {
                let bus = EventBus::observed(EventBus::CAPACITY, observer.clone());
                let renderer = tokio::spawn(super::progress::render(bus.subscribe()));
                BusGuard {
                    bus,
                    renderer: Some(renderer),
                    cancel: None,
                }
            },
            Self::Tui { bus, .. } => {
                // The same renderer as the command line turns events into tracing
                // lines; the TUI's log buffer captures them for its Logs view. The bus
                // is shared here (the app and the flow hold clones), so it never closes
                // on this guard alone: `close` cancels the renderer, which first drains
                // the queued events so the last progress lines still reach the buffer.
                let cancel = CancellationToken::new();
                let renderer = tokio::spawn(super::progress::render_until(
                    bus.subscribe(),
                    cancel.clone(),
                ));
                BusGuard {
                    bus: bus.clone(),
                    renderer: Some(renderer),
                    cancel: Some(cancel),
                }
            },
        }
    }

    /// What stops a raced flow and is noted by a shielded one, caught from now on.
    pub(crate) fn interrupt(&self) -> Interrupt {
        match self {
            Self::Cli(_) => Interrupt::catch(),
            Self::Tui { detach, .. } => {
                let detach = detach.clone();
                Interrupt::on(Box::pin(async move {
                    detach.cancelled_owned().await;
                    Ok(())
                }))
            },
        }
    }

    /// The flag Runpod provisioning checks between its steps, set from now on.
    ///
    /// # Errors
    ///
    /// Returns an error when the Ctrl-C handler cannot be installed.
    pub(crate) fn provisioning_flag(&self) -> anyhow::Result<Flag> {
        match self {
            Self::Cli(_) => {
                let mut sigint = signal(SignalKind::interrupt()).context("cannot catch Ctrl-C")?;
                let interrupted = Arc::new(AtomicBool::new(false));
                let seen = Arc::clone(&interrupted);
                let watcher = tokio::spawn(async move {
                    if sigint.recv().await.is_some() {
                        seen.store(true, Ordering::SeqCst);
                    }
                });
                Ok(Flag {
                    interrupted,
                    watcher: Some(watcher),
                })
            },
            Self::Tui { abandon, .. } => Ok(Flag {
                interrupted: Arc::clone(abandon),
                watcher: None,
            }),
        }
    }

    /// Emits a summary line: the command line prints it on stdout, the TUI gets
    /// it as a [`Report::Line`].
    pub(crate) fn line(&self, line: &str) {
        match self {
            Self::Cli(_) => println!("{line}"),
            Self::Tui { report, .. } => report(Report::Line(line.to_string())),
        }
    }

    /// Says that the training run ended stopped, its snapshot at `step`.
    pub(crate) fn run_stopped(&self, step: u64) {
        match self {
            Self::Cli(_) => {},
            Self::Tui { report, .. } => report(Report::RunStopped(step)),
        }
    }

    /// Says that training run `id` was created: the TUI binds its task to it.
    pub(crate) fn run_created(&self, id: &str) {
        match self {
            Self::Cli(_) => {},
            Self::Tui { report, .. } => report(Report::RunCreated(id.to_string())),
        }
    }
}

/// How long `close` waits for a cancelled renderer to drain its queued events.
const RENDER_GRACE: Duration = Duration::from_secs(2);

/// The bus of a flow, and the task rendering it if any.
pub(crate) struct BusGuard {
    /// Where the flow publishes.
    pub(crate) bus: EventBus,
    renderer: Option<JoinHandle<()>>,
    /// Set when the bus is shared beyond this guard (the TUI): the renderer only stops
    /// on this signal, since the bus never closes here.
    cancel: Option<CancellationToken>,
}

impl BusGuard {
    /// Drops the bus. When the bus is this guard's alone, waits for the renderer to
    /// show what is left; when it is shared (the TUI), signals the renderer to drain
    /// the queued events and stop, then waits for it, bounded by [`RENDER_GRACE`].
    pub(crate) async fn close(mut self) {
        let renderer = self.renderer.take();
        let cancel = self.cancel.take();
        // Drops the bus; with the renderer taken, the drop leaves it running.
        drop(self);
        let Some(mut renderer) = renderer else {
            return;
        };
        match cancel {
            Some(cancel) => {
                cancel.cancel();
                // Await the drain, but do not leave it running on the shared bus past
                // the grace period: abort it and reap the handle.
                if tokio::time::timeout(RENDER_GRACE, &mut renderer)
                    .await
                    .is_err()
                {
                    renderer.abort();
                    renderer.await.ok();
                }
            },
            None => {
                renderer.await.ok();
            },
        }
    }
}

impl Drop for BusGuard {
    /// Stops the renderer of a guard dropped without [`BusGuard::close`] (a flow
    /// interrupted before it closed its bus): a shared bus's renderer drains what is
    /// queued and ends, any other is aborted, so no renderer outlives its flow.
    fn drop(&mut self) {
        let Some(renderer) = self.renderer.take() else {
            return;
        };
        match self.cancel.take() {
            Some(cancel) => cancel.cancel(),
            None => renderer.abort(),
        }
    }
}

/// The interrupted flag of Runpod provisioning, and the task that sets it if any.
pub(crate) struct Flag {
    /// Set once the flow must stop provisioning.
    pub(crate) interrupted: Arc<AtomicBool>,
    watcher: Option<JoinHandle<()>>,
}

impl Flag {
    /// Stops watching for the interruption.
    pub(crate) fn close(self) {
        drop(self);
    }
}

impl Drop for Flag {
    /// Stops the watcher, so a flag dropped without [`Flag::close`] never leaves
    /// it running.
    fn drop(&mut self) {
        if let Some(watcher) = self.watcher.take() {
            watcher.abort();
        }
    }
}

type Signal = Pin<Box<dyn Future<Output = io::Result<()>> + Send>>;

/// An interruption (Ctrl-C on the command line), caught from its creation on so
/// that it no longer stops the process: a flow either runs to its end regardless
/// ([`Interrupt::shield`]) or stops at the first interruption ([`Interrupt::race`]).
pub(crate) enum Interrupt {
    /// Waiting for the interruption.
    Listening(Signal),
    /// The interruption came.
    Caught,
    /// The interruption cannot be caught: Ctrl-C stops the process as usual.
    Off,
}

impl Interrupt {
    /// Interrupted when `signal` resolves.
    fn on(signal: Signal) -> Self {
        Self::Listening(signal)
    }

    /// Starts catching Ctrl-C now.
    fn catch() -> Self {
        let mut signal: Signal = Box::pin(tokio::signal::ctrl_c());
        // The handler is installed on the first poll; a later poll registers the
        // real waker.
        let mut cx = TaskContext::from_waker(Waker::noop());
        match signal.as_mut().poll(&mut cx) {
            Poll::Pending => Self::Listening(signal),
            Poll::Ready(result) => Self::after(result),
        }
    }

    fn after(result: io::Result<()>) -> Self {
        match result {
            Ok(()) => Self::Caught,
            Err(error) => {
                tracing::warn!("cannot catch Ctrl-C: {error}");
                Self::Off
            },
        }
    }

    /// Whether the interruption came.
    pub(crate) fn caught(&self) -> bool {
        matches!(self, Self::Caught)
    }

    /// Runs `flow` to its end, noting an interruption meanwhile.
    pub(crate) async fn shield<T>(&mut self, flow: impl Future<Output = T>) -> T {
        let mut flow = pin!(flow);
        let result = match self {
            Self::Listening(signal) => tokio::select! {
                output = &mut flow => return output,
                result = signal.as_mut() => result,
            },
            Self::Caught | Self::Off => return flow.await,
        };
        *self = Self::after(result);
        flow.await
    }

    /// Runs `flow`, or `None` when the interruption comes first (or already came).
    pub(crate) async fn race<T>(&mut self, flow: impl Future<Output = T>) -> Option<T> {
        let mut flow = pin!(flow);
        let result = match self {
            Self::Listening(signal) => tokio::select! {
                output = &mut flow => return Some(output),
                result = signal.as_mut() => result,
            },
            Self::Caught => return None,
            Self::Off => return Some(flow.await),
        };
        *self = Self::after(result);
        if self.caught() {
            None
        } else {
            Some(flow.await)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tokio::signal::unix::{SignalKind, signal};

    use super::*;
    use crate::events::{Event, Stage};

    /// Sends SIGINT to this process, as Ctrl-C does.
    async fn ctrl_c() -> io::Result<()> {
        let status = tokio::process::Command::new("kill")
            .args(["-INT", &std::process::id().to_string()])
            .status()
            .await?;
        if status.success() {
            Ok(())
        } else {
            Err(io::Error::other("kill failed"))
        }
    }

    #[tokio::test]
    async fn ctrl_c_never_stops_a_shielded_flow_and_stops_a_raced_one()
    -> Result<(), Box<dyn std::error::Error>> {
        let _signals = crate::test_support::SIGNALS.lock().await;
        let limit = Duration::from_secs(10);
        let mut interrupt = Frontend::Cli(None).interrupt();
        assert!(matches!(interrupt, Interrupt::Listening(_)));
        // Ctrl-C arrives while the flow waits: it still ends, and Ctrl-C is noted.
        let mut received = signal(SignalKind::interrupt())?;
        let shielded = async {
            ctrl_c().await?;
            received.recv().await;
            tokio::time::sleep(Duration::from_millis(50)).await;
            Ok::<_, io::Error>("ended")
        };
        let ended = tokio::time::timeout(limit, interrupt.shield(shielded)).await??;
        assert_eq!(ended, "ended");
        assert!(interrupt.caught());
        // Once Ctrl-C was pressed, a race does not even start its flow.
        assert_eq!(interrupt.race(std::future::ready(())).await, None);

        let mut interrupt = Frontend::Cli(None).interrupt();
        let raced = async {
            ctrl_c().await?;
            std::future::pending::<()>().await;
            Ok::<_, io::Error>(())
        };
        let raced = tokio::time::timeout(limit, interrupt.race(raced)).await?;
        assert!(raced.is_none());
        assert!(interrupt.caught());
        Ok(())
    }

    fn tui(detach: &CancellationToken) -> Frontend {
        Frontend::Tui {
            bus: EventBus::with_capacity(8),
            detach: detach.clone(),
            abandon: Arc::new(AtomicBool::new(false)),
            report: Arc::new(|_| {}),
        }
    }

    #[tokio::test]
    async fn a_detach_token_never_stops_a_shielded_flow_and_stops_a_raced_one()
    -> Result<(), Box<dyn std::error::Error>> {
        let limit = Duration::from_secs(10);
        let token = CancellationToken::new();
        let mut interrupt = tui(&token).interrupt();
        assert!(matches!(interrupt, Interrupt::Listening(_)));
        let shielded = async {
            token.cancel();
            tokio::time::sleep(Duration::from_millis(50)).await;
            "ended"
        };
        assert_eq!(
            tokio::time::timeout(limit, interrupt.shield(shielded)).await?,
            "ended"
        );
        assert!(interrupt.caught());
        assert_eq!(interrupt.race(std::future::ready(())).await, None);

        let token = CancellationToken::new();
        let mut interrupt = tui(&token).interrupt();
        let raced = async {
            token.cancel();
            std::future::pending::<()>().await;
        };
        assert!(
            tokio::time::timeout(limit, interrupt.race(raced))
                .await?
                .is_none()
        );
        assert!(interrupt.caught());
        Ok(())
    }

    #[tokio::test]
    async fn a_tui_front_end_shares_its_bus_flag_and_reports() -> anyhow::Result<()> {
        let lines = Arc::new(std::sync::Mutex::new(Vec::new()));
        let seen = Arc::clone(&lines);
        let abandon = Arc::new(AtomicBool::new(false));
        let front = Frontend::Tui {
            bus: EventBus::with_capacity(8),
            detach: CancellationToken::new(),
            abandon: Arc::clone(&abandon),
            report: Arc::new(move |report| {
                if let Ok(mut lines) = seen.lock() {
                    lines.push(report);
                }
            }),
        };
        let Frontend::Tui { bus, .. } = &front else {
            anyhow::bail!("not a TUI front end");
        };
        let mut events = bus.subscribe();
        let guard = front.open_bus();
        guard.bus.publish(Event::StageStarted {
            stage: Stage::Split,
            total: 2,
        });
        assert!(events.try_recv().is_ok(), "the task's own bus");
        guard.close().await;
        let flag = front.provisioning_flag()?;
        abandon.store(true, Ordering::SeqCst);
        assert!(flag.interrupted.load(Ordering::SeqCst));
        flag.close();
        front.line("split: 1 train, 1 eval, 0 excluded, 0 orphaned");
        let reported = lines.lock().map(|lines| lines.clone()).unwrap_or_default();
        assert_eq!(
            reported,
            [Report::Line(
                "split: 1 train, 1 eval, 0 excluded, 0 orphaned".into()
            )]
        );
        Ok(())
    }

    #[tokio::test]
    async fn a_cli_bus_is_rendered_until_it_is_closed() -> Result<(), tokio::time::error::Elapsed> {
        let guard = Frontend::Cli(None).open_bus();
        let mut extra = guard.bus.subscribe();
        guard.bus.publish(Event::StageStarted {
            stage: Stage::Split,
            total: 1,
        });
        assert!(extra.try_recv().is_ok());
        tokio::time::timeout(Duration::from_secs(10), guard.close()).await
    }

    /// Waits until `done` holds, yielding to the runtime, for at most ten seconds.
    async fn eventually(done: impl Fn() -> bool) -> Result<(), tokio::time::error::Elapsed> {
        tokio::time::timeout(Duration::from_secs(10), async {
            while !done() {
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        })
        .await
    }

    #[tokio::test]
    async fn a_cli_flag_dropped_without_close_stops_its_watcher() -> anyhow::Result<()> {
        // No other test may send Ctrl-C meanwhile: it would end the watcher too.
        let _signals = crate::test_support::SIGNALS.lock().await;
        let flag = Frontend::Cli(None).provisioning_flag()?;
        let interrupted = Arc::clone(&flag.interrupted);
        assert_eq!(
            Arc::strong_count(&interrupted),
            3,
            "the flag and its watcher"
        );
        drop(flag);
        // The watcher holds the other clone until it ends.
        eventually(|| Arc::strong_count(&interrupted) == 1).await?;
        Ok(())
    }

    #[tokio::test]
    async fn a_bus_guard_dropped_without_close_stops_its_renderer() -> anyhow::Result<()> {
        // The command line: a clone keeps the bus open, so only the guard ends it.
        let guard = Frontend::Cli(None).open_bus();
        let bus = guard.bus.clone();
        assert_eq!(bus.receiver_count(), 1, "the renderer");
        drop(guard);
        eventually(|| bus.receiver_count() == 0).await?;

        // The TUI: the task's bus outlives the guard.
        let front = tui(&CancellationToken::new());
        let Frontend::Tui { bus, .. } = &front else {
            anyhow::bail!("not a TUI front end");
        };
        let guard = front.open_bus();
        assert_eq!(bus.receiver_count(), 1, "the renderer");
        drop(guard);
        eventually(|| bus.receiver_count() == 0).await?;
        Ok(())
    }

    #[tokio::test]
    async fn a_cli_provisioning_flag_starts_clear() -> anyhow::Result<()> {
        let flag = Frontend::Cli(None).provisioning_flag()?;
        assert!(!flag.interrupted.load(Ordering::SeqCst));
        flag.close();
        Ok(())
    }
}