Use AutoAbortJoinHandle in AsyncFileWriter

to cancel the task on failure.

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>
This commit is contained in:
Jiahao XU 2022-06-08 20:46:00 +10:00
parent 12931fc024
commit d6a372a160
No known key found for this signature in database
GPG key ID: 591C0B03040416D6

View file

@ -229,7 +229,9 @@ pub trait Template: Serialize {
#[derive(Debug)] #[derive(Debug)]
pub struct AsyncFileWriter { pub struct AsyncFileWriter {
handle: task::JoinHandle<io::Result<()>>, /// Use AutoAbortJoinHandle so that the task
/// will be cancelled on failure.
handle: AutoAbortJoinHandle<io::Result<()>>,
tx: mpsc::Sender<Bytes>, tx: mpsc::Sender<Bytes>,
} }
@ -240,7 +242,7 @@ impl AsyncFileWriter {
let mut file = fs::File::create(path)?; let mut file = fs::File::create(path)?;
let (tx, mut rx) = mpsc::channel::<Bytes>(100); let (tx, mut rx) = mpsc::channel::<Bytes>(100);
let handle = task::spawn_blocking(move || { let handle = AutoAbortJoinHandle(task::spawn_blocking(move || {
while let Some(bytes) = rx.blocking_recv() { while let Some(bytes) = rx.blocking_recv() {
file.write_all(&*bytes)?; file.write_all(&*bytes)?;
} }
@ -249,7 +251,7 @@ impl AsyncFileWriter {
file.flush()?; file.flush()?;
Ok(()) Ok(())
}); }));
Ok(Self { handle, tx }) Ok(Self { handle, tx })
} }
@ -293,8 +295,8 @@ impl AsyncFileWriter {
Self::wait(&mut self.handle).await Self::wait(&mut self.handle).await
} }
async fn wait(handle: &mut task::JoinHandle<io::Result<()>>) -> io::Result<()> { async fn wait(handle: &mut AutoAbortJoinHandle<io::Result<()>>) -> io::Result<()> {
match handle.await { match (&mut handle.0).await {
Ok(res) => res, Ok(res) => res,
Err(join_err) => Err(io::Error::new(io::ErrorKind::Other, join_err)), Err(join_err) => Err(io::Error::new(io::ErrorKind::Other, join_err)),
} }