You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
99 lines
2.5 KiB
99 lines
2.5 KiB
use super::*;
|
|
use std::{
|
|
num::NonZeroUsize,
|
|
sync::Arc,
|
|
marker::{
|
|
Send,
|
|
Sync,
|
|
},
|
|
path::{Path,PathBuf,},
|
|
ffi::OsStr,
|
|
};
|
|
use tokio::{
|
|
prelude::*,
|
|
sync::{
|
|
Semaphore,
|
|
mpsc,
|
|
},
|
|
io,
|
|
};
|
|
use futures::future::{
|
|
join_all,
|
|
Future,
|
|
};
|
|
|
|
async fn maybe_await<T>(from: Option<T>) -> Option<<T as Future>::Output>
|
|
where T: Future
|
|
{
|
|
if let Some(v) = from {
|
|
Some(v.await)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
async fn do_work(process: impl AsRef<Path>, file: impl AsRef<OsStr>) -> Result<Vec<String>, process::Error>
|
|
{
|
|
let process = process.as_ref();
|
|
let file = file.as_ref();
|
|
|
|
let (tx, mut rx) = mpsc::channel(16);
|
|
println!("[p]: Processing {:?}", file);
|
|
let collector = tokio::spawn(async move {
|
|
let mut stack = fixed_stack::FixedStack::new(100);
|
|
while let Some(value) = rx.recv().await {
|
|
stack.push(value);
|
|
}
|
|
stack
|
|
});
|
|
tokio::task::yield_now().await;
|
|
match process::contained_spawn(process, std::iter::once(file), tx).await {
|
|
Ok(_) => Ok(collector.await.expect("Child panic").into_iter().collect()),
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
pub async fn work<I,T,U>(process: U, files: I, children: Option<NonZeroUsize>) -> Result<(), Box<dyn std::error::Error>>
|
|
where I: IntoIterator<Item=T>,
|
|
T: AsRef<OsStr> + Send + Sync + 'static,
|
|
U: Into<PathBuf>
|
|
{
|
|
let (tx,mut rx) = mpsc::channel::<(T, Vec<String>)>(children.as_ref().map(|&x| usize::from(x)).unwrap_or(16));
|
|
let semaphore = children.map(|children| Arc::new(Semaphore::new(children.into())));
|
|
let process = Arc::new(process.into());
|
|
|
|
let display = tokio::spawn(async move {
|
|
let mut stdout = io::stdout();
|
|
while let Some((file, values)) = rx.recv().await {
|
|
for line in values.into_iter()
|
|
{
|
|
let line = format!("[i] {:?}: {}\n", file.as_ref(), line);
|
|
let _ = stdout.write_all(line.as_bytes()).await;
|
|
|
|
}
|
|
}
|
|
});
|
|
|
|
for failed in join_all(files.into_iter()
|
|
.map(|filename| {
|
|
let semaphore = semaphore.clone();
|
|
let process = Arc::clone(&process);
|
|
let mut tx = tx.clone();
|
|
tokio::spawn(async move {
|
|
let _lock = maybe_await(semaphore.map(|x| x.acquire_owned())).await;
|
|
match do_work(process.as_ref(), &filename).await {
|
|
Ok(strings) => tx.send((filename, strings)).await.map_err(|_| "Child panic").unwrap(),
|
|
Err(error) => eprintln!("[!] {:?}: {}", filename.as_ref(), error),
|
|
}
|
|
})
|
|
})).await
|
|
.into_iter()
|
|
.filter_map(|x| x.err())
|
|
{
|
|
eprintln!("Child panic: {}", failed);
|
|
}
|
|
drop(tx);
|
|
display.await?;
|
|
Ok(())
|
|
}
|