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};
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)
}
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));
}
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),
}));
}
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));
}
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);
}
}
}
}