mirror of
https://github.com/cargo-bins/cargo-binstall.git
synced 2025-05-05 03:30:03 +00:00
911c52d8e1
unless done is called. Also moves creation of the dir/file into the blocking thread to avoid blocking. Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>
126 lines
3.8 KiB
Rust
126 lines
3.8 KiB
Rust
use std::fs;
|
|
use std::io::{self, Write};
|
|
use std::path::Path;
|
|
|
|
use bytes::Bytes;
|
|
use scopeguard::{guard, Always, ScopeGuard};
|
|
use tokio::{sync::mpsc, task::spawn_blocking};
|
|
|
|
use super::AutoAbortJoinHandle;
|
|
|
|
enum Content {
|
|
/// Data to write to file
|
|
Data(Bytes),
|
|
|
|
/// Abort the writing and remove the file.
|
|
Abort,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct AsyncFileWriterInner {
|
|
/// Use AutoAbortJoinHandle so that the task
|
|
/// will be cancelled on failure.
|
|
handle: AutoAbortJoinHandle<io::Result<()>>,
|
|
tx: mpsc::Sender<Content>,
|
|
}
|
|
|
|
impl AsyncFileWriterInner {
|
|
fn new(path: &Path) -> io::Result<Self> {
|
|
let path = path.to_owned();
|
|
let (tx, rx) = mpsc::channel::<Content>(100);
|
|
|
|
let handle = AutoAbortJoinHandle::new(spawn_blocking(move || {
|
|
// close rx on error so that tx.send will return an error
|
|
let mut rx = guard(rx, |mut rx| {
|
|
rx.close();
|
|
});
|
|
|
|
fs::create_dir_all(path.parent().unwrap())?;
|
|
let mut file = fs::File::create(&path)?;
|
|
|
|
// remove it unless the operation isn't aborted and no write
|
|
// fails.
|
|
let remove_guard = guard(path, |path| {
|
|
fs::remove_file(path).ok();
|
|
});
|
|
|
|
while let Some(content) = rx.blocking_recv() {
|
|
match content {
|
|
Content::Data(bytes) => file.write_all(&*bytes)?,
|
|
Content::Abort => return Err(io::Error::new(io::ErrorKind::Other, "Aborted")),
|
|
}
|
|
}
|
|
|
|
file.flush()?;
|
|
|
|
// Operation isn't aborted and all writes succeed,
|
|
// disarm the remove_guard.
|
|
ScopeGuard::into_inner(remove_guard);
|
|
|
|
Ok(())
|
|
}));
|
|
|
|
Ok(Self { handle, tx })
|
|
}
|
|
|
|
/// Upon error, this writer shall not be reused.
|
|
/// Otherwise, `Self::done` would panic.
|
|
async fn write(&mut self, bytes: Bytes) -> io::Result<()> {
|
|
if self.tx.send(Content::Data(bytes)).await.is_err() {
|
|
// task failed
|
|
Err(Self::wait(&mut self.handle).await.expect_err(
|
|
"Implementation bug: write task finished successfully before all writes are done",
|
|
))
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
async fn done(mut self) -> io::Result<()> {
|
|
// Drop tx as soon as possible so that the task would wrap up what it
|
|
// was doing and flush out all the pending data.
|
|
drop(self.tx);
|
|
|
|
Self::wait(&mut self.handle).await
|
|
}
|
|
|
|
async fn wait(handle: &mut AutoAbortJoinHandle<io::Result<()>>) -> io::Result<()> {
|
|
match handle.await {
|
|
Ok(res) => res,
|
|
Err(join_err) => Err(io::Error::new(io::ErrorKind::Other, join_err)),
|
|
}
|
|
}
|
|
|
|
fn abort(self) {
|
|
let tx = self.tx;
|
|
// If Self::write fail, then the task is already tear down,
|
|
// tx closed and no need to abort.
|
|
if !tx.is_closed() {
|
|
// Use send here because blocking_send would panic if used
|
|
// in async threads.
|
|
tokio::spawn(async move {
|
|
tx.send(Content::Abort).await.ok();
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
/// AsyncFileWriter removes the file if `done` isn't called.
|
|
#[derive(Debug)]
|
|
pub struct AsyncFileWriter(ScopeGuard<AsyncFileWriterInner, fn(AsyncFileWriterInner), Always>);
|
|
|
|
impl AsyncFileWriter {
|
|
pub fn new(path: &Path) -> io::Result<Self> {
|
|
AsyncFileWriterInner::new(path).map(|inner| Self(guard(inner, AsyncFileWriterInner::abort)))
|
|
}
|
|
|
|
/// Upon error, this writer shall not be reused.
|
|
/// Otherwise, `Self::done` would panic.
|
|
pub async fn write(&mut self, bytes: Bytes) -> io::Result<()> {
|
|
self.0.write(bytes).await
|
|
}
|
|
|
|
pub async fn done(self) -> io::Result<()> {
|
|
ScopeGuard::into_inner(self.0).done().await
|
|
}
|
|
}
|