use crate::args;
use crate::artifacts::{Artifact, HashedArtifact};
use crate::errors::*;
use crate::plot::{Ctx, Plot, PlotExtras};
use std::collections::BTreeMap;
use std::fs;
use std::fs::File;
use std::io::{self, Read, Write};
use zstd::stream::write::Encoder;
pub const ZSTD_COMPRESSION_LEVEL: i32 = 3;
pub struct ArchiveBuilder<W: Write> {
builder: tar::Builder<W>,
}
impl<W: Write> ArchiveBuilder<W> {
pub fn new(writer: W) -> Self {
Self {
builder: tar::Builder::new(writer),
}
}
pub fn append(&mut self, path: &str, data: &[u8]) -> Result<()> {
let mut header = tar::Header::new_gnu();
header.set_size(data.len() as u64);
header.set_mode(0o640);
debug!("Adding to archive: {:?} ({} bytes)", path, data.len());
self.builder.append_data(&mut header, path, data)?;
Ok(())
}
pub fn finish(self) -> Result<()> {
debug!("Finishing archive...");
self.builder.into_inner()?;
Ok(())
}
}
pub struct CompiledArchive {
pub plot: String,
pub artifacts: BTreeMap<String, HashedArtifact>,
}
impl CompiledArchive {
pub async fn build<R: Read>(plot: &mut Plot, cache_from: Option<R>) -> Result<CompiledArchive> {
let mut artifacts = BTreeMap::new();
if let Some(f) = cache_from {
debug!("Loading existing plot as cache...");
Ctx::load_as_download_cache(f, plot, &mut artifacts)
.await
.context("Failed to load existing plot as cache")?;
debug!("Finished loading existing plot");
}
debug!("Setting up compressed writer...");
info!("Resolving plot...");
let PlotExtras {
mut artifacts,
signing_keys,
sessions: _,
} = plot
.resolve_extras(artifacts)
.await
.context("Failed to resolve plot into runtime state")?;
for (key, value) in &mut plot.artifacts {
if let Artifact::File(artifact) = value {
info!("Reading artifact from disk: {:?}", artifact.path);
let buf = fs::read(&artifact.path)?;
artifacts.insert(key.to_string(), HashedArtifact::new(buf));
*value = Artifact::Memory
}
}
if !signing_keys.is_empty() {
info!("Embedding generated/referenced secret keys in plot");
plot.signing_keys = Some(
signing_keys
.into_iter()
.map(|(k, v)| (k, v.into()))
.collect(),
);
}
let plot = serde_json::to_string(&plot)?;
Ok(CompiledArchive { plot, artifacts })
}
pub fn write<W: Write>(&self, w: W) -> Result<()> {
info!("Writing archive...");
let w = Encoder::new(w, ZSTD_COMPRESSION_LEVEL)
.context("Failed to setup zstd stream")?
.auto_finish();
let mut tar = ArchiveBuilder::new(w);
tar.append("plot.json", self.plot.as_bytes())?;
for (key, value) in &self.artifacts {
tar.append(key, value.as_bytes())?;
}
tar.finish()?;
Ok(())
}
}
pub async fn build<W: Write>(
w: W,
plot: &mut Plot,
cache_from: Option<File>,
) -> Result<CompiledArchive> {
let archive = CompiledArchive::build(plot, cache_from).await?;
archive.write(w)?;
Ok(archive)
}
pub async fn run(args: args::Build) -> Result<()> {
let mut plot = Plot::load_from_path(&args.plot.path)?;
if args.no_build {
serde_json::to_writer_pretty(io::stdout(), &plot)?;
println!();
return Ok(());
}
let cache_from = if let Some(path) = args.plot.cache_from {
info!("Opening existing plot as cache: {:?}", path);
let f = File::open(&path)
.with_context(|| anyhow!("Failed to open cache plot at {:?}", path))?;
Some(f)
} else {
None
};
let output_path = args
.output
.as_ref()
.context("Missing mandatory option for output: --output ./plot.tar.zst")?;
let f = File::create(output_path)
.with_context(|| anyhow!("Failed to open output file: {:?}", args.output))?;
if let Err(err) = build(f, &mut plot, cache_from).await {
fs::remove_file(output_path).with_context(|| {
anyhow!("Failed to delete output file of failed build: {output_path:?}")
})?;
Err(err)
} else {
Ok(())
}
}