tastty-driver 0.1.0

Terminal automation driver built on tastty
//! Async wait future for [`Session::wait_async`](crate::Session::wait_async).
//!
//! Built on `std` primitives only (no tokio); runtime-agnostic.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use tastty::Terminal;

use crate::wait::blocking::{snapshot, try_wait};
use crate::wait::condition::WaitCondition;
use crate::wait::error::WaitError;
use crate::wait::outcome::WaitOutcome;
use crate::wait::output::OutputNotifier;
use crate::wait::probe::{CompiledCondition, Probe, probe};

/// Future returned by [`Session::wait_async`](crate::Session::wait_async).
///
/// Holds an [`Arc`] of the session's [`OutputNotifier`] alongside a slot
/// per wake source for the monotonic key under which its
/// [`Waker`](std::task::Waker) is registered. On each poll the future
/// registers (or refreshes) its waker on the output notifier (woken on
/// every parser tick, including the post-exit kick the reaper issues),
/// then probes its condition; if not yet matched, it also registers on
/// the notifier's timer thread (woken at the timeout deadline, or sooner
/// if a [`WaitCondition::stable`] sub-condition's settle window is
/// closer).
///
/// Re-poll order is fixed: register-tick -> probe -> deadline -> exit ->
/// register-timer, with the tick registration first so a parser tick
/// firing after `register_tick` and before `Poll::Pending` is returned
/// still produces a re-poll. If the originating session has already
/// shut the notifier down, registration reports closure and the future
/// resolves with [`WaitError::SessionClosed`] instead of parking forever.
/// Dropping the future removes its waker entries from the notifier so
/// later wakes do not visit a slot whose future is gone.
pub(crate) struct WaitFuture {
    inner: Option<Inner>,
}

struct Inner {
    terminal: Arc<Terminal>,
    output: Arc<OutputNotifier>,
    condition: WaitCondition,
    compiled: CompiledCondition,
    exit_short_circuits: bool,
    start: Instant,
    deadline: Instant,
    tick_key: Option<u64>,
    timer_key: Option<u64>,
    initial_compile_error: Option<WaitError>,
}

impl WaitFuture {
    pub(crate) fn new(
        terminal: Arc<Terminal>,
        output: Arc<OutputNotifier>,
        condition: WaitCondition,
        timeout: Duration,
    ) -> Self {
        let start = Instant::now();
        let deadline = start + timeout;
        let (compiled, exit_short_circuits, initial_compile_error) =
            match CompiledCondition::compile(&condition) {
                Ok(compiled) => {
                    let short = compiled.is_exit_or_stable();
                    (compiled, short, None)
                }
                // Compilation only fails for invalid regex / nested any_of.
                // Surface the error on the first poll so the constructor
                // stays infallible (matching the synchronous `wait` shape
                // that returns the error rather than panicking at the call
                // site). Use a placeholder compiled condition that is
                // never probed because `initial_compile_error` short-
                // circuits the first poll.
                Err(error) => (
                    CompiledCondition::compile(&WaitCondition::exit())
                        .expect("exit condition compiles"),
                    true,
                    Some(error),
                ),
            };
        Self {
            inner: Some(Inner {
                terminal,
                output,
                condition,
                compiled,
                exit_short_circuits,
                start,
                deadline,
                tick_key: None,
                timer_key: None,
                initial_compile_error,
            }),
        }
    }
}

impl Future for WaitFuture {
    type Output = Result<WaitOutcome, WaitError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        let inner = this
            .inner
            .as_mut()
            .expect("WaitFuture polled after completion");

        if let Some(error) = inner.initial_compile_error.take() {
            this.inner = None;
            return Poll::Ready(Err(error));
        }

        // Register the tick waker BEFORE probing so a parser tick that
        // fires during this poll guarantees a re-poll. The notifier
        // drains all tick wakers atomically, so a fresh registration is
        // either preserved (no tick yet) or drained-and-woken (tick
        // happened) but never silently dropped.
        let tick_registration_closed = !inner.output.register_tick(&mut inner.tick_key, cx.waker());

        let probe_result = inner
            .terminal
            .with_screen(|screen| probe(&inner.terminal, screen, &mut inner.compiled));
        match probe_result {
            Ok(Probe::Matched(wait_match)) => {
                let snap = snapshot(&inner.terminal);
                let exit_status = try_wait(&inner.terminal).ok().flatten();
                let outcome = WaitOutcome {
                    snapshot: snap,
                    elapsed: inner.start.elapsed(),
                    wait_match,
                    exit_status,
                };
                this.inner = None;
                return Poll::Ready(Ok(outcome));
            }
            Ok(Probe::NotYet) => {}
            Err(error) => {
                this.inner = None;
                return Poll::Ready(Err(error));
            }
        }

        let now = Instant::now();
        if now >= inner.deadline {
            let condition = inner.condition.clone();
            let elapsed = inner.start.elapsed();
            let snap = snapshot(&inner.terminal);
            this.inner = None;
            return Poll::Ready(Err(WaitError::Timeout {
                condition,
                elapsed,
                snapshot: Box::new(snap),
            }));
        }

        // Exit and Stable are passively-satisfiable from a dead process,
        // so a clean exit during their wait does not abort. Every other
        // condition is content- or cursor-driven, so an observed exit
        // without a match is terminal. The reaper kicks the output
        // notifier after each exit transition, so a future that registers
        // its tick waker (above) before this `try_wait` cannot miss the
        // wake even if exit transitions in the gap.
        if !inner.exit_short_circuits
            && let Ok(Some(exit_status)) = try_wait(&inner.terminal)
        {
            let condition = inner.condition.clone();
            let snap = snapshot(&inner.terminal);
            this.inner = None;
            return Poll::Ready(Err(WaitError::ProcessExitedBeforeMatch {
                condition,
                exit_status,
                snapshot: Box::new(snap),
            }));
        }

        if tick_registration_closed {
            let error = inner.session_closed_error();
            this.inner = None;
            return Poll::Ready(Err(error));
        }

        // Schedule a deadline wake at the earlier of overall timeout and
        // the next stable-settle deadline (when one is being tracked).
        // Without this, a Stable condition with a quiet child would never
        // wake, and a non-Stable condition with no further output would
        // miss its timeout.
        let timer_deadline = match inner.compiled.next_stable_deadline() {
            Some(stable) => inner.deadline.min(stable),
            None => inner.deadline,
        };
        if !inner
            .output
            .register_deadline(&mut inner.timer_key, timer_deadline, cx.waker())
        {
            let error = inner.session_closed_error();
            this.inner = None;
            return Poll::Ready(Err(error));
        }

        Poll::Pending
    }
}

impl Inner {
    fn session_closed_error(&self) -> WaitError {
        WaitError::SessionClosed {
            condition: self.condition.clone(),
            elapsed: self.start.elapsed(),
            snapshot: Box::new(snapshot(&self.terminal)),
        }
    }
}

impl Drop for WaitFuture {
    fn drop(&mut self) {
        if let Some(inner) = self.inner.take() {
            if let Some(key) = inner.tick_key {
                inner.output.unregister_tick(key);
            }
            if let Some(key) = inner.timer_key {
                inner.output.unregister_deadline(key);
            }
        }
    }
}