enumerate-ordered/src/main.rs

#[macro_use] extern crate log;

use color_eyre::{
    eyre::{
	self,
	eyre,
	WrapErr as _,
    },
    SectionExt as _,
    Help as _,
};
use tokio::{
    sync::{
	mpsc,
    },
};

use jemallocator::Jemalloc;

#[global_allocator]
static GLOBAL: Jemalloc = Jemalloc;

mod args;
mod order;
mod work;
mod walk;

fn init_logging() -> eyre::Result<()>
{
    color_eyre::install()?;
    pretty_env_logger::init();
    Ok(())
}

async fn work_on(cfg: work::Config, mut into: mpsc::Receiver<work::FileInfo>) -> eyre::Result<work::FSTimeMap>
{
    use work::*;
    let mut map = FSTimeMap::new(cfg);
    while let Some(file) = into.recv().await {
	if cfg!(debug_assertions) {
	    trace!("insert +{}", map.len());
	}
	map.insert(file);
    }
    Ok(map)
}

async fn walk_paths<I, P>(paths: I, cfg: walk::Config, worker_cfg: &std::sync::Arc<work::Config>, to: mpsc::Sender<work::FileInfo>) -> eyre::Result<usize>
where I: futures::stream::Stream<Item = P>,
      P: AsRef<std::path::Path>
{
    use futures::prelude::*;
    let children: Vec<usize> = 
	paths.map(|path| futures::stream::once(walk::start_walk(cfg.clone(), std::sync::Arc::clone(worker_cfg), path, to.clone())).boxed_local()).flatten_unordered(None).try_collect().await?;
    
    Ok(children.into_iter().sum())
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
    init_logging().wrap_err("Failed to set logging handlers")?;

    let (tx, rx) = mpsc::channel(4096);

    //TODO: Read main config from args

    let args = args::parse_args()
	.wrap_err("Failed to parse command line args")
	.with_suggestion(|| "Try `--help`")?;
    
    let worker_cfg = {
	//TODO: Read worker config from main config
	std::sync::Arc::new(work::Config::default())
    };
    let walker_cfg = {
	//TODO: Read walker config from main config
	walk::Config::default()
    };
    
    // Spin up ordering task.
    let ordering = {
	let cfg = (*worker_cfg).clone();
	tokio::spawn(async move {
	    trace!("spun up ordering backing thread");
	    work_on(cfg, rx).await //TODO: Parse config from args
	})
    };

    trace!("Starting recursive walk of input locations");
    //TODO: Trace directory trees from paths in `args` and/or `stdin` and pass results to `tx`
    let walk = walk_paths(args.paths(), walker_cfg, &worker_cfg, tx);
    tokio::pin!(walk);
    
    let set = async move {
	ordering.await.wrap_err("Ordering task panic")?
            .wrap_err(eyre!("Failed to collect ordered files"))
    };
    tokio::pin!(set);

    /*let set = tokio::select! {
    n = walk => {
    let n =n.wrap_err("Walker failed")?;
    info!("Walked {} files", n);
},
    res = set => {
    let set = res.wrap_err("Ordering task failed before walk completed")?;
    return Err(eyre!("Ordering task exited before walker task"));
}
};*/
    let (walk, set) = {
	let (w, s) = tokio::join!(walk, set);
	(w?, s?)
    };
    info!("Walked {} files", walk);

    // Write the output in a blocking task - There's not much benefit from async here.
    tokio::task::spawn_blocking(move || -> eyre::Result<()> {
	use std::io::Write;
	use std::os::unix::prelude::*;
	let mut stdout = {
	    let lock = std::io::stdout().lock();
	    std::io::BufWriter::new(lock)
	};
	trace!("Writing ordered results to stdout... (buffered, sync)");
	for info in set.into_iter()
	{
	    stdout.write_all(info.path().as_os_str().as_bytes())
		.and_then(|_| stdout.write_all(&[b'\n']))
		.wrap_err("Failed to write raw pathname for entry to stdout")
		.with_context(|| format!("{:?}", info.path()).header("Pathname was"))?;
	}
	
	stdout.flush().wrap_err("Failed to flush buffered output to stdout")?;
	Ok(())
    }).await.wrap_err("Writer (blocking) task panic")?
	.wrap_err("Failed to write results to stdout")?;

    trace!("Finished output task");
    
    Ok(())
}