File size: 3,894 Bytes
ea39c0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | use std::collections::HashMap;
use std::fs::File;
use std::io;
use std::sync::Arc;
use codex_file_system::FILE_READ_CHUNK_SIZE;
use tokio::sync::Mutex;
const MAX_OPEN_FILE_READS: usize = 128;
#[derive(Debug, Eq, PartialEq)]
pub(crate) struct FileReadBlock {
pub(crate) bytes: Vec<u8>,
pub(crate) eof: bool,
}
#[derive(Clone, Default)]
pub(crate) struct FileReadHandleManager {
handles: Arc<Mutex<HashMap<String, Arc<File>>>>,
}
impl FileReadHandleManager {
pub(crate) async fn open(
&self,
handle_id: String,
file: tokio::fs::File,
) -> io::Result<String> {
let file = Arc::new(file.into_std().await);
let mut handles = self.handles.lock().await;
if handles.contains_key(&handle_id) {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("file read handle `{handle_id}` already exists"),
));
}
if handles.len() >= MAX_OPEN_FILE_READS {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("at most {MAX_OPEN_FILE_READS} file reads may be open per connection"),
));
}
handles.insert(handle_id.clone(), file);
Ok(handle_id)
}
pub(crate) async fn read_block(
&self,
handle_id: &str,
offset: u64,
len: usize,
) -> io::Result<FileReadBlock> {
validate_read_block_len(len)?;
let file = {
let handles = self.handles.lock().await;
handles
.get(handle_id)
.cloned()
.ok_or_else(|| unknown_handle_error(handle_id))?
};
let result =
match tokio::task::spawn_blocking(move || read_block_at(&file, offset, len)).await {
Ok(result) => result,
Err(error) => Err(io::Error::other(format!(
"file read task stopped unexpectedly: {error}"
))),
};
if result.is_err() {
self.close(handle_id).await;
}
result
}
pub(crate) async fn close(&self, handle_id: &str) {
self.handles.lock().await.remove(handle_id);
}
pub(crate) async fn close_all(&self) {
self.handles.lock().await.clear();
}
}
fn read_block_at(file: &File, offset: u64, len: usize) -> io::Result<FileReadBlock> {
let mut bytes = vec![0; len];
let mut bytes_read = 0;
while bytes_read < len {
let read_offset = offset.checked_add(bytes_read as u64).ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "file read offset overflowed")
})?;
match read_file_at(file, &mut bytes[bytes_read..], read_offset) {
Ok(0) => break,
Ok(read) => bytes_read += read,
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
bytes.truncate(bytes_read);
Ok(FileReadBlock {
eof: bytes_read < len,
bytes,
})
}
#[cfg(unix)]
fn read_file_at(file: &File, bytes: &mut [u8], offset: u64) -> io::Result<usize> {
std::os::unix::fs::FileExt::read_at(file, bytes, offset)
}
#[cfg(windows)]
fn read_file_at(file: &File, bytes: &mut [u8], offset: u64) -> io::Result<usize> {
std::os::windows::fs::FileExt::seek_read(file, bytes, offset)
}
fn validate_read_block_len(len: usize) -> io::Result<()> {
if !(1..=FILE_READ_CHUNK_SIZE).contains(&len) {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("file read block length must be between 1 and {FILE_READ_CHUNK_SIZE}"),
));
}
Ok(())
}
fn unknown_handle_error(handle_id: &str) -> io::Error {
io::Error::new(
io::ErrorKind::NotFound,
format!("unknown file read handle `{handle_id}`"),
)
}
|