berkus
/
cacache-rs
mirror of https://github.com/zkat/cacache-rs.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

Fixes: https://github.com/zkat/cacache-rs/issues/84
This commit is contained in:
Jeff Mendez 2024-06-25 11:42:11 -04:00 committed by GitHub
parent ab5f1c9185
commit 146a593c8e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 190 additions and 131 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
src/async_lib.rs
View File

@ -100,8 +100,8 @@ pub fn unwrap_joinhandle_value<T>(value: T) -> T {
pub use tokio::task::JoinHandle; pub use tokio::task::JoinHandle;
#[cfg(feature = "tokio")] #[cfg(feature = "tokio")]
#[inline] #[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; use tempfile::NamedTempFile;
@ -110,19 +110,28 @@ use crate::errors::IoErrorExt;
#[cfg(feature = "async-std")] #[cfg(feature = "async-std")]
#[inline] #[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(); let cloned = tmp_path.clone();
spawn_blocking(|| NamedTempFile::new_in(tmp_path))
.await Some(
.with_context(|| format!("Failed to create a temp file at {}", cloned.display())) spawn_blocking(|| NamedTempFile::new_in(tmp_path))
.await
.with_context(|| format!("Failed to create a temp file at {}", cloned.display())),
)
} }
#[cfg(feature = "tokio")] #[cfg(feature = "tokio")]
#[inline] #[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(); 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,
}
} }

290
src/content/write.rs
View File

@ -19,6 +19,7 @@ use tempfile::NamedTempFile;
use crate::async_lib::{AsyncWrite, JoinHandle}; use crate::async_lib::{AsyncWrite, JoinHandle};
use crate::content::path; use crate::content::path;
use crate::errors::{IoErrorExt, Result}; use crate::errors::{IoErrorExt, Result};
use crate::Error;
#[cfg(feature = "mmap")] #[cfg(feature = "mmap")]
pub const MAX_MMAP_SIZE: usize = 1024 * 1024; pub const MAX_MMAP_SIZE: usize = 1024 * 1024;
@ -171,16 +172,25 @@ impl AsyncWriter {
tmp_path.display() tmp_path.display()
) )
})?; })?;
let mut tmpfile = crate::async_lib::create_named_tempfile(tmp_path).await?;
let mmap = make_mmap(&mut tmpfile, size)?; match crate::async_lib::create_named_tempfile(tmp_path).await {
Ok(AsyncWriter(Mutex::new(State::Idle(Some(Inner { Some(tmpfile) => {
cache: cache_path, let mut tmpfile = tmpfile?;
builder: IntegrityOpts::new().algorithm(algo), let mmap = make_mmap(&mut tmpfile, size)?;
mmap, Ok(AsyncWriter(Mutex::new(State::Idle(Some(Inner {
tmpfile, cache: cache_path,
buf: vec![], builder: IntegrityOpts::new().algorithm(algo),
last_op: None, mmap,
}))))) tmpfile,
buf: vec![],
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> { pub async fn close(self) -> Result<Integrity> {
@ -247,9 +257,11 @@ impl AsyncWriter {
}, },
// Poll the asynchronous operation the file is currently blocked on. // Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => { 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,108 +282,119 @@ impl AsyncWrite for AsyncWriter {
cx: &mut Context<'_>, cx: &mut Context<'_>,
buf: &[u8], buf: &[u8],
) -> Poll<std::io::Result<usize>> { ) -> Poll<std::io::Result<usize>> {
let state = &mut *self.0.lock().unwrap(); match self.0.lock() {
Ok(mut state) => {
let state = &mut *state;
loop { loop {
match state { match state {
State::Idle(opt) => { State::Idle(opt) => {
// Grab a reference to the inner representation of the file or return an error // Grab a reference to the inner representation of the file or return an error
// if the file is closed. // if the file is closed.
let inner = opt let inner = opt
.as_mut() .as_mut()
.ok_or_else(|| crate::errors::io_error("file closed"))?; .ok_or_else(|| crate::errors::io_error("file closed"))?;
// Check if the operation has completed. // Check if the operation has completed.
if let Some(Operation::Write(res)) = inner.last_op.take() { if let Some(Operation::Write(res)) = inner.last_op.take() {
let n = res?; let n = res?;
// If more data was written than is available in the buffer, let's retry // If more data was written than is available in the buffer, let's retry
// the write operation. // the write operation.
if n <= buf.len() { if n <= buf.len() {
return Poll::Ready(Ok(n)); return Poll::Ready(Ok(n));
} }
} else {
let mut inner = opt.take().unwrap();
// Set the length of the inner buffer to the length of the provided buffer.
if inner.buf.len() < buf.len() {
inner.buf.reserve(buf.len() - inner.buf.len());
}
unsafe {
inner.buf.set_len(buf.len());
}
// Copy the data to write into the inner buffer.
inner.buf[..buf.len()].copy_from_slice(buf);
// Start the operation asynchronously.
*state = State::Busy(crate::async_lib::spawn_blocking(|| {
inner.builder.input(&inner.buf);
if let Some(mmap) = &mut inner.mmap {
mmap.copy_from_slice(&inner.buf);
inner.last_op = Some(Operation::Write(Ok(inner.buf.len())));
State::Idle(Some(inner))
} else { } else {
let res = inner.tmpfile.write(&inner.buf); let mut inner = opt.take().unwrap();
inner.last_op = Some(Operation::Write(res));
State::Idle(Some(inner)) // Set the length of the inner buffer to the length of the provided buffer.
if inner.buf.len() < buf.len() {
inner.buf.reserve(buf.len() - inner.buf.len());
}
unsafe {
inner.buf.set_len(buf.len());
}
// Copy the data to write into the inner buffer.
inner.buf[..buf.len()].copy_from_slice(buf);
// Start the operation asynchronously.
*state = State::Busy(crate::async_lib::spawn_blocking(|| {
inner.builder.input(&inner.buf);
if let Some(mmap) = &mut inner.mmap {
mmap.copy_from_slice(&inner.buf);
inner.last_op = Some(Operation::Write(Ok(inner.buf.len())));
State::Idle(Some(inner))
} else {
let res = inner.tmpfile.write(&inner.buf);
inner.last_op = Some(Operation::Write(res));
State::Idle(Some(inner))
}
}));
} }
})); }
// Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => {
let next_state = crate::async_lib::unwrap_joinhandle_value(
futures::ready!(Pin::new(task).poll(cx)),
);
update_state(state, next_state);
}
} }
} }
// Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => {
*state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
task
)
.poll(cx)))
}
} }
_ => Poll::Pending,
} }
} }
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> { 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) => {
loop { let state = &mut *state;
match state { loop {
State::Idle(opt) => { match state {
// Grab a reference to the inner representation of the file or return if the State::Idle(opt) => {
// file is closed. // Grab a reference to the inner representation of the file or return if the
let inner = match opt.as_mut() { // file is closed.
None => return Poll::Ready(Ok(())), let inner = match opt.as_mut() {
Some(s) => s, None => return Poll::Ready(Ok(())),
}; Some(s) => s,
// Check if the operation has completed.
if let Some(Operation::Flush(res)) = inner.last_op.take() {
return Poll::Ready(res);
} else {
let mut inner = opt.take().unwrap();
if let Some(mmap) = &inner.mmap {
match mmap.flush_async() {
Ok(_) => (),
Err(e) => return Poll::Ready(Err(e)),
}; };
}
// Start the operation asynchronously. // Check if the operation has completed.
*state = State::Busy(crate::async_lib::spawn_blocking(|| { if let Some(Operation::Flush(res)) = inner.last_op.take() {
let res = inner.tmpfile.flush(); return Poll::Ready(res);
inner.last_op = Some(Operation::Flush(res)); } else {
State::Idle(Some(inner)) let mut inner = opt.take().unwrap();
}));
if let Some(mmap) = &inner.mmap {
match mmap.flush_async() {
Ok(_) => (),
Err(e) => return Poll::Ready(Err(e)),
};
}
// Start the operation asynchronously.
*state = State::Busy(crate::async_lib::spawn_blocking(|| {
let res = inner.tmpfile.flush();
inner.last_op = Some(Operation::Flush(res));
State::Idle(Some(inner))
}));
}
}
// Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => {
let next_state = crate::async_lib::unwrap_joinhandle_value(
futures::ready!(Pin::new(task).poll(cx)),
);
update_state(state, next_state);
}
} }
} }
// Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => {
*state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
task
)
.poll(cx)))
}
} }
_ => Poll::Pending,
} }
} }
@ -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"))] #[cfg(any(feature = "async-std", feature = "tokio"))]
impl AsyncWriter { impl AsyncWriter {
#[inline] #[inline]
@ -393,32 +438,37 @@ impl AsyncWriter {
self: Pin<&mut Self>, self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>, cx: &mut std::task::Context<'_>,
) -> Poll<std::io::Result<()>> { ) -> 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) => {
// Grab a reference to the inner representation of the file or return if the
// file is closed.
let inner = match opt.take() {
None => return Poll::Ready(Ok(())),
Some(s) => s,
};
loop { // Start the operation asynchronously.
match state { *state = State::Busy(crate::async_lib::spawn_blocking(|| {
State::Idle(opt) => { drop(inner);
// Grab a reference to the inner representation of the file or return if the State::Idle(None)
// file is closed. }));
let inner = match opt.take() { }
None => return Poll::Ready(Ok(())), // Poll the asynchronous operation the file is currently blocked on.
Some(s) => s, State::Busy(task) => {
}; let next_state = crate::async_lib::unwrap_joinhandle_value(
futures::ready!(Pin::new(task).poll(cx)),
);
// Start the operation asynchronously. update_state(state, next_state);
*state = State::Busy(crate::async_lib::spawn_blocking(|| { }
drop(inner); }
State::Idle(None)
}));
}
// Poll the asynchronous operation the file is currently blocked on.
State::Busy(task) => {
*state = crate::async_lib::unwrap_joinhandle_value(futures::ready!(Pin::new(
task
)
.poll(cx)))
} }
} }
_ => Poll::Pending,
} }
} }
} }