use crate::config::HudiConfigs;
use bytes::Bytes;
use object_store::path::Path as ObjPath;
use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use std::io::{Error, ErrorKind, Result};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
pub const DEFAULT_STREAM_WINDOW_SIZE: u64 = 16 * 1024 * 1024;
pub const CONFIG_DFS_BUFFER_MAX_SIZE: &str = "hoodie.memory.dfs.buffer.max.size";
pub const CONFIG_HFILE_WHOLE_READ_MAX_SIZE_MB: &str = "hoodie.metadata.file.cache.max.size.mb";
pub const DEFAULT_HFILE_WHOLE_READ_MAX_SIZE_MB: u64 = 50;
pub const HFILE_WHOLE_READ_WITH_KEYS_MAX_SIZE: u64 = 512 * 1024;
pub fn hfile_whole_read_max_size(hudi_configs: &HudiConfigs) -> Result<u64> {
let mb = match hudi_configs
.as_options()
.get(CONFIG_HFILE_WHOLE_READ_MAX_SIZE_MB)
.map(|raw| raw.trim().parse::<u64>())
{
None => DEFAULT_HFILE_WHOLE_READ_MAX_SIZE_MB,
Some(Ok(mb)) => mb,
Some(Err(_)) => {
return Err(Error::other(format!(
"{CONFIG_HFILE_WHOLE_READ_MAX_SIZE_MB} must be a non-negative integer \
count of megabytes"
)));
}
};
Ok(mb.saturating_mul(1024 * 1024))
}
pub fn stream_window_size(hudi_configs: &HudiConfigs) -> Result<u64> {
let Some(raw) = hudi_configs
.as_options()
.get(CONFIG_DFS_BUFFER_MAX_SIZE)
.cloned()
else {
return Ok(DEFAULT_STREAM_WINDOW_SIZE);
};
match raw.trim().parse::<u64>() {
Ok(0) | Err(_) => Err(Error::other(format!(
"{CONFIG_DFS_BUFFER_MAX_SIZE} must be a positive integer byte count, got '{raw}'"
))),
Ok(size) => Ok(size),
}
}
async fn get_range(
object_store: &Arc<dyn ObjectStore>,
location: &ObjPath,
offset: u64,
length: u64,
) -> Result<Bytes> {
let end = offset.checked_add(length).ok_or_else(|| {
Error::other(format!(
"ranged read offset {offset} + length {length} overflows u64"
))
})?;
object_store
.get_range(location, offset..end)
.await
.map_err(|e| Error::other(format!("ranged object-store read failed: {e}")))
}
#[derive(Clone, Debug)]
pub struct LogBlockFetcher {
object_store: Arc<dyn ObjectStore>,
location: ObjPath,
reads: Arc<FetchCounts>,
}
#[derive(Debug, Default)]
pub struct FetchCounts {
calls: AtomicU64,
bytes: AtomicU64,
}
impl FetchCounts {
pub fn calls(&self) -> u64 {
self.calls.load(Ordering::Relaxed)
}
pub fn bytes(&self) -> u64 {
self.bytes.load(Ordering::Relaxed)
}
fn record(&self, calls: u64, bytes: u64) {
self.calls.fetch_add(calls, Ordering::Relaxed);
self.bytes.fetch_add(bytes, Ordering::Relaxed);
}
}
impl LogBlockFetcher {
pub fn new(object_store: Arc<dyn ObjectStore>, location: ObjPath) -> Self {
Self {
object_store,
location,
reads: Arc::new(FetchCounts::default()),
}
}
pub fn reads(&self) -> &FetchCounts {
&self.reads
}
pub async fn read_content(&self, offset: u64, length: u64) -> Result<Bytes> {
let bytes = get_range(&self.object_store, &self.location, offset, length).await?;
self.reads.record(1, bytes.len() as u64);
Ok(bytes)
}
pub fn location(&self) -> &ObjPath {
&self.location
}
pub async fn read_contents(&self, ranges: &[std::ops::Range<u64>]) -> Result<Vec<Bytes>> {
let fetched = self
.object_store
.get_ranges(&self.location, ranges)
.await
.map_err(|e| {
Error::other(format!(
"batched ranged read of {} range(s) from '{}' failed: {e}",
ranges.len(),
self.location
))
})?;
self.reads
.record(1, fetched.iter().map(|b| b.len() as u64).sum());
Ok(fetched)
}
}
#[derive(Debug)]
pub struct StorageReader {
object_store: Arc<dyn ObjectStore>,
location: ObjPath,
file_len: u64,
pos: u64,
whole: Option<Bytes>,
window: Bytes,
window_start: u64,
window_size: u64,
}
impl StorageReader {
pub async fn new(object_store: Arc<dyn ObjectStore>, object_meta: ObjectMeta) -> Result<Self> {
let get_result = object_store
.get(&object_meta.location)
.await
.map_err(|e| Error::other(format!("object-store get failed: {e}")))?;
let bytes = get_result
.bytes()
.await
.map_err(|e| Error::other(format!("object-store read failed: {e}")))?;
let file_len = bytes.len() as u64;
Ok(Self {
object_store,
location: object_meta.location,
file_len,
pos: 0,
whole: Some(bytes),
window: Bytes::new(),
window_start: 0,
window_size: DEFAULT_STREAM_WINDOW_SIZE,
})
}
pub fn new_streaming(
object_store: Arc<dyn ObjectStore>,
object_meta: ObjectMeta,
window_size: u64,
) -> Self {
Self {
object_store,
location: object_meta.location,
file_len: object_meta.size,
pos: 0,
whole: None,
window: Bytes::new(),
window_start: 0,
window_size,
}
}
pub fn block_fetcher(&self) -> LogBlockFetcher {
LogBlockFetcher::new(self.object_store.clone(), self.location.clone())
}
pub fn file_len(&self) -> u64 {
self.file_len
}
pub fn position(&self) -> u64 {
self.pos
}
pub fn seek_to(&mut self, pos: u64) {
self.pos = pos;
}
fn window_covers(&self, start: u64, end: u64) -> bool {
start >= self.window_start && end <= self.window_start + self.window.len() as u64
}
pub fn has_resident(&self, start: u64, end: u64) -> bool {
if self.whole.is_some() {
return end <= self.file_len;
}
self.window_covers(start, end)
}
async fn fill_window(&mut self) -> Result<()> {
let end = self.pos.saturating_add(self.window_size).min(self.file_len);
self.window =
get_range(&self.object_store, &self.location, self.pos, end - self.pos).await?;
self.window_start = self.pos;
Ok(())
}
pub async fn read_bytes(&mut self, len: u64) -> Result<Bytes> {
if len == 0 {
return Ok(Bytes::new());
}
let end = self.pos.checked_add(len).ok_or_else(|| {
Error::other(format!(
"read of {len} byte(s) at offset {} overflows u64",
self.pos
))
})?;
if end > self.file_len {
return Err(Error::new(
ErrorKind::UnexpectedEof,
format!(
"read of {len} byte(s) at offset {} runs past the end of '{}' ({} bytes)",
self.pos, self.location, self.file_len
),
));
}
if let Some(whole) = self.whole.as_ref() {
let bytes = whole.slice(self.pos as usize..end as usize);
self.pos = end;
return Ok(bytes);
}
if len > self.window_size {
let bytes = get_range(&self.object_store, &self.location, self.pos, len).await?;
self.pos = end;
return Ok(bytes);
}
if !self.window_covers(self.pos, end) {
self.fill_window().await?;
}
let off = (self.pos - self.window_start) as usize;
let resident = self.window.len() - off;
if resident < len as usize {
return Err(Error::new(
ErrorKind::UnexpectedEof,
format!(
"a window fetched at offset {} holds {resident} byte(s), short of the {len} \
asked for: '{}' was truncated while it was being read",
self.pos, self.location
),
));
}
let bytes = self.window.slice(off..off + len as usize);
self.pos = end;
Ok(bytes)
}
pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
if buf.is_empty() {
return Ok(());
}
let bytes = self.read_bytes(buf.len() as u64).await?;
buf.copy_from_slice(&bytes);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use object_store::local::LocalFileSystem;
use std::io::Write;
async fn make_store(
len: usize,
) -> (Arc<dyn ObjectStore>, ObjectMeta, Vec<u8>, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("data.bin");
let bytes: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
std::fs::File::create(&path)
.unwrap()
.write_all(&bytes)
.unwrap();
let store: Arc<dyn ObjectStore> = Arc::new(LocalFileSystem::new_with_prefix(&dir).unwrap());
let location = ObjPath::from("data.bin");
let meta = store.head(&location).await.unwrap();
(store, meta, bytes, dir)
}
async fn read_all(reader: &mut StorageReader, chunk: u64) -> Vec<u8> {
let mut out = Vec::new();
loop {
let want = chunk.min(reader.file_len() - reader.position());
if want == 0 {
return out;
}
out.extend_from_slice(&reader.read_bytes(want).await.unwrap());
}
}
#[tokio::test]
async fn test_streaming_read_matches_eager_across_window_refills() {
let len = (DEFAULT_STREAM_WINDOW_SIZE as usize) * 2 + 4096;
let (store, meta, expected, _dir) = make_store(len).await;
let mut eager = StorageReader::new(store.clone(), meta.clone())
.await
.unwrap();
assert_eq!(read_all(&mut eager, 64 * 1024).await, expected);
let mut streaming = StorageReader::new_streaming(store, meta, DEFAULT_STREAM_WINDOW_SIZE);
assert_eq!(read_all(&mut streaming, 64 * 1024).await, expected);
}
#[tokio::test]
async fn test_streaming_seek_backwards_and_past_end() {
let len = (DEFAULT_STREAM_WINDOW_SIZE as usize) + 1024;
let (store, meta, expected, _dir) = make_store(len).await;
let mut reader = StorageReader::new_streaming(store, meta, DEFAULT_STREAM_WINDOW_SIZE);
reader.seek_to(len as u64 - 512);
let tail = reader.read_bytes(512).await.unwrap();
assert_eq!(&tail[..], &expected[len - 512..]);
reader.seek_to(0);
let head = reader.read_bytes(512).await.unwrap();
assert_eq!(&head[..], &expected[..512]);
reader.seek_to(len as u64);
assert_eq!(
reader.read_bytes(16).await.unwrap_err().kind(),
ErrorKind::UnexpectedEof
);
reader.seek_to(len as u64 * 2);
assert_eq!(
reader.read_bytes(1).await.unwrap_err().kind(),
ErrorKind::UnexpectedEof
);
assert_eq!(reader.position(), len as u64 * 2);
}
#[test]
fn test_stream_window_size_defaults_to_hudis_own_default() {
let empty: Vec<(&str, &str)> = vec![];
assert_eq!(
stream_window_size(&HudiConfigs::new(empty)).unwrap(),
DEFAULT_STREAM_WINDOW_SIZE
);
assert_eq!(DEFAULT_STREAM_WINDOW_SIZE, 16 * 1024 * 1024);
}
#[tokio::test]
async fn test_configured_window_is_used_for_refills() {
let window = 4096u64;
let configs = HudiConfigs::new([(CONFIG_DFS_BUFFER_MAX_SIZE, window.to_string())]);
assert_eq!(stream_window_size(&configs).unwrap(), window);
let len = (window as usize) * 3 + 17;
let (store, meta, expected, _dir) = make_store(len).await;
let mut reader = StorageReader::new_streaming(store, meta, window);
assert_eq!(read_all(&mut reader, 512).await, expected);
assert!(
reader.window.len() as u64 <= window,
"a refill fetched {} bytes for a {window}-byte window",
reader.window.len()
);
}
#[tokio::test]
async fn test_read_longer_than_the_window_does_not_grow_it() {
let window = 4096u64;
let len = (window as usize) * 4;
let (store, meta, expected, _dir) = make_store(len).await;
let mut reader = StorageReader::new_streaming(store, meta, window);
reader.read_bytes(16).await.unwrap();
assert_eq!(reader.window.len() as u64, window);
reader.seek_to(1000);
let big = reader.read_bytes(window * 2).await.unwrap();
assert_eq!(&big[..], &expected[1000..1000 + (window * 2) as usize]);
assert_eq!(reader.position(), 1000 + window * 2);
assert!(
reader.window.len() as u64 <= window,
"an oversized read must not grow the resident window, now {} bytes",
reader.window.len()
);
}
#[test]
fn test_unusable_window_size_is_rejected() {
for bad in ["0", "-1", "16MB", ""] {
let configs = HudiConfigs::new([(CONFIG_DFS_BUFFER_MAX_SIZE, bad)]);
let err = stream_window_size(&configs)
.expect_err("'{bad}' is not a positive byte count and must not be accepted");
assert!(
format!("{err}").contains(CONFIG_DFS_BUFFER_MAX_SIZE),
"the error must name the config, got: {err}"
);
}
}
#[tokio::test]
async fn test_block_fetcher_reads_only_its_range() {
let len = 8192;
let (store, meta, expected, _dir) = make_store(len).await;
let reader = StorageReader::new_streaming(store, meta, DEFAULT_STREAM_WINDOW_SIZE);
let content = reader
.block_fetcher()
.read_content(1000, 256)
.await
.unwrap();
assert_eq!(&content[..], &expected[1000..1256]);
}
}