prealloc writing done: --save-raw

fuck_this_bullshit
Avril 4 years ago
parent 2583ae674c
commit 00798609b2
Signed by: flanchan
GPG Key ID: 284488987C31F630

@ -33,6 +33,7 @@ const OPTIONS_NORMAL: &'static [&'static str] = &[
"-v Verbose mode. Output extra information.", "-v Verbose mode. Output extra information.",
#[cfg(feature="inspect")] "--save <file> Dump the collected data to this file for further inspection.", #[cfg(feature="inspect")] "--save <file> Dump the collected data to this file for further inspection.",
#[cfg(feature="inspect")] "-D Dump the collected data to `stdout` (see `--save`.)", #[cfg(feature="inspect")] "-D Dump the collected data to `stdout` (see `--save`.)",
#[cfg(feature="prealloc")] "--save-raw <file> Dump the collected data to this file uncompressed.",
"- Stop parsing arguments, treat all the rest as paths.", "- Stop parsing arguments, treat all the rest as paths.",
]; ];
@ -140,12 +141,17 @@ fn parse<I: IntoIterator<Item=String>>(args: I) -> eyre::Result<Mode>
cfg.output_level = config::OutputLevel::Verbose; cfg.output_level = config::OutputLevel::Verbose;
}, },
#[cfg(feature="inspect")] "-D" => { #[cfg(feature="inspect")] "-D" => {
cfg.serialise_output = Some(None); cfg.serialise_output = Some(config::OutputSerialisationMode::Stdout);
}, },
#[cfg(feature="inspect")] "--save" => { #[cfg(feature="inspect")] "--save" => {
let file = args.next().ok_or(eyre!("`--save` expects a filename parameter")) let file = args.next().ok_or(eyre!("`--save` expects a filename parameter"))
.with_suggestion(suggestion_intended_arg.clone())?; .with_suggestion(suggestion_intended_arg.clone())?;
cfg.serialise_output = Some(Some(file.into())); cfg.serialise_output = Some(config::OutputSerialisationMode::File(file.into()));
},
#[cfg(feature="prealloc")] "--save-raw" => {
let file = args.next().ok_or(eyre!("`--save-raw` expects a filename parameter"))
.with_suggestion(suggestion_intended_arg.clone())?;
cfg.serialise_output = Some(config::OutputSerialisationMode::PreallocFile(file.into()));
}, },
"--recursive" => { "--recursive" => {
let max = args.next().ok_or(eyre!("`--recursive` expects a parameter")) let max = args.next().ok_or(eyre!("`--recursive` expects a parameter"))

@ -96,6 +96,14 @@ impl From<usize> for Recursion
} }
} }
#[cfg(feature="inspect")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OutputSerialisationMode
{
Stdout,
File(PathBuf),
#[cfg(feature="prealloc")] PreallocFile(PathBuf),
}
/// Configuration for this run /// Configuration for this run
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@ -108,7 +116,7 @@ pub struct Config
pub output_level: OutputLevel, pub output_level: OutputLevel,
#[cfg(feature="inspect")] #[cfg(feature="inspect")]
pub serialise_output: Option<Option<PathBuf>>, // Some(None) means dump to `stdout` pub serialise_output: Option<OutputSerialisationMode>,
} }
impl Config impl Config
@ -131,7 +139,7 @@ impl Config
#[inline] pub fn is_using_stdout(&self) -> bool #[inline] pub fn is_using_stdout(&self) -> bool
{ {
#[cfg(feature="inspect")] { #[cfg(feature="inspect")] {
return if let Some(None) = self.serialise_output { return if let Some(OutputSerialisationMode::Stdout) = self.serialise_output {
true true
} else { } else {
false false

@ -72,10 +72,11 @@ async fn normal(cfg: config::Config) -> eyre::Result<()>
match cfg.serialise_output.as_ref().map(|ser_out| { match cfg.serialise_output.as_ref().map(|ser_out| {
type BoxedWrite = Box<dyn tokio::io::AsyncWrite + Unpin>; type BoxedWrite = Box<dyn tokio::io::AsyncWrite + Unpin>;
use futures::FutureExt; use futures::FutureExt;
use config::OutputSerialisationMode;
match ser_out { match ser_out {
Some(output_file) => { OutputSerialisationMode::File(output_file) => {
use tokio::fs::OpenOptions; use tokio::fs::OpenOptions;
async move { (async move {
let stream = OpenOptions::new() let stream = OpenOptions::new()
.write(true) .write(true)
.truncate(true) .truncate(true)
@ -85,16 +86,42 @@ async fn normal(cfg: config::Config) -> eyre::Result<()>
.with_section(|| format!("{:?}", output_file).header("File was"))?; .with_section(|| format!("{:?}", output_file).header("File was"))?;
Ok::<BoxedWrite, eyre::Report>(Box::new(stream)) Ok::<BoxedWrite, eyre::Report>(Box::new(stream))
}.boxed() }.boxed(), None)
},
OutputSerialisationMode::Stdout => (async move { Ok::<BoxedWrite, _>(Box::new(tokio::io::stdout())) }.boxed(), None),
#[cfg(feature="prealloc")] OutputSerialisationMode::PreallocFile(output_file) => {
(async move {
Ok::<BoxedWrite, _>(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))
}, },
None => async move { Ok::<BoxedWrite, _>(Box::new(tokio::io::stdout())) }.boxed(),
} }
}) { }) {
Some(stream_fut) => { // We use tuple item `.1` here to indicate if we're in normal write mode.
// None -> normal
// Some(path) -> prealloc
Some((stream_fut, None)) => {
let stream = stream_fut.await?; let stream = stream_fut.await?;
serial::write_async(stream, &graph).await serial::write_async(stream, &graph).await
.wrap_err(eyre!("Failed to serialise graph to stream"))?; .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;
tokio::task::spawn_blocking(move || {
serial::write_sync_map(&mut file, &graph)
}).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"))?;
},
_ => (), _ => (),
} }

Loading…
Cancel
Save