Skip to content

Commit

Permalink
fs: add read/write_at for fs::File
Browse files Browse the repository at this point in the history
  • Loading branch information
SteveLauC committed Mar 25, 2024
1 parent deff252 commit ba4c86a
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 0 deletions.
50 changes: 50 additions & 0 deletions tokio/src/fs/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -538,6 +538,56 @@ impl File {
pub fn set_max_buf_size(&mut self, max_buf_size: usize) {
self.max_buf_size = max_buf_size;
}

/// Reads a number of bytes starting from a given offset, returns the number
/// of bytes.
pub async fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize> {
#[cfg(unix)]
fn _read_at(std: &StdFile, n: usize, offset: u64) -> io::Result<Vec<u8>> {
use std::os::unix::fs::FileExt;
let mut buf: Vec<u8> = vec![0; n];
let n_read = std.read_at(&mut buf, offset)?;
buf.truncate(n_read);

Ok(buf)
}
#[cfg(windows)]
fn _read_at(std: &StdFile, n: usize, offset: u64) -> io::Result<Vec<u8>> {
use std::os::windows::fs::FileExt;
let mut buf: Vec<u8> = vec![0; n];
let n_read = std.seek_read(&mut buf, offset)?;
buf.truncate(n_read);

Ok(buf)
}

let std = self.std.clone();
let n = buf.len();
let bytes_read = asyncify(move || _read_at(&std, n, offset)).await?;
let len = bytes_read.len();
buf[..len].copy_from_slice(&bytes_read);

Ok(len)
}

/// Writes a number of bytes starting from a given offset, returns the number
/// of bytes written.
pub async fn write_at(&self, buf: &[u8], offset: u64) -> io::Result<usize> {
#[cfg(unix)]
fn _write_at(std: &StdFile, buf: &[u8], offset: u64) -> io::Result<usize> {
use std::os::unix::fs::FileExt;
std.write_at(buf, offset)
}
#[cfg(windows)]
fn _write_at(std: &StdFile, buf: &[u8], offset: u64) -> io::Result<usize> {
use std::os::windows::fs::FileExt;
std.seek_write(buf, offset)
}

let std = self.std.clone();
let buf_clone = buf.to_vec();
asyncify(move || _write_at(&std, &buf_clone, offset)).await
}
}

impl AsyncRead for File {
Expand Down
12 changes: 12 additions & 0 deletions tokio/src/fs/mocks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,18 @@ mock! {
impl std::os::unix::io::FromRawFd for File {
unsafe fn from_raw_fd(h: std::os::unix::io::RawFd) -> Self;
}

#[cfg(unix)]
impl std::os::unix::fs::FileExt for File {
fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize>;
fn write_at(&self, buf: &[u8], offset: u64) -> io::Result<usize>;
}

#[cfg(windows)]
impl std::os::windows::fs::FileExt for File {
fn seek_read(&self, buf: &mut [u8], offset: u64) -> io::Result<usize>;
fn seek_write(&self, buf: &[u8], offset: u64) -> io::Result<usize>;
}
}

impl Read for MockFile {
Expand Down
20 changes: 20 additions & 0 deletions tokio/tests/io_read_at.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASM does support this, but it is nightly

use tempfile::tempdir;
use tokio::fs;
use tokio::io::AsyncSeekExt;

#[tokio::test]
async fn read_at() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("a.txt");
fs::write(&file_path, b"HelloWorld").await.unwrap();
let mut file = fs::File::open(file_path.as_path()).await.unwrap();

let mut buf = [0_u8; 10];
assert_eq!(file.read_at(&mut buf, 5).await.unwrap(), 5);
assert_eq!(&buf[..5], b"World");

assert_eq!(file.stream_position().await.unwrap(), 0);
}
25 changes: 25 additions & 0 deletions tokio/tests/io_write_at.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASM does support this, but it is nightly

use tempfile::tempdir;
use tokio::fs;
use tokio::fs::OpenOptions;
use tokio::io::AsyncSeekExt;

#[tokio::test]
async fn write_at() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("a.txt");
fs::write(&file_path, b"Hello File").await.unwrap();
let mut file = OpenOptions::new()
.write(true)
.open(file_path.as_path())
.await
.unwrap();

assert_eq!(file.write_at(b"World", 5).await.unwrap(), 5);
let contents = fs::read(file_path.as_path()).await.unwrap();
assert_eq!(contents, b"HelloWorld");

assert_eq!(file.stream_position().await.unwrap(), 0);
}

0 comments on commit ba4c86a

Please sign in to comment.