#![feature(try_blocks)] #![allow(dead_code)] #[macro_use] extern crate pin_project; #[macro_use] extern crate lazy_static; #[macro_use] extern crate cfg_if; #[macro_use] extern crate ad_hoc_iter; #[cfg(feature="inspect")] use serde::{Serialize, Deserialize}; #[cfg(feature="jemalloc")] use jemallocator::Jemalloc; #[cfg(feature="jemalloc")] #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; use std::convert::{TryFrom, TryInto}; use color_eyre::{ eyre::{ self, eyre, WrapErr as _, }, Help as _, SectionExt as _, }; #[macro_use] mod ext; pub use ext::prelude::*; use std::sync::Arc; mod bytes; mod data; mod config; mod state; mod arg; mod work; mod info; #[cfg(feature="inspect")] mod serial; #[cfg(feature="defer-drop")] mod defer_drop; #[cfg(feature="inspect")] async fn write_graph(graph: Arc) -> eyre::Result<()> { let cfg = config::get_global(); match cfg.serialise_output.as_ref().map(|ser_out| { cfg_eprintln!(Verbose; cfg, "Writing graph to stream..."); type BoxedWrite = Box; use futures::FutureExt; use config::OutputSerialisationMode; let should_comp = ser_out.should_compress(); match ser_out { OutputSerialisationMode::File(output_file) | OutputSerialisationMode::RawFile(output_file) => { use tokio::fs::OpenOptions; (async move { let stream = OpenOptions::new() .write(true) .truncate(true) .create(true) .open(output_file).await .wrap_err(eyre!("Failed to open file for writing")) .with_section(|| format!("{:?}", output_file).header("File was"))?; Ok::(Box::new(stream)) }.boxed(), None, should_comp) }, OutputSerialisationMode::Stdout | OutputSerialisationMode::RawStdout => (async move { Ok::(Box::new(tokio::io::stdout())) }.boxed(), Option::<&std::path::PathBuf>::None, should_comp), #[cfg(feature="prealloc")] OutputSerialisationMode::PreallocFile(output_file) => { (async move { Ok::(Box::new(tokio::io::sink())) // we use a sink as a shim stream since it will never be used when tuple item `.1` is Some() }.boxed(), Some(output_file), false) }, } }) { // We use tuple item `.1` here to indicate if we're in normal write mode. // None -> normal // Some(path) -> prealloc // `.2` indicates if we should compress while in normal write mode. Some((stream_fut, None, compress)) => { let stream = stream_fut.await?; serial::write_async(stream, graph.as_ref(), compress).await .wrap_err(eyre!("Failed to serialise graph to stream"))?; }, #[cfg(feature="prealloc")] Some((_task_fut, Some(output_file), _)) => { use tokio::fs::OpenOptions; let file = OpenOptions::new() .write(true) .read(true) //needed for map I think? .truncate(true) .create(true) .open(&output_file).await .wrap_err(eyre!("Failed to open file for mapping")) .with_section(|| format!("{:?}", output_file).header("File was"))?; let mut file = file.into_std().await; cfg_eprintln!(Verbose; cfg, "Opened file for prealloc-write"); tokio::task::spawn_blocking(move || { serial::write_sync_map(&mut file, graph.as_ref()) }).await.wrap_err(eyre!("Prealloc panicked while dumping")) .with_section(|| format!("{:?}", output_file).header("File was"))? .wrap_err(eyre!("Prealloc failed to dump graph to file")) .with_section(|| format!("{:?}", output_file).header("File was"))?; }, _ => (), } Ok(()) } async fn normal(cfg: config::Config) -> eyre::Result<()> { let state = state::State::new(cfg .validate() .wrap_err(eyre!("Invalid config")) .with_suggestion(|| "Try running `--help`")?); let (graph, cfg) = tokio::select!{ x = work::work_on_all(state) => {x} _ = tokio::signal::ctrl_c() => { return Err(eyre!("Interrupt signalled, exiting")); } }; let cfg = cfg.make_global(); let (treemap, graph) = tokio::task::spawn_blocking(move || { cfg_println!(Verbose; cfg, "Computing hierarchy..."); // Create treemap #[cfg(feature="treemap")] let treemap = { //Check config for if the treemap flag is present. If it is, get the width and height from there too. if let Some(&(x,y)) = cfg.inspection.treemap.as_ref() { Some(info::treemap(cfg, &graph, ( x as f64, y as f64, ))) } else { None } }; #[cfg(not(feature="treemap"))] let treemap = Option::::None; // Create recursive graph let mut graph = graph.into_hierarchical(); cfg_println!(Verbose; cfg, "Computing sizes..."); graph.compute_recursive_sizes(); (treemap, graph) }).await.expect("Failed to compute hierarchy from graph"); //#[cfg(debug_assertions)] cfg_eprintln!(Verbose; cfg, "{:?}", graph); #[allow(unused_imports)] use futures::future::OptionFuture; let (graph, writer) = { cfg_if!{ if #[cfg(feature="inspect")] { let graph = Arc::new(graph); let (g, w) = if cfg.serialise_output.is_some() { (Arc::clone(&graph), Some(tokio::spawn(write_graph(graph)))) } else { (graph, None) }; (g, OptionFuture::from(w)) } else { (graph, ())//OptionFuture::from(Option::>::None)) } } }; #[cfg(feature="treemap")] if let Some(treemap) = treemap { //TODO: Print treemap todo!("{:?}",treemap); } info::print_basic_max_info(&cfg, &graph); #[cfg(feature="inspect")] match writer.await { Some(Ok(Ok(_))) if cfg.serialise_output.is_some() => cfg_eprintln!(Verbose; cfg, "Written successfully"), Some(Ok(error @ Err(_))) => return error.wrap_err(eyre!("Failed to write graph to output stream")), Some(Err(_)) => cfg_eprintln!(Silent; cfg, "Panic while writing graph to stream"), _ => (), } #[cfg(not(feature="inspect"))] drop(writer); Ok(()) } async fn parse_mode() -> eyre::Result<()> { match arg::parse_args() .wrap_err(eyre!("Failed to parse args")) .with_suggestion(|| "Try running `--help`")? { arg::Mode::Normal(cfg) => { #[cfg(debug_assertions)] eprintln!("cfg: {:#?}", cfg); normal(cfg).await }, arg::Mode::Help => arg::help(), } } #[tokio::main] async fn main() -> eyre::Result<()> { color_eyre::install()?; parse_mode().await .wrap_err(eyre!("Fatal error"))?; /* let max_size = graph.directories().map(|x| x.size()).max(); println!("Max size: {:?}", max_size);*/ //println!("{:?}", graph); Ok(()) }