use std::path::Path;
use async_trait::async_trait;
use crate::error::DownloadError;
#[async_trait]
pub trait Sink: Send + Sync {
async fn write_at(&self, offset: u64, bytes: &[u8]) -> Result<(), DownloadError>;
async fn finalize(&self) -> Result<(), DownloadError> {
Ok(())
}
fn staging_path(&self) -> Option<&Path> {
None
}
}
#[derive(Debug, Default)]
pub struct InMemorySink {
inner: tokio::sync::Mutex<Inner>,
}
#[derive(Debug, Default)]
struct Inner {
buf: Vec<u8>,
finalized: bool,
}
impl InMemorySink {
pub fn new() -> Self {
InMemorySink::default()
}
pub async fn contents(&self) -> Vec<u8> {
self.inner.lock().await.buf.clone()
}
pub async fn is_finalized(&self) -> bool {
self.inner.lock().await.finalized
}
}
#[async_trait]
impl Sink for InMemorySink {
async fn write_at(&self, offset: u64, bytes: &[u8]) -> Result<(), DownloadError> {
let mut inner = self.inner.lock().await;
let end = offset as usize + bytes.len();
if inner.buf.len() < end {
inner.buf.resize(end, 0);
}
inner.buf[offset as usize..end].copy_from_slice(bytes);
Ok(())
}
async fn finalize(&self) -> Result<(), DownloadError> {
self.inner.lock().await.finalized = true;
Ok(())
}
}
pub const TMP_SUFFIX: &str = ".download.tmp";
pub const STATE_SUFFIX: &str = ".download.tmp.state";
pub fn staging_path_for(final_path: &Path) -> std::path::PathBuf {
let mut s = final_path.as_os_str().to_owned();
s.push(TMP_SUFFIX);
std::path::PathBuf::from(s)
}
#[derive(Debug)]
pub struct FileSink {
final_path: std::path::PathBuf,
tmp_path: std::path::PathBuf,
file: tokio::sync::Mutex<Option<std::fs::File>>,
}
impl FileSink {
pub fn new(final_path: impl Into<std::path::PathBuf>) -> Self {
let final_path = final_path.into();
let tmp_path = staging_path_for(&final_path);
FileSink {
final_path,
tmp_path,
file: tokio::sync::Mutex::new(None),
}
}
pub fn final_path(&self) -> &Path {
&self.final_path
}
pub fn tmp_path(&self) -> &Path {
&self.tmp_path
}
}
#[async_trait]
impl Sink for FileSink {
async fn write_at(&self, offset: u64, bytes: &[u8]) -> Result<(), DownloadError> {
use std::io::{Seek, SeekFrom, Write};
let mut guard = self.file.lock().await;
if guard.is_none() {
if let Some(parent) = self.tmp_path.parent() {
std::fs::create_dir_all(parent).map_err(DownloadError::sink)?;
}
let f = std::fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&self.tmp_path)
.map_err(DownloadError::sink)?;
*guard = Some(f);
}
let f = guard.as_mut().expect("file opened above");
f.seek(SeekFrom::Start(offset))
.map_err(DownloadError::sink)?;
f.write_all(bytes).map_err(DownloadError::sink)?;
Ok(())
}
async fn finalize(&self) -> Result<(), DownloadError> {
{
let mut guard = self.file.lock().await;
if let Some(f) = guard.as_mut() {
f.sync_all().map_err(DownloadError::sink)?;
}
*guard = None; }
std::fs::rename(&self.tmp_path, &self.final_path).map_err(DownloadError::sink)?;
Ok(())
}
fn staging_path(&self) -> Option<&Path> {
Some(&self.tmp_path)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn writes_placed_by_offset_out_of_order() {
let sink = InMemorySink::new();
sink.write_at(3, b"DEF").await.unwrap();
sink.write_at(0, b"ABC").await.unwrap();
assert_eq!(sink.contents().await, b"ABCDEF");
assert!(!sink.is_finalized().await);
sink.finalize().await.unwrap();
assert!(sink.is_finalized().await);
}
#[tokio::test]
async fn overlapping_write_overwrites() {
let sink = InMemorySink::new();
sink.write_at(0, b"ABCDEF").await.unwrap();
sink.write_at(2, b"xy").await.unwrap();
assert_eq!(sink.contents().await, b"ABxyEF");
}
fn temp_dir(tag: &str) -> std::path::PathBuf {
let d = std::env::temp_dir().join(format!(
"dig-download-sink-{tag}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&d).unwrap();
d
}
#[tokio::test]
async fn file_sink_stages_then_atomically_finalizes() {
let dir = temp_dir("finalize");
let final_path = dir.join("resource.dig");
let sink = FileSink::new(&final_path);
sink.write_at(3, b"DEF").await.unwrap();
sink.write_at(0, b"ABC").await.unwrap();
assert!(sink.tmp_path().exists());
assert!(!final_path.exists());
assert_eq!(sink.tmp_path(), staging_path_for(&final_path));
sink.finalize().await.unwrap();
assert!(final_path.exists());
assert!(!sink.tmp_path().exists());
assert_eq!(std::fs::read(&final_path).unwrap(), b"ABCDEF");
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn file_sink_resume_reattaches_without_truncating() {
let dir = temp_dir("resume");
let final_path = dir.join("resource.dig");
{
let sink = FileSink::new(&final_path);
sink.write_at(3, b"DEF").await.unwrap();
}
assert!(staging_path_for(&final_path).exists());
let sink2 = FileSink::new(&final_path);
sink2.write_at(0, b"ABC").await.unwrap();
sink2.finalize().await.unwrap();
assert_eq!(std::fs::read(&final_path).unwrap(), b"ABCDEF");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn staging_path_appends_suffix() {
let p = staging_path_for(Path::new("/data/x.dig"));
assert!(p.to_string_lossy().ends_with(".dig.download.tmp"));
}
}