use std::sync::Arc;
use bytes::Bytes;
use memmap2::Mmap;
use tracing::{debug, trace, warn};
use crate::client::{OpenLocalBlockGuard, WorkerClient};
use crate::metrics::{self, name};
use crate::proto::proto::security::Capability;
use super::{AccessHint, ShortCircuitError};
struct MmapChunk {
mmap: Arc<Mmap>,
file_size: usize,
}
impl AsRef<[u8]> for MmapChunk {
#[inline]
fn as_ref(&self) -> &[u8] {
&self.mmap[..self.file_size]
}
}
pub struct LocalBlockReader {
block_id: i64,
file_size: usize,
mmap: Arc<Mmap>,
_guard: Option<OpenLocalBlockGuard>,
}
impl LocalBlockReader {
pub async fn open(
client: &WorkerClient,
block_id: i64,
block_size: i64,
capability: Option<Capability>,
hint: AccessHint,
thp: bool,
) -> Result<Self, ShortCircuitError> {
let (resp, guard) = client
.open_local_block(block_id, block_size, capability)
.await
.map_err(|e| {
metrics::counter(name::CLIENT_SC_OPENLOCAL_FAIL).inc(1);
ShortCircuitError::OpenLocalBlock(Box::new(e))
})?;
let path = resp.path.ok_or_else(|| {
metrics::counter(name::CLIENT_SC_OPENLOCAL_FAIL).inc(1);
ShortCircuitError::MissingPath
})?;
let file_size = resp.block_size.unwrap_or(block_size).max(0) as usize;
let file = std::fs::File::open(&path).map_err(|e| {
metrics::counter(name::CLIENT_SC_FILE_OPEN_FAIL).inc(1);
ShortCircuitError::FileOpen(e)
})?;
let mmap = unsafe { Mmap::map(&file) }.map_err(|e| {
metrics::counter(name::CLIENT_SC_MMAP_FAIL).inc(1);
ShortCircuitError::Mmap(e)
})?;
drop(file);
if file_size > mmap.len() {
metrics::counter(name::CLIENT_SC_MMAP_FAIL).inc(1);
warn!(
block_id = block_id,
logical = file_size,
mapped = mmap.len(),
"short-circuit: logical block size exceeds mmap length, \
falling back to gRPC (metadata/block drift?)"
);
return Err(ShortCircuitError::SizeMismatch {
logical: file_size,
mapped: mmap.len(),
});
}
apply_advice(&mmap, hint);
if thp {
apply_hugepage(&mmap);
}
metrics::counter(name::CLIENT_SC_OPEN_SUCCESS).inc(1);
metrics::gauge(name::CLIENT_SC_ACTIVE_READERS)
.set(metrics::gauge(name::CLIENT_SC_ACTIVE_READERS).get() + 1);
debug!(
block_id = block_id,
path = %path,
file_size = file_size,
mmap_len = mmap.len(),
?hint,
"LocalBlockReader opened (whole-block mmap)"
);
Ok(Self {
block_id,
file_size,
mmap: Arc::new(mmap),
_guard: Some(guard),
})
}
#[inline]
pub fn block_id(&self) -> i64 {
self.block_id
}
#[inline]
pub fn file_size(&self) -> usize {
self.file_size
}
#[inline]
fn bounds_check(&self, offset: usize, len: usize) -> Result<(), ShortCircuitError> {
let end = offset.checked_add(len);
match end {
Some(end) if end <= self.file_size => Ok(()),
_ => Err(ShortCircuitError::OutOfRange {
off: offset,
len,
file_size: self.file_size,
}),
}
}
pub fn read(&self, offset: usize, len: usize) -> Result<&[u8], ShortCircuitError> {
self.bounds_check(offset, len)?;
metrics::counter(name::CLIENT_SC_READ_CALLS).inc(1);
metrics::counter(name::CLIENT_SC_READ_BYTES).inc(len as i64);
trace!(block_id = self.block_id, offset, len, "sc read (slice)");
Ok(&self.mmap[offset..offset + len])
}
pub fn read_bytes(&self, offset: usize, len: usize) -> Result<Bytes, ShortCircuitError> {
self.bounds_check(offset, len)?;
metrics::counter(name::CLIENT_SC_READ_CALLS).inc(1);
metrics::counter(name::CLIENT_SC_READ_BYTES).inc(len as i64);
let full = Bytes::from_owner(MmapChunk {
mmap: Arc::clone(&self.mmap),
file_size: self.file_size,
});
Ok(full.slice(offset..offset + len))
}
pub fn read_to_slice(&self, offset: usize, dst: &mut [u8]) -> Result<usize, ShortCircuitError> {
let len = dst.len();
self.bounds_check(offset, len)?;
metrics::counter(name::CLIENT_SC_READ_CALLS).inc(1);
metrics::counter(name::CLIENT_SC_READ_BYTES).inc(len as i64);
dst.copy_from_slice(&self.mmap[offset..offset + len]);
Ok(len)
}
pub fn prefetch(&self, offset: usize, len: usize) -> Result<(), ShortCircuitError> {
self.bounds_check(offset, len)?;
metrics::counter(name::CLIENT_SC_PREFETCH_CALLS).inc(1);
metrics::counter(name::CLIENT_SC_PREFETCH_BYTES).inc(len as i64);
if len == 0 {
return Ok(());
}
advise_willneed(&self.mmap, offset, len);
Ok(())
}
pub fn prefetch_many(
&self,
ranges: &[(usize, usize)],
coalesce_gap: usize,
) -> Result<(), ShortCircuitError> {
metrics::counter(name::CLIENT_SC_PREFETCH_CALLS).inc(1);
if ranges.is_empty() {
return Ok(());
}
for &(off, len) in ranges {
self.bounds_check(off, len)?;
}
let merged = coalesce_ranges(ranges, coalesce_gap);
let mut total_bytes: i64 = 0;
for (_off, len) in &merged {
total_bytes += *len as i64;
}
metrics::counter(name::CLIENT_SC_PREFETCH_BYTES).inc(total_bytes);
for (off, len) in merged {
if len == 0 {
continue;
}
advise_willneed(&self.mmap, off, len);
}
Ok(())
}
#[inline]
pub fn mmap_len(&self) -> usize {
self.mmap.len()
}
#[cfg(test)]
fn open_from_path_for_test(
path: &std::path::Path,
file_size: usize,
hint: AccessHint,
) -> Result<Self, ShortCircuitError> {
let file = std::fs::File::open(path).map_err(ShortCircuitError::FileOpen)?;
let mmap = unsafe { Mmap::map(&file) }.map_err(ShortCircuitError::Mmap)?;
drop(file);
let file_size = file_size.min(mmap.len());
apply_advice(&mmap, hint);
Ok(Self {
block_id: 0,
file_size,
mmap: Arc::new(mmap),
_guard: None,
})
}
}
impl Drop for LocalBlockReader {
fn drop(&mut self) {
metrics::gauge(name::CLIENT_SC_ACTIVE_READERS)
.set((metrics::gauge(name::CLIENT_SC_ACTIVE_READERS).get() - 1).max(0));
}
}
fn coalesce_ranges(ranges: &[(usize, usize)], gap: usize) -> Vec<(usize, usize)> {
let mut sorted: Vec<(usize, usize)> = ranges.iter().copied().filter(|(_, l)| *l > 0).collect();
if sorted.is_empty() {
return Vec::new();
}
sorted.sort_by_key(|(off, _)| *off);
let mut merged: Vec<(usize, usize)> = Vec::with_capacity(sorted.len());
let (mut cur_off, mut cur_len) = sorted[0];
for &(off, len) in &sorted[1..] {
let cur_end = cur_off + cur_len;
if off <= cur_end.saturating_add(gap) {
let new_end = cur_end.max(off + len);
cur_len = new_end - cur_off;
} else {
merged.push((cur_off, cur_len));
cur_off = off;
cur_len = len;
}
}
merged.push((cur_off, cur_len));
merged
}
#[cfg(unix)]
fn apply_advice(mmap: &Mmap, hint: AccessHint) {
use memmap2::Advice;
let advice = match hint {
AccessHint::Sequential => Advice::Sequential,
AccessHint::Random => Advice::Random,
AccessHint::Default => return, };
if let Err(e) = mmap.advise(advice) {
debug!(error = %e, ?hint, "madvise(advice) failed (non-fatal)");
}
}
#[cfg(not(unix))]
fn apply_advice(_mmap: &Mmap, _hint: AccessHint) {}
#[cfg(target_os = "linux")]
fn apply_hugepage(mmap: &Mmap) {
use memmap2::Advice;
if let Err(e) = mmap.advise(Advice::HugePage) {
debug!(error = %e, "madvise(HUGEPAGE) failed (non-fatal)");
}
}
#[cfg(not(target_os = "linux"))]
fn apply_hugepage(_mmap: &Mmap) {}
#[cfg(unix)]
fn advise_willneed(mmap: &Mmap, offset: usize, len: usize) {
use memmap2::Advice;
match mmap.advise_range(Advice::WillNeed, offset, len) {
Ok(()) => {
metrics::counter(name::CLIENT_SC_PREFETCH_MADVISE).inc(1);
}
Err(e) => {
debug!(error = %e, offset, len, "madvise(WILLNEED) failed (non-fatal)");
}
}
}
#[cfg(not(unix))]
fn advise_willneed(_mmap: &Mmap, _offset: usize, _len: usize) {}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn local_block_reader_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<LocalBlockReader>();
assert_send_sync::<std::sync::Arc<LocalBlockReader>>();
}
fn write_temp(data: &[u8]) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
let uniq = format!(
"gfs_sc_test_{}_{}.bin",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
p.push(uniq);
let mut f = std::fs::File::create(&p).unwrap();
f.write_all(data).unwrap();
f.sync_all().unwrap();
p
}
fn reader_for(data: &[u8]) -> (LocalBlockReader, std::path::PathBuf) {
let path = write_temp(data);
let r = LocalBlockReader::open_from_path_for_test(&path, data.len(), AccessHint::Random)
.unwrap();
(r, path)
}
#[test]
fn read_matches_source_bytes() {
let data: Vec<u8> = (0..40_000u32).map(|i| (i % 251) as u8).collect();
let (r, path) = reader_for(&data);
for &(off, len) in &[
(0usize, 0usize),
(0, 1),
(0, data.len()),
(123, 4096),
(4095, 4098), (data.len() - 1, 1),
(data.len(), 0), ] {
let got = r.read(off, len).unwrap();
assert_eq!(got, &data[off..off + len], "off={off} len={len}");
}
std::fs::remove_file(path).ok();
}
#[test]
fn three_apis_agree() {
let data: Vec<u8> = (0..10_000u32).map(|i| (i * 7 % 256) as u8).collect();
let (r, path) = reader_for(&data);
let (off, len) = (1000usize, 2048usize);
let a = r.read(off, len).unwrap().to_vec();
let b = r.read_bytes(off, len).unwrap();
let mut c = vec![0u8; len];
let n = r.read_to_slice(off, &mut c).unwrap();
assert_eq!(n, len);
assert_eq!(a, b.as_ref());
assert_eq!(a, c);
assert_eq!(a, &data[off..off + len]);
std::fs::remove_file(path).ok();
}
#[test]
fn read_bytes_outlives_reader() {
let data: Vec<u8> = (0..5000u32).map(|i| (i % 256) as u8).collect();
let (r, path) = reader_for(&data);
let held = r.read_bytes(100, 500).unwrap();
let expected = data[100..600].to_vec();
drop(r); assert_eq!(held.as_ref(), expected.as_slice());
let clone = held.clone();
assert_eq!(clone.as_ref(), expected.as_slice());
std::fs::remove_file(path).ok();
}
#[test]
fn out_of_range_is_error() {
let data = vec![0u8; 1000];
let (r, path) = reader_for(&data);
assert!(matches!(
r.read(900, 200),
Err(ShortCircuitError::OutOfRange { .. })
));
assert!(matches!(
r.read(1001, 0),
Err(ShortCircuitError::OutOfRange { .. })
));
assert!(matches!(
r.read(usize::MAX, 1),
Err(ShortCircuitError::OutOfRange { .. })
));
assert!(r.read(1000, 0).is_ok());
let mut dst = vec![0u8; 200];
assert!(matches!(
r.read_to_slice(900, &mut dst),
Err(ShortCircuitError::OutOfRange { .. })
));
std::fs::remove_file(path).ok();
}
#[test]
fn prefetch_does_not_change_bytes() {
let data: Vec<u8> = (0..20_000u32).map(|i| (i % 256) as u8).collect();
let (r, path) = reader_for(&data);
let before = r.read(0, data.len()).unwrap().to_vec();
r.prefetch(0, 4096).unwrap();
r.prefetch(8192, 4096).unwrap();
r.prefetch_many(&[(0, 1000), (1000, 1000), (15000, 100)], 64 * 1024)
.unwrap();
let after = r.read(0, data.len()).unwrap().to_vec();
assert_eq!(before, after);
assert!(matches!(
r.prefetch(data.len() - 10, 100),
Err(ShortCircuitError::OutOfRange { .. })
));
assert!(matches!(
r.prefetch_many(&[(0, 10), (data.len(), 10)], 0),
Err(ShortCircuitError::OutOfRange { .. })
));
std::fs::remove_file(path).ok();
}
#[test]
fn logical_size_clamped_to_mapping() {
let data = vec![7u8; 100];
let path = write_temp(&data);
let r = LocalBlockReader::open_from_path_for_test(&path, 1_000_000, AccessHint::Default)
.unwrap();
assert_eq!(r.file_size(), 100, "logical size must clamp to file len");
assert!(matches!(
r.read(0, 200),
Err(ShortCircuitError::OutOfRange { .. })
));
std::fs::remove_file(path).ok();
}
#[test]
fn coalesce_merges_adjacent_and_within_gap() {
let ranges = [(0usize, 100usize), (100, 50), (200, 10), (1000, 10)];
let merged = coalesce_ranges(&ranges, 64);
assert_eq!(merged, vec![(0, 210), (1000, 10)]);
}
#[test]
fn coalesce_sorts_unordered_input() {
let ranges = [(1000, 10), (0, 10)];
let merged = coalesce_ranges(&ranges, 0);
assert_eq!(merged, vec![(0, 10), (1000, 10)]);
}
#[test]
fn coalesce_drops_zero_length() {
let ranges = [(0, 0), (10, 5), (20, 0)];
let merged = coalesce_ranges(&ranges, 0);
assert_eq!(merged, vec![(10, 5)]);
}
#[test]
fn coalesce_overlapping_ranges() {
let ranges = [(0, 100), (50, 100)];
let merged = coalesce_ranges(&ranges, 0);
assert_eq!(merged, vec![(0, 150)]);
}
#[test]
fn coalesce_empty_input() {
assert!(coalesce_ranges(&[], 64).is_empty());
}
}