master
Avril 4 years ago
parent 5521799550
commit 4459d58df9
Signed by: flanchan
GPG Key ID: 284488987C31F630

@ -8,10 +8,13 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features] [features]
default = ["parallel", "limit-concurrency", "splash"] default = ["parallel", "limit-concurrency", "splash", "limit-recursion"]
# Limit the concurrent operations of parallel mode to 4096 # Limit the concurrent operations of parallel mode to 4096
limit-concurrency = ["parallel"] limit-concurrency = ["parallel"]
# Handle directories recursively
recursive = []
limit-recursion = ["recursive"]
splash = [] splash = []
parallel = ["tokio", "futures"] parallel = ["tokio", "futures"]
threads = ["parallel", "tokio/rt-threaded"] threads = ["parallel", "tokio/rt-threaded"]

@ -28,6 +28,8 @@ use map::*;
mod temp; mod temp;
mod error; mod error;
mod arg; mod arg;
#[cfg(feature="recursive")]
mod recurse;
cfg_if!{ cfg_if!{
if #[cfg(feature="splash")] { if #[cfg(feature="splash")] {
@ -57,11 +59,20 @@ fn args_or_out<T: ExactSizeIterator>(i: T, low: usize) -> T
#[cfg(feature="parallel")] #[cfg(feature="parallel")]
#[cfg_attr(feature="parallel", tokio::main)] #[cfg_attr(feature="parallel", tokio::main)]
async fn main() -> eyre::Result<()> { async fn main() -> eyre::Result<()> {
reyre!(init(), "Failed to initialise")?; use futures::{
stream,
prelude::*,
};
reyre!(parallel::main(futures::stream::iter(args_or_out(std::env::args(), 2) reyre!(init(), "Failed to initialise")?;
.skip(1) reyre!(parallel::main(stream::iter(args_or_out(std::env::args(), 2)
.dedup())).await, .skip(1)
.dedup())
.filter_map(|file| {
async move {
Some(parallel::expand_dir(file).await)
}
}).flatten()).await,
"Jobs failed") "Jobs failed")
} }

@ -12,9 +12,11 @@ use futures::{
Future, Future,
OptionFuture, OptionFuture,
FutureExt, FutureExt,
BoxFuture,
join_all, join_all,
}, },
stream::{ stream::{
self,
Stream, Stream,
StreamExt, StreamExt,
}, },
@ -22,6 +24,7 @@ use futures::{
use tokio::{ use tokio::{
sync::{ sync::{
Semaphore, Semaphore,
mpsc,
}, },
fs::{ fs::{
OpenOptions, OpenOptions,
@ -39,6 +42,8 @@ cfg_if!{
} }
} }
#[cfg(feature="recursive")] use recurse::{MAX_DEPTH, Recursion};
fn gensem() -> Option<Arc<Semaphore>> fn gensem() -> Option<Arc<Semaphore>>
{ {
trace!("Limiting concurrency to {:?}", MAX_WORKERS); trace!("Limiting concurrency to {:?}", MAX_WORKERS);
@ -90,7 +95,7 @@ async fn work<P: AsRef<Path>>(apath: P, sem: Option<Arc<Semaphore>>) -> Result<(
} }
async fn join_stream<I: Stream>(stream: I) -> impl Iterator<Item=<I::Item as Future>::Output> + ExactSizeIterator async fn join_stream<I: Stream>(stream: I) -> impl Iterator<Item=<I::Item as Future>::Output> + ExactSizeIterator
where I::Item: Future where I::Item: Future
{ {
//gotta be a better way than heap allocating here, right? //gotta be a better way than heap allocating here, right?
stream.then(|x| async move { x.await }).collect::<Vec<_>>().await.into_iter() stream.then(|x| async move { x.await }).collect::<Vec<_>>().await.into_iter()
@ -108,7 +113,6 @@ pub async fn main<I: Stream<Item=String>>(list: I) -> eyre::Result<()>
for (i, res) in (0usize..).zip(join_stream(list.map(|file| tokio::spawn(work(file, sem.clone())))) for (i, res) in (0usize..).zip(join_stream(list.map(|file| tokio::spawn(work(file, sem.clone()))))
.map(|x| {trace!("--- {} Finished ---", x.len()); x}).await) .map(|x| {trace!("--- {} Finished ---", x.len()); x}).await)
{ {
//trace!("Done on {:?}", res);
match res { match res {
Ok(Ok((path, true))) => info!("<{:?}> OK (processed)", path), Ok(Ok((path, true))) => info!("<{:?}> OK (processed)", path),
Ok(Ok((path, false))) => info!("<{:?}> OK (skipped)", path), Ok(Ok((path, false))) => info!("<{:?}> OK (skipped)", path),
@ -156,3 +160,73 @@ pub async fn main<I: Stream<Item=String>>(list: I) -> eyre::Result<()>
Ok(()) Ok(())
} }
#[cfg(feature="recursive")]
fn push_dir<'a>(path: &'a Path, depth: usize, to: mpsc::Sender<String>) -> BoxFuture<'a, tokio::io::Result<()>>
{
async move {
let mut dir = fs::read_dir(path).await?;
let mut workers = match dir.size_hint() {
(0, Some(0)) | (0, None) => Vec::new(),
(x, None) | (_, Some(x)) => Vec::with_capacity(x),
};
let can_recurse = match MAX_DEPTH {
Recursion::All => true,
Recursion::N(n) if depth < usize::from(n) => true,
_ => false,
};
while let Some(item) = dir.next_entry().await? {
let mut to = to.clone();
workers.push(async move {
match path.join(item.file_name()).into_os_string().into_string() {
Ok(name) => {
if item.file_type().await?.is_dir() {
if can_recurse {
if let Err(e) = push_dir(name.as_ref(), depth+1, to).await {
error!("Walking dir {:?} failed: {}", item.file_name(), e);
}
}
} else {
to.send(name).await.unwrap();
}
},
Err(err) => {
error!("Couldn't process file {:?} because it contains invalid UTF-8", err);
},
}
Ok::<_, std::io::Error>(())
});
}
join_all(workers).await;
Ok(())
}.boxed()
}
pub async fn expand_dir(p: String) -> impl Stream<Item=String>
{
cfg_if!{
if #[cfg(feature="recursive")] {
let (mut tx, rx) = mpsc::channel(16);
tokio::spawn(async move {
let path = Path::new(&p);
if path.is_dir() {
if let Err(err) = push_dir(path, 0, tx).await {
error!("Walking dir {:?} failed: {}", path, err);
}
} else {
tx.send(p).await.unwrap();
}
});
rx
} else {
stream::iter(iter::once(p).filter_map(|p| {
if Path::new(&p).is_dir() {
warn!("{:?} is a directory, skipping", p);
None
} else {
Some(p)
}
}))
}
}
}

@ -0,0 +1,33 @@
//! Recursion stuffs
use super::*;
use std::{
num::NonZeroUsize,
fmt,
};
#[derive(Debug)]
pub enum Recursion
{
All,
N(NonZeroUsize),
}
cfg_if!{
if #[cfg(feature="limit-recursion")] {
pub const MAX_DEPTH: Recursion = Recursion::N(unsafe{NonZeroUsize::new_unchecked(256)});
} else {
pub const MAX_DEPTH: Recursion = Recursion::All;
}
}
impl fmt::Display for Recursion
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
{
match self {
Self::N(n) => write!(f, "{}", n),
Self::All => write!(f, "unlimited"),
_ => write!(f, "no"),
}
}
}

@ -39,9 +39,12 @@ For verbose output, set `RUST_LOG` env var to one of the following:
Made by {} with <3 (Licensed GPL 3.0 or later)"#, arg::program_name(), env!("CARGO_PKG_VERSION"), env!("CARGO_PKG_AUTHORS")); Made by {} with <3 (Licensed GPL 3.0 or later)"#, arg::program_name(), env!("CARGO_PKG_VERSION"), env!("CARGO_PKG_AUTHORS"));
println!("\nEnabled extensions: "); println!("\nEnabled extensions: ");
feature!(in nightly, "\tCompiled with Rust nightly extensions"); feature!(in nightly, "\tCompiled with Rust nightly extensions");
println!(); println!("Features:");
feature!("parallel", "\tWill run up to {} operations in parallel", parallel::MAX_WORKERS.map(|x| Cow::Owned(x.to_string())).unwrap_or(Cow::Borrowed("unlimited"))); feature!("parallel", "\tWill run up to {} operations in parallel", parallel::MAX_WORKERS.map(|x| Cow::Owned(x.to_string())).unwrap_or(Cow::Borrowed("unlimited")));
feature!("limit-concurrency", "Concurrency is capped"); feature!("limit-concurrency", "Concurrency is capped");
feature!("threads", "\tUsing thread-pool"); feature!("threads", "\tUsing thread-pool");
feature!("recursive", "\tRecursivly process files up to {} directories deep", recurse::MAX_DEPTH);
feature!("limit-recursion", "Concurrency is capped");
std::process::exit(1) std::process::exit(1)
} }

Loading…
Cancel
Save