use crate::batch::{FrameStream, batch_frames};
use crate::{
BatchOptions, BatchStream, Executor, ExecutorError, Indexer, Input, InputCompletion, Lister,
Output, StreamExt, Writer,
};
use alloc::{boxed::Box, vec, vec::Vec};
use core::{
fmt,
future::{Future, poll_fn},
marker::PhantomData,
pin::Pin,
task::Poll,
};
use std::{
ffi::OsString,
io::{self, Cursor},
process::Stdio,
};
use tokio::{process::Child, sync::watch};
#[derive(Debug)]
pub struct PipelineError {
pub stage: usize,
pub program: OsString,
pub error: ExecutorError,
}
impl fmt::Display for PipelineError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"Pipeline stage {} ({}): {}",
self.stage,
self.program.to_string_lossy(),
self.error
)
}
}
impl core::error::Error for PipelineError {
fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
Some(&self.error)
}
}
pub type PipelineStream = BatchStream<PipelineError>;
mod sealed {
pub trait Sealed {}
}
pub trait PipelineProgram: sealed::Sealed + Into<PipelineStage> {}
pub trait GraphProducer: PipelineProgram {}
pub trait GraphConsumer: PipelineProgram {}
macro_rules! programs {
(producer: $($producer:ty),*; consumer: $($consumer:ty),*; both: $($both:ty),*) => {
$(impl sealed::Sealed for $producer {}
impl PipelineProgram for $producer {}
impl GraphProducer for $producer {})*
$(impl sealed::Sealed for $consumer {}
impl PipelineProgram for $consumer {}
impl GraphConsumer for $consumer {})*
$(impl sealed::Sealed for $both {}
impl PipelineProgram for $both {}
impl GraphProducer for $both {}
impl GraphConsumer for $both {})*
};
}
programs! {
producer: crate::Adapter, crate::Emitter, crate::Fetcher, Lister, crate::Reader;
consumer: Writer, Indexer;
both: crate::Matcher, crate::Reasoner
}
#[derive(Debug)]
pub struct Pipeline<P> {
stages: Vec<PipelineStage>,
batching: Option<BatchOptions>,
tail: PhantomData<fn() -> P>,
}
impl<P: PipelineProgram> Pipeline<P> {
pub fn new(program: P) -> Self {
Self {
stages: vec![program.into()],
batching: None,
tail: PhantomData,
}
}
}
impl<P: GraphProducer> Pipeline<P> {
pub fn pipe<N: GraphConsumer>(mut self, program: N) -> Pipeline<N> {
self.stages.push(program.into());
Pipeline {
stages: self.stages,
batching: self.batching,
tail: PhantomData,
}
}
#[must_use]
pub fn with_batching(mut self, options: BatchOptions) -> Self {
self.batching = Some(options);
self
}
pub async fn execute(self) -> Result<PipelineStream, PipelineError> {
let batching = self
.batching
.unwrap_or(self.stages.last().unwrap().batching);
Ok(batch_frames(self.execute_frames().await?, batching))
}
async fn execute_frames(self) -> Result<FrameStream<PipelineError>, PipelineError> {
let mut running = start(self.stages, true).await?;
let mut lines = running.lines.take();
let tail = running.tail.clone();
Ok(Box::pin(async_stream::try_stream! {
let mut completed = false;
if let Some(ref mut lines) = lines {
loop {
let event = tokio::select! {
result = running.wait(), if !completed => {
completed = true;
result.map(|_| None)
},
line = lines.next() => Ok(Some(line)),
}?;
let Some(line) = event else { continue };
match line {
Some(Ok(line)) => yield line,
Some(Err(error)) => {
running.cancel().await;
Err(tail.error(error))?;
},
None => break,
}
}
}
if !completed {
running.wait().await?;
}
}))
}
}
impl Pipeline<Writer> {
pub async fn execute(self) -> Result<Cursor<Vec<u8>>, PipelineError> {
let mut running = start(self.stages, false).await?;
Ok(Cursor::new(running.wait().await?))
}
}
impl Pipeline<Indexer> {
pub async fn execute(self) -> Result<(), PipelineError> {
start(self.stages, false).await?.wait().await?;
Ok(())
}
}
#[doc(hidden)]
#[derive(Debug)]
pub struct PipelineStage {
program: OsString,
kind: StageKind,
error: Option<ExecutorError>,
external_input: bool,
external_writer: bool,
batching: BatchOptions,
}
#[derive(Debug)]
enum StageKind {
Native {
executor: Executor,
input: Input,
output: Output,
},
LimitedLister(Box<Lister>),
}
impl PipelineStage {
pub(crate) fn native(
mut executor: Executor,
input: Input,
output: Output,
error: Option<ExecutorError>,
) -> Self {
Self {
program: executor.command().as_std().get_program().to_os_string(),
batching: executor.batch_options(),
external_input: !matches!(input, Input::Ignored),
external_writer: matches!(output, Output::AsyncWrite(_)),
kind: StageKind::Native {
executor,
input,
output,
},
error,
}
}
pub(crate) fn limited_lister(
lister: Lister,
program: OsString,
external_writer: bool,
error: Option<ExecutorError>,
batching: BatchOptions,
) -> Self {
Self {
program,
kind: StageKind::LimitedLister(Box::new(lister)),
error,
external_input: false,
external_writer,
batching,
}
}
}
pub(crate) fn graph_formats(input: Option<&str>, output: Option<&str>) -> Option<ExecutorError> {
for (option, format) in [("input", input), ("output", output)] {
if let Some(format) = format {
if format != "jsonl" {
return Some(invalid(alloc::format!(
"pipeline graph {option} format must be jsonl, got {format}"
)));
}
}
}
None
}
fn invalid(message: impl Into<alloc::string::String>) -> ExecutorError {
io::Error::new(io::ErrorKind::InvalidInput, message.into()).into()
}
macro_rules! stage {
($program:ty, $value:ident, $input:expr, $output:expr, $input_format:expr, $output_format:expr) => {
impl From<$program> for crate::pipeline::PipelineStage {
fn from($value: $program) -> Self {
let error = crate::pipeline::graph_formats($input_format, $output_format);
Self::native($value.executor, $input, $output, error)
}
}
};
}
pub(crate) use stage;
#[derive(Clone)]
struct StageInfo {
index: usize,
program: OsString,
}
impl StageInfo {
fn error(&self, error: impl Into<ExecutorError>) -> PipelineError {
PipelineError {
stage: self.index,
program: self.program.clone(),
error: error.into(),
}
}
}
type Job = Pin<Box<dyn Future<Output = Result<Vec<u8>, PipelineError>> + Send>>;
struct Running {
jobs: Vec<Option<Job>>,
stop: watch::Sender<bool>,
lines: Option<FrameStream>,
tail: StageInfo,
failure: Option<PipelineError>,
output: Vec<u8>,
}
impl Running {
async fn next(&mut self) -> Option<(usize, Result<Vec<u8>, PipelineError>)> {
poll_fn(|cx| {
let mut pending = false;
for (index, job) in self.jobs.iter_mut().enumerate().rev() {
if let Some(future) = job {
match future.as_mut().poll(cx) {
Poll::Ready(result) => {
*job = None;
return Poll::Ready(Some((index, result)));
},
Poll::Pending => pending = true,
}
}
}
if pending {
Poll::Pending
} else {
Poll::Ready(None)
}
})
.await
}
async fn wait(&mut self) -> Result<Vec<u8>, PipelineError> {
while let Some((index, result)) = self.next().await {
match result {
Ok(bytes) if index + 1 == self.jobs.len() => self.output = bytes,
Ok(_) => {},
Err(error) if self.failure.is_none() => {
self.failure = Some(error);
self.stop.send_replace(true);
},
Err(_) => {},
}
}
match self.failure.take() {
Some(error) => Err(error),
None => Ok(core::mem::take(&mut self.output)),
}
}
async fn cancel(&mut self) {
self.stop.send_replace(true);
while self.next().await.is_some() {}
}
}
struct Spawned {
child: Child,
input: Input,
output: Output,
info: StageInfo,
}
async fn run_stage(
mut stage: Spawned,
mut stop: watch::Receiver<bool>,
source: Option<StageInfo>,
) -> Result<Vec<u8>, PipelineError> {
let cancelled = *stop.borrow();
let completion = if cancelled {
None
} else {
tokio::select! {
biased;
result = crate::executor::communicate_child(&mut stage.child, &mut stage.input, &mut stage.output) => Some(result),
_ = stop.changed() => None,
}
};
match completion {
Some(result) => {
let completion = result.map_err(|error| stage.info.error(error))?;
let info = if matches!(completion.input, InputCompletion::SourceFailed(_)) {
source.as_ref().unwrap_or(&stage.info)
} else {
&stage.info
};
completion
.into_result()
.map(Cursor::into_inner)
.map_err(|error| info.error(error))
},
None => {
let _ = stage.child.start_kill();
let _ = stage.child.wait().await;
Ok(Vec::new())
},
}
}
async fn start(
mut stages: Vec<PipelineStage>,
capture_graph: bool,
) -> Result<Running, PipelineError> {
let count = stages.len();
for (index, stage) in stages.iter_mut().enumerate() {
let info = StageInfo {
index,
program: stage.program.clone(),
};
if let Some(error) = stage.error.take() {
return Err(info.error(error));
}
if index != 0 && stage.external_input {
return Err(info.error(invalid(
"piped stages must use Input::Ignored; their input comes from the preceding stage",
)));
}
if index + 1 != count && stage.external_writer {
return Err(info.error(invalid("an intermediate pipeline stage cannot also forward stdout to an AsyncWrite destination")));
}
}
let tail = StageInfo {
index: count - 1,
program: stages.last().unwrap().program.clone(),
};
let (stop, receiver) = watch::channel(false);
let mut running = Running {
jobs: Vec::new(),
stop,
lines: None,
tail,
failure: None,
output: Vec::new(),
};
let mut source_info = None;
let mut limited_source = None;
if matches!(stages[0].kind, StageKind::LimitedLister(_)) {
let stage = stages.remove(0);
let info = StageInfo {
index: 0,
program: stage.program,
};
let StageKind::LimitedLister(mut lister) = stage.kind else {
unreachable!()
};
if stages.is_empty() {
running.lines = Some(
lister
.execute_frames()
.await
.map_err(|error| info.error(error))?,
);
return Ok(running);
}
limited_source = Some(lister);
source_info = Some(info);
}
let base = usize::from(source_info.is_some());
for index in 0..stages.len().saturating_sub(1) {
let (reader, writer) = io::pipe().map_err(|error| {
StageInfo {
index: index + base,
program: stages[index].program.clone(),
}
.error(error)
})?;
let StageKind::Native { executor, .. } = &mut stages[index].kind else {
unreachable!()
};
executor.command().stdout(Stdio::from(writer));
let StageKind::Native { executor, .. } = &mut stages[index + 1].kind else {
unreachable!()
};
executor.command().stdin(Stdio::from(reader));
}
if let Some(lister) = limited_source {
let source = lister
.into_pipeline_source()
.await
.map_err(|error| source_info.as_ref().unwrap().error(error))?;
let StageKind::Native {
executor, input, ..
} = &mut stages[0].kind
else {
unreachable!()
};
*input = Input::Jsonl(source);
executor.command().stdin(input.as_stdio());
}
let mut spawned: Vec<Spawned> = Vec::new();
let mut plans = stages.into_iter().enumerate();
while let Some((index, stage)) = plans.next() {
let info = StageInfo {
index: index + base,
program: stage.program,
};
let StageKind::Native {
mut executor,
input,
mut output,
} = stage.kind
else {
unreachable!()
};
if info.index + 1 != count {
output = Output::Ignored;
}
let result = executor.spawn().await;
drop(executor);
let mut child = match result {
Ok(child) => child,
Err(error) => {
drop(plans);
drop(input);
for stage in &mut spawned {
stage.input = Input::Ignored;
let _ = stage.child.start_kill();
}
for stage in &mut spawned {
let _ = stage.child.wait().await;
}
return Err(info.error(error));
},
};
if info.index + 1 == count && capture_graph && matches!(output, Output::Captured) {
running.lines = child.stdout.take().map(crate::jsonl::jsonl_frames);
}
spawned.push(Spawned {
child,
input,
output,
info,
});
}
for (index, stage) in spawned.into_iter().enumerate() {
running.jobs.push(Some(Box::pin(run_stage(
stage,
receiver.clone(),
if index == 0 {
source_info.clone()
} else {
None
},
))));
}
Ok(running)
}