leyline/core/response/
download.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, Ordering};
4
5use tokio::io::AsyncWriteExt;
6
7use super::Response;
8use crate::core::error::{Error, Result};
9use crate::util::atomic;
10
11struct PartialFile {
12    path: PathBuf,
13    armed: bool,
14}
15
16impl PartialFile {
17    fn new(path: PathBuf) -> Self {
18        Self { path, armed: true }
19    }
20
21    fn disarm(mut self) -> PathBuf {
22        self.armed = false;
23        std::mem::take(&mut self.path)
24    }
25}
26
27impl Drop for PartialFile {
28    fn drop(&mut self) {
29        if self.armed {
30            drop(std::fs::remove_file(&self.path));
31        }
32    }
33}
34
35struct CallerWaits(Arc<AtomicBool>);
36
37impl Default for CallerWaits {
38    fn default() -> Self {
39        Self(Arc::new(AtomicBool::new(true)))
40    }
41}
42
43impl Drop for CallerWaits {
44    fn drop(&mut self) {
45        self.0.store(false, Ordering::Release);
46    }
47}
48
49impl Response {
50    pub async fn download_to(self, path: impl AsRef<Path>, limit: Option<u64>) -> Result<u64> {
51        let path = path.as_ref().to_path_buf();
52        let partial = PartialFile::new(atomic::temp_path(&path)?);
53        let written = self.write_file(&partial.path, &path, limit).await?;
54        let caller = CallerWaits::default();
55        let waiting = Arc::clone(&caller.0);
56        tokio::task::spawn_blocking(move || {
57            if !waiting.load(Ordering::Acquire) {
58                drop(partial);
59                return Err(Error::from(std::io::Error::from(
60                    std::io::ErrorKind::Interrupted,
61                )));
62            }
63            let temp = partial.disarm();
64            atomic::commit(&temp, &path)
65        })
66        .await
67        .map_err(|err| Error::from(std::io::Error::other(err)))??;
68        drop(caller);
69        Ok(written)
70    }
71
72    async fn write_file(self, temp: &Path, target: &Path, limit: Option<u64>) -> Result<u64> {
73        let mut file = tokio::fs::File::from_std(atomic::create_temp(
74            temp,
75            target,
76            atomic::FileMode::KeepExisting,
77        )?);
78        let written = self.copy_decoded_to(&mut file, limit).await?;
79        file.flush().await?;
80        file.sync_all().await?;
81        Ok(written)
82    }
83}
leyline/core/response/download.rs

leyline/core/response/
download.rs
