fix(tokio): add safe access join handles (#85)

Fixes: https://github.com/zkat/cacache-rs/issues/84

src/async_lib.rs

@@ -100,8 +100,8 @@ pub fn unwrap_joinhandle_value<T>(value: T) -> T {
 pub use tokio::task::JoinHandle;
 #[cfg(feature = "tokio")]
 #[inline]
-pub fn unwrap_joinhandle_value<T>(value: Result<T, tokio::task::JoinError>) -> T {
+pub fn unwrap_joinhandle_value<T>(value: T) -> T {
-    value.unwrap()
+    value
 }
 
 use tempfile::NamedTempFile;
@@ -110,19 +110,28 @@ use crate::errors::IoErrorExt;
 
 #[cfg(feature = "async-std")]
 #[inline]
-pub async fn create_named_tempfile(tmp_path: std::path::PathBuf) -> crate::Result<NamedTempFile> {
+pub async fn create_named_tempfile(
+    tmp_path: std::path::PathBuf,
+) -> Option<crate::Result<NamedTempFile>> {
     let cloned = tmp_path.clone();
+
+    Some(
         spawn_blocking(|| NamedTempFile::new_in(tmp_path))
             .await
-        .with_context(|| format!("Failed to create a temp file at {}", cloned.display()))
+            .with_context(|| format!("Failed to create a temp file at {}", cloned.display())),
+    )
 }
 
 #[cfg(feature = "tokio")]
 #[inline]
-pub async fn create_named_tempfile(tmp_path: std::path::PathBuf) -> crate::Result<NamedTempFile> {
+pub async fn create_named_tempfile(
+    tmp_path: std::path::PathBuf,
+) -> Option<crate::Result<NamedTempFile>> {
     let cloned = tmp_path.clone();
-    spawn_blocking(|| NamedTempFile::new_in(tmp_path))
+    match spawn_blocking(|| NamedTempFile::new_in(tmp_path)).await {
-        .await
+        Ok(ctx) => Some(
-        .unwrap()
+            ctx.with_context(|| format!("Failed to create a temp file at {}", cloned.display())),
-        .with_context(|| format!("Failed to create a temp file at {}", cloned.display()))
+        ),
+        _ => None,
+    }
 }

src/content/write.rs

@@ -19,6 +19,7 @@ use tempfile::NamedTempFile;
 use crate::async_lib::{AsyncWrite, JoinHandle};
 use crate::content::path;
 use crate::errors::{IoErrorExt, Result};
+use crate::Error;
 
 #[cfg(feature = "mmap")]
 pub const MAX_MMAP_SIZE: usize = 1024 * 1024;
@@ -171,7 +172,10 @@ impl AsyncWriter {
                     tmp_path.display()
                 )
             })?;
-        let mut tmpfile = crate::async_lib::create_named_tempfile(tmp_path).await?;
+
+        match crate::async_lib::create_named_tempfile(tmp_path).await {
+            Some(tmpfile) => {
+                let mut tmpfile = tmpfile?;
                 let mmap = make_mmap(&mut tmpfile, size)?;
                 Ok(AsyncWriter(Mutex::new(State::Idle(Some(Inner {
                     cache: cache_path,
@@ -182,6 +186,12 @@ impl AsyncWriter {
                     last_op: None,
                 })))))
             }
+            _ => Err(Error::IoError(
+                std::io::Error::new(std::io::ErrorKind::Other, "temp file create error"),
+                "Possible memory issues for file handle".into(),
+            )),
+        }
+    }
 
     pub async fn close(self) -> Result<Integrity> {
         // NOTE: How do I even get access to `inner` safely???
@@ -247,9 +257,11 @@ impl AsyncWriter {
                     },
                     // Poll the asynchronous operation the file is currently blocked on.
                     State::Busy(task) => {
-                        *state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(
+                        let next_state = crate::async_lib::unwrap_joinhandle_value(
-                            Pin::new(task).poll(cx)
+                            futures::ready!(Pin::new(task).poll(cx)),
-                        ))
+                        );
+
+                        update_state(state, next_state);
                     }
                 }
             }
@@ -270,7 +282,9 @@ impl AsyncWrite for AsyncWriter {
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<std::io::Result<usize>> {
-        let state = &mut *self.0.lock().unwrap();
+        match self.0.lock() {
+            Ok(mut state) => {
+                let state = &mut *state;
 
                 loop {
                     match state {
@@ -321,18 +335,23 @@ impl AsyncWrite for AsyncWriter {
                         }
                         // Poll the asynchronous operation the file is currently blocked on.
                         State::Busy(task) => {
-                    *state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
+                            let next_state = crate::async_lib::unwrap_joinhandle_value(
-                        task
+                                futures::ready!(Pin::new(task).poll(cx)),
-                    )
+                            );
-                    .poll(cx)))
+
+                            update_state(state, next_state);
                         }
                     }
                 }
             }
+            _ => Poll::Pending,
+        }
+    }
 
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        let state = &mut *self.0.lock().unwrap();
+        match self.0.lock() {
-
+            Ok(mut state) => {
+                let state = &mut *state;
                 loop {
                     match state {
                         State::Idle(opt) => {
@@ -366,14 +385,18 @@ impl AsyncWrite for AsyncWriter {
                         }
                         // Poll the asynchronous operation the file is currently blocked on.
                         State::Busy(task) => {
-                    *state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
+                            let next_state = crate::async_lib::unwrap_joinhandle_value(
-                        task
+                                futures::ready!(Pin::new(task).poll(cx)),
-                    )
+                            );
-                    .poll(cx)))
+
+                            update_state(state, next_state);
                         }
                     }
                 }
             }
+            _ => Poll::Pending,
+        }
+    }
 
     #[cfg(feature = "async-std")]
     fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
@@ -386,6 +409,28 @@ impl AsyncWrite for AsyncWriter {
     }
 }
 
+#[cfg(feature = "tokio")]
+/// Update the state.
+fn update_state(
+    current_state: &mut State,
+    next_state: std::result::Result<State, tokio::task::JoinError>,
+) {
+    match next_state {
+        Ok(next) => {
+            *current_state = next;
+        }
+        _ => {
+            *current_state = State::Idle(None);
+        }
+    }
+}
+
+#[cfg(not(feature = "tokio"))]
+/// Update the state.
+fn update_state(current_state: &mut State, next_state: State) {
+    *current_state = next_state;
+}
+
 #[cfg(any(feature = "async-std", feature = "tokio"))]
 impl AsyncWriter {
     #[inline]
@@ -393,8 +438,9 @@ impl AsyncWriter {
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> Poll<std::io::Result<()>> {
-        let state = &mut *self.0.lock().unwrap();
+        match self.0.lock() {
-
+            Ok(mut state) => {
+                let state = &mut *state;
                 loop {
                     match state {
                         State::Idle(opt) => {
@@ -413,14 +459,18 @@ impl AsyncWriter {
                         }
                         // Poll the asynchronous operation the file is currently blocked on.
                         State::Busy(task) => {
-                    *state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
+                            let next_state = crate::async_lib::unwrap_joinhandle_value(
-                        task
+                                futures::ready!(Pin::new(task).poll(cx)),
-                    )
+                            );
-                    .poll(cx)))
+
+                            update_state(state, next_state);
                         }
                     }
                 }
             }
+            _ => Poll::Pending,
+        }
+    }
 }
 
 #[cfg(feature = "mmap")]