tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-17 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

fs: add set_max_buf_size to tokio::fs::File (#6411)

Browse Source
This commit is contained in:
Motoyuki Kimura
2024-03-22 17:28:29 +03:30
committed by GitHub
parent bb25a06f34
commit f9d78fbe92
6 changed files with 79 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/fs/file.rs
+34 -4
View File
@@ -3,7 +3,7 @@
//! [`File`]: File
use crate::fs::{asyncify, OpenOptions};
use crate::io::blocking::Buf;
use crate::io::blocking::{Buf, DEFAULT_MAX_BUF_SIZE};
use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
use crate::sync::Mutex;
@@ -90,6 +90,7 @@ use std::fs::File as StdFile;
pub struct File {
std: Arc<StdFile>,
inner: Mutex<Inner>,
max_buf_size: usize,
}
struct Inner {
@@ -241,6 +242,7 @@ impl File {
last_write_err: None,
pos: 0,
}),
max_buf_size: DEFAULT_MAX_BUF_SIZE,
}
}
@@ -508,6 +510,34 @@ impl File {
let std = self.std.clone();
asyncify(move || std.set_permissions(perm)).await
}
/// Set the maximum buffer size for the underlying [`AsyncRead`] / [`AsyncWrite`] operation.
///
/// Although Tokio uses a sensible default value for this buffer size, this function would be
/// useful for changing that default depending on the situation.
///
/// # Examples
///
/// ```no_run
/// use tokio::fs::File;
/// use tokio::io::AsyncWriteExt;
///
/// # async fn dox() -> std::io::Result<()> {
/// let mut file = File::open("foo.txt").await?;
///
/// // Set maximum buffer size to 8 MiB
/// file.set_max_buf_size(8 * 1024 * 1024);
///
/// let mut buf = vec![1; 1024 * 1024 * 1024];
///
/// // Write the 1 GiB buffer in chunks up to 8 MiB each.
/// file.write_all(&mut buf).await?;
/// # Ok(())
/// # }
/// ```
pub fn set_max_buf_size(&mut self, max_buf_size: usize) {
self.max_buf_size = max_buf_size;
}
}
impl AsyncRead for File {
@@ -531,7 +561,7 @@ impl AsyncRead for File {
return Poll::Ready(Ok(()));
}
buf.ensure_capacity_for(dst);
buf.ensure_capacity_for(dst, me.max_buf_size);
let std = me.std.clone();
inner.state = State::Busy(spawn_blocking(move || {
@@ -668,7 +698,7 @@ impl AsyncWrite for File {
None
};
let n = buf.copy_from(src);
let n = buf.copy_from(src, me.max_buf_size);
let std = me.std.clone();
let blocking_task_join_handle = spawn_mandatory_blocking(move || {
@@ -739,7 +769,7 @@ impl AsyncWrite for File {
None
};
let n = buf.copy_from_bufs(bufs);
let n = buf.copy_from_bufs(bufs, me.max_buf_size);
let std = me.std.clone();
let blocking_task_join_handle = spawn_mandatory_blocking(move || {
tokio/src/fs/file/tests.rs
+2 -2
View File
@@ -231,7 +231,7 @@ fn flush_while_idle() {
#[cfg_attr(miri, ignore)] // takes a really long time with miri
fn read_with_buffer_larger_than_max() {
// Chunks
let chunk_a = crate::io::blocking::MAX_BUF;
let chunk_a = crate::io::blocking::DEFAULT_MAX_BUF_SIZE;
let chunk_b = chunk_a * 2;
let chunk_c = chunk_a * 3;
let chunk_d = chunk_a * 4;
@@ -303,7 +303,7 @@ fn read_with_buffer_larger_than_max() {
#[cfg_attr(miri, ignore)] // takes a really long time with miri
fn write_with_buffer_larger_than_max() {
// Chunks
let chunk_a = crate::io::blocking::MAX_BUF;
let chunk_a = crate::io::blocking::DEFAULT_MAX_BUF_SIZE;
let chunk_b = chunk_a * 2;
let chunk_c = chunk_a * 3;
let chunk_d = chunk_a * 4;
tokio/src/fs/mocks.rs
+1
View File
@@ -30,6 +30,7 @@ mock! {
pub fn open(pb: PathBuf) -> io::Result<Self>;
pub fn set_len(&self, size: u64) -> io::Result<()>;
pub fn set_permissions(&self, _perm: Permissions) -> io::Result<()>;
pub fn set_max_buf_size(&self, max_buf_size: usize);
pub fn sync_all(&self) -> io::Result<()>;
pub fn sync_data(&self) -> io::Result<()>;
pub fn try_clone(&self) -> io::Result<Self>;
tokio/src/io/blocking.rs
+10 -10
View File
@@ -23,7 +23,7 @@ pub(crate) struct Buf {
pos: usize,
}
pub(crate) const MAX_BUF: usize = 2 * 1024 * 1024;
pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024;
#[derive(Debug)]
enum State<T> {
@@ -64,7 +64,7 @@ where
return Poll::Ready(Ok(()));
}
buf.ensure_capacity_for(dst);
buf.ensure_capacity_for(dst, DEFAULT_MAX_BUF_SIZE);
let mut inner = self.inner.take().unwrap();
self.state = State::Busy(sys::run(move || {
@@ -111,7 +111,7 @@ where
assert!(buf.is_empty());
let n = buf.copy_from(src);
let n = buf.copy_from(src, DEFAULT_MAX_BUF_SIZE);
let mut inner = self.inner.take().unwrap();
self.state = State::Busy(sys::run(move || {
@@ -214,10 +214,10 @@ impl Buf {
n
}
pub(crate) fn copy_from(&mut self, src: &[u8]) -> usize {
pub(crate) fn copy_from(&mut self, src: &[u8], max_buf_size: usize) -> usize {
assert!(self.is_empty());
let n = cmp::min(src.len(), MAX_BUF);
let n = cmp::min(src.len(), max_buf_size);
self.buf.extend_from_slice(&src[..n]);
n
@@ -227,10 +227,10 @@ impl Buf {
&self.buf[self.pos..]
}
pub(crate) fn ensure_capacity_for(&mut self, bytes: &ReadBuf<'_>) {
pub(crate) fn ensure_capacity_for(&mut self, bytes: &ReadBuf<'_>, max_buf_size: usize) {
assert!(self.is_empty());
let len = cmp::min(bytes.remaining(), MAX_BUF);
let len = cmp::min(bytes.remaining(), max_buf_size);
if self.buf.len() < len {
self.buf.reserve(len - self.buf.len());
@@ -274,10 +274,10 @@ cfg_fs! {
ret
}
pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>]) -> usize {
pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>], max_buf_size: usize) -> usize {
assert!(self.is_empty());
let mut rem = MAX_BUF;
let mut rem = max_buf_size;
for buf in bufs {
if rem == 0 {
break
@@ -288,7 +288,7 @@ cfg_fs! {
rem -= len;
}
MAX_BUF - rem
max_buf_size - rem
}
}
}
tokio/src/io/stdio_common.rs
+10 -9
View File
@@ -4,7 +4,7 @@ use std::pin::Pin;
use std::task::{Context, Poll};
/// # Windows
/// [`AsyncWrite`] adapter that finds last char boundary in given buffer and does not write the rest,
/// if buffer contents seems to be `utf8`. Otherwise it only trims buffer down to `MAX_BUF`.
/// if buffer contents seems to be `utf8`. Otherwise it only trims buffer down to `DEFAULT_MAX_BUF_SIZE`.
/// That's why, wrapped writer will always receive well-formed utf-8 bytes.
/// # Other platforms
/// Passes data to `inner` as is.
@@ -45,12 +45,13 @@ where
// 2. If buffer is small, it will not be shrunk.
// That's why, it's "textness" will not change, so we don't have
// to fixup it.
if cfg!(not(any(target_os = "windows", test))) || buf.len() <= crate::io::blocking::MAX_BUF
if cfg!(not(any(target_os = "windows", test)))
|| buf.len() <= crate::io::blocking::DEFAULT_MAX_BUF_SIZE
{
return call_inner(buf);
}
buf = &buf[..crate::io::blocking::MAX_BUF];
buf = &buf[..crate::io::blocking::DEFAULT_MAX_BUF_SIZE];
// Now there are two possibilities.
// If caller gave is binary buffer, we **should not** shrink it
@@ -108,7 +109,7 @@ where
#[cfg(test)]
#[cfg(not(loom))]
mod tests {
use crate::io::blocking::MAX_BUF;
use crate::io::blocking::DEFAULT_MAX_BUF_SIZE;
use crate::io::AsyncWriteExt;
use std::io;
use std::pin::Pin;
@@ -123,7 +124,7 @@ mod tests {
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
assert!(buf.len() <= MAX_BUF);
assert!(buf.len() <= DEFAULT_MAX_BUF_SIZE);
assert!(std::str::from_utf8(buf).is_ok());
Poll::Ready(Ok(buf.len()))
}
@@ -158,7 +159,7 @@ mod tests {
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
assert!(buf.len() <= MAX_BUF);
assert!(buf.len() <= DEFAULT_MAX_BUF_SIZE);
self.write_history.push(buf.len());
Poll::Ready(Ok(buf.len()))
}
@@ -178,7 +179,7 @@ mod tests {
#[test]
#[cfg_attr(miri, ignore)]
fn test_splitter() {
let data = str::repeat("█", MAX_BUF);
let data = str::repeat("█", DEFAULT_MAX_BUF_SIZE);
let mut wr = super::SplitByUtf8BoundaryIfWindows::new(TextMockWriter);
let fut = async move {
wr.write_all(data.as_bytes()).await.unwrap();
@@ -197,7 +198,7 @@ mod tests {
// was not shrunk too much.
let checked_count = super::MAGIC_CONST * super::MAX_BYTES_PER_CHAR;
let mut data: Vec<u8> = str::repeat("a", checked_count).into();
data.extend(std::iter::repeat(0b1010_1010).take(MAX_BUF - checked_count + 1));
data.extend(std::iter::repeat(0b1010_1010).take(DEFAULT_MAX_BUF_SIZE - checked_count + 1));
let mut writer = LoggingMockWriter::new();
let mut splitter = super::SplitByUtf8BoundaryIfWindows::new(&mut writer);
crate::runtime::Builder::new_current_thread()
@@ -214,7 +215,7 @@ mod tests {
data.len()
);
// Check that at most MAX_BYTES_PER_CHAR + 1 (i.e. 5) bytes were shrunk
// from the buffer: one because it was outside of MAX_BUF boundary, and
// from the buffer: one because it was outside of DEFAULT_MAX_BUF_SIZE boundary, and
// up to one "utf8 code point".
assert!(data.len() - writer.write_history[0] <= super::MAX_BYTES_PER_CHAR + 1);
}
tokio/tests/fs_file.rs
+22
View File
@@ -180,6 +180,28 @@ fn tempfile() -> NamedTempFile {
NamedTempFile::new().unwrap()
}
#[tokio::test]
async fn set_max_buf_size_read() {
let mut tempfile = tempfile();
tempfile.write_all(HELLO).unwrap();
let mut file = File::open(tempfile.path()).await.unwrap();
let mut buf = [0; 1024];
file.set_max_buf_size(1);
// A single read operation reads a maximum of 1 byte.
assert_eq!(file.read(&mut buf).await.unwrap(), 1);
}
#[tokio::test]
async fn set_max_buf_size_write() {
let tempfile = tempfile();
let mut file = File::create(tempfile.path()).await.unwrap();
file.set_max_buf_size(1);
// A single write operation writes a maximum of 1 byte.
assert_eq!(file.write(HELLO).await.unwrap(), 1);
}
#[tokio::test]
#[cfg(unix)]
async fn file_debug_fmt() {