use std::{
fmt, fs, io,
path::{Path, PathBuf},
sync::atomic::{AtomicU64, Ordering},
};
use ic_host_artifacts::artifact::{ArtifactError, CopyError, copy_reader};
use ic_host_fs::{
durable::{NamedWriteError, PublicationMode, WriteOptions, write_with},
read::open_regular_file,
};
use super::cache_fs::remove_path_if_present;
const STAGE_ATTEMPTS: usize = 64;
static STAGE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
#[derive(Debug)]
pub(super) struct OwnedDirectory {
path: PathBuf,
armed: bool,
}
impl OwnedDirectory {
pub(super) const fn new(path: PathBuf) -> Self {
Self { path, armed: true }
}
pub(super) fn path(&self) -> &Path {
&self.path
}
pub(super) const fn is_armed(&self) -> bool {
self.armed
}
pub(super) const fn disarm(&mut self) {
self.armed = false;
}
pub(super) fn preserve(mut self) {
self.disarm();
}
pub(super) fn try_cleanup(&mut self) -> io::Result<()> {
remove_path_if_present(&self.path)?;
self.disarm();
Ok(())
}
pub(super) fn cleanup(mut self) -> io::Result<()> {
self.try_cleanup()
}
}
impl Drop for OwnedDirectory {
fn drop(&mut self) {
if self.armed {
let _ = remove_path_if_present(&self.path);
self.disarm();
}
}
}
#[derive(Debug)]
pub struct ArtifactStagingDirectory {
directory: OwnedDirectory,
}
impl ArtifactStagingDirectory {
pub fn new(parent: &Path) -> io::Result<Self> {
for _ in 0..STAGE_ATTEMPTS {
let sequence = STAGE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
let root = parent.join(format!(
".ic-build-artifact-stage-{}-{sequence}",
std::process::id()
));
match fs::create_dir(&root) {
Ok(()) => {
return Ok(Self {
directory: OwnedDirectory::new(root),
});
}
Err(source) if source.kind() == io::ErrorKind::AlreadyExists => {}
Err(source) => return Err(source),
}
}
Err(io::Error::new(
io::ErrorKind::AlreadyExists,
format!(
"failed to allocate artifact staging beneath {}",
parent.display()
),
))
}
#[must_use]
pub fn path(&self) -> &Path {
self.directory.path()
}
pub fn close(mut self) -> io::Result<()> {
self.directory.disarm();
remove_path_if_present(self.directory.path())
}
}
#[derive(Debug)]
pub struct CapturedArtifact {
staging: ArtifactStagingDirectory,
path: PathBuf,
}
#[derive(Debug)]
pub struct CaptureRetentionError {
pub source: io::Error,
pub retained_path: PathBuf,
pub staging_directory: PathBuf,
}
impl fmt::Display for CaptureRetentionError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
formatter,
"capture retained at {} (staging {}): {}",
self.retained_path.display(),
self.staging_directory.display(),
self.source
)
}
}
impl std::error::Error for CaptureRetentionError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.source)
}
}
impl CapturedArtifact {
pub fn capture(source: &Path, parent: &Path, maximum_bytes: usize) -> io::Result<Self> {
let source = open_regular_file(source).map_err(io::Error::from)?;
let limit = maximum_bytes as u64;
if source.metadata()?.len() > limit {
return Err(ArtifactError::LimitExceeded { limit }.into());
}
let staging = ArtifactStagingDirectory::new(parent)?;
let path = staging.path().join("input.artifact");
write_with(
&path,
WriteOptions {
mode: PublicationMode::Replace,
permissions: 0o666,
},
|file| copy_reader(source, file, limit).map(|_| ()),
)
.map_err(|error| match error {
NamedWriteError::Producer {
source: CopyError::Input(source),
cleanup_error: None,
} => io::Error::from(source),
other => io::Error::other(other),
})?;
Ok(Self { staging, path })
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
#[must_use]
pub fn retain(mut self) -> PathBuf {
self.staging.directory.disarm();
self.path
}
pub fn retain_to(self, destination: &Path) -> Result<PathBuf, CaptureRetentionError> {
let staging = self.staging.path().to_owned();
let retained = self.retain();
fs::rename(&retained, destination).map_err(|source| CaptureRetentionError {
source,
retained_path: retained.clone(),
staging_directory: staging.clone(),
})?;
fs::remove_dir(&staging).map_err(|source| CaptureRetentionError {
source,
retained_path: destination.to_owned(),
staging_directory: staging,
})?;
Ok(destination.to_owned())
}
pub fn close(self) -> io::Result<()> {
self.staging.close()
}
}
#[cfg(test)]
mod tests;