use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use super::FileProcessor;
use super::strategies::Strategy;
use crate::errors::TaleError;
use crate::memory_budget::{MemoryAllocation, MemoryBudget, MemoryPressure};
use crate::metrics::*;
use crate::readers::strategies::ChunkConfig;
use crate::readers::{IsStrategy, StaticStrategy};
#[derive(Debug)]
pub struct FileChunk {
pub data: Vec<u8>,
pub start_offset: u64,
pub end_offset: u64,
pub starts_at_line_boundary: bool,
pub ends_at_line_boundary: bool,
}
impl FileChunk {
pub fn new(data: Vec<u8>, start_offset: u64, end_offset: u64) -> Self {
let starts_at_line_boundary = start_offset == 0 || data.first() != Some(&b'\n');
let ends_at_line_boundary = data.last() == Some(&b'\n');
Self {
data,
start_offset,
end_offset,
starts_at_line_boundary,
ends_at_line_boundary,
}
}
pub fn size(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn lines(&self) -> impl Iterator<Item = &str> {
let data_str = std::str::from_utf8(&self.data).unwrap_or("");
data_str.lines()
}
pub fn find_last_line_boundary(&self) -> Option<usize> {
self.data.iter().rposition(|&b| b == b'\n')
}
pub fn split_at_last_line(&mut self) -> Option<Vec<u8>> {
if let Some(boundary) = self.find_last_line_boundary() {
let remainder = self.data.split_off(boundary + 1);
self.end_offset = self.start_offset + self.data.len() as u64;
self.ends_at_line_boundary = true;
Some(remainder)
} else {
None
}
}
}
#[derive(Debug)]
pub struct ChunkedFileReader {
file: File,
file_size: u64,
current_position: u64,
_path: PathBuf,
pending_data: Vec<u8>,
strategy: Strategy,
metrics: ChunkMetrics,
memory_budget: Option<MemoryBudget>,
current_allocation: Option<MemoryAllocation>,
reader_id: String,
}
impl ChunkedFileReader {
pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, TaleError> {
let file_size = std::fs::metadata(&path)?.len();
#[cfg(not(test))]
let strategy = Strategy::from_config(crate::config::config(), Some(file_size));
#[cfg(test)]
let strategy = Strategy::from_config(&crate::config::config(), Some(file_size));
let path = path.as_ref().to_path_buf();
let reader_id = format!(
"chunked_reader_{}",
path.file_name().and_then(|n| n.to_str()).unwrap_or("unknown")
);
let mut file = File::open(&path)?;
file.seek(SeekFrom::End(0))?;
file.seek(SeekFrom::Start(0))?;
let memory_budget = if let Some(max_memory) = crate::config::config().max_memory {
Some(MemoryBudget::new(max_memory)?)
} else {
MemoryBudget::from_system_memory(10.0).ok()
};
Ok(Self {
file,
file_size,
current_position: 0,
_path: path,
pending_data: Vec::new(),
strategy,
metrics: ChunkMetrics::new(),
memory_budget,
current_allocation: None,
reader_id,
})
}
pub fn with_strategy<P: AsRef<Path>>(path: P, strategy: Strategy) -> Result<Self, TaleError> {
let mut reader = Self::new(path)?;
reader.strategy = strategy;
Ok(reader)
}
pub fn with_memory_budget<P: AsRef<Path>>(path: P, memory_budget: MemoryBudget) -> Result<Self, TaleError> {
let mut reader = Self::new(path)?;
reader.memory_budget = Some(memory_budget);
Ok(reader)
}
pub fn static_optimal<P: AsRef<Path>>(path: P) -> Result<Self, TaleError> {
let mut reader = Self::new(&path)?;
let file_size = reader.file_size;
let strategy = StaticStrategy::optimal_for_file(file_size);
reader.strategy = Strategy::Static(strategy);
Ok(reader)
}
pub fn new_with_config<P: AsRef<Path>>(path: P, config: ChunkConfig) -> Result<Self, TaleError> {
let mut reader = Self::new(&path)?;
let strategy = StaticStrategy::with_config(config);
reader.strategy = Strategy::Static(strategy);
Ok(reader)
}
pub fn with_optimal_config<P: AsRef<Path>>(path: P) -> Result<Self, TaleError> {
Self::static_optimal(path)
}
pub fn file_size(&self) -> u64 {
self.file_size
}
pub fn position(&self) -> u64 {
self.current_position
}
pub fn is_at_end(&self) -> bool {
self.current_position >= self.file_size
}
pub fn read_chunk(&mut self) -> Result<Option<FileChunk>, TaleError> {
if self.is_at_end() && self.pending_data.is_empty() {
return Ok(None);
}
if self.metrics.chunks_seen % crate::defaults::processing::ADAPTATION_INTERVAL == 0
&& self.strategy.should_adapt(&self.metrics)
{
let current_size = self.strategy.initial_chunk_size();
self.strategy.adapt_size(&self.metrics, current_size);
}
let mut chunk_size = self.strategy.initial_chunk_size();
if let Some(ref budget) = self.memory_budget {
if let Ok(pressure) = budget.current_pressure() {
let factor = pressure.chunk_size_factor();
chunk_size = (chunk_size as f64 * factor) as usize;
chunk_size = chunk_size.max(4096);
if matches!(pressure, MemoryPressure::Critical) {
eprintln!(
"⚠️ Critical memory pressure - reducing chunk size to {} bytes",
chunk_size
);
}
}
let total_allocation_needed = chunk_size + self.pending_data.len();
self.current_allocation = None;
match budget.try_allocate(total_allocation_needed, &self.reader_id) {
Ok(Some(allocation)) => {
self.current_allocation = Some(allocation);
}
Ok(None) => {
let emergency_size = chunk_size / 4; if emergency_size >= 1024 {
chunk_size = emergency_size;
let emergency_allocation =
budget.try_allocate(emergency_size + self.pending_data.len(), &self.reader_id)?;
if let Some(allocation) = emergency_allocation {
self.current_allocation = Some(allocation);
eprintln!("🆘 Emergency memory allocation - using {} byte chunks", chunk_size);
} else {
return Err(TaleError::MemoryError(
"Cannot allocate memory even for emergency chunk size".to_string(),
));
}
} else {
return Err(TaleError::MemoryError(
"Out of memory - chunk size would be too small".to_string(),
));
}
}
Err(e) => return Err(e),
}
}
let pending_len = self.pending_data.len();
let mut buffer = vec![0u8; chunk_size];
let bytes_read = if self.is_at_end() {
0
} else {
let start = std::time::Instant::now();
let read = self.file.read(&mut buffer).map_err(TaleError::from)?;
let read_duration = start.elapsed();
if read > 0 {
let line_count = buffer[..read].iter().filter(|&&b| b == b'\n').count();
self.metrics.record_chunk_processing(read, read_duration, line_count);
}
read
};
if bytes_read == 0 && self.pending_data.is_empty() {
return Ok(None);
}
buffer.truncate(bytes_read);
if !self.pending_data.is_empty() {
let mut combined = std::mem::take(&mut self.pending_data);
combined.extend_from_slice(&buffer);
buffer = combined;
}
let start_offset = self.current_position - pending_len as u64;
self.current_position += bytes_read as u64;
let mut chunk = FileChunk::new(buffer, start_offset, self.current_position);
if !chunk.ends_at_line_boundary
&& !self.is_at_end()
&& let Some(remainder) = chunk.split_at_last_line()
{
self.pending_data = remainder;
}
Ok(Some(chunk))
}
pub fn seek(&mut self, pos: SeekFrom) -> Result<u64, TaleError> {
let new_pos = self.file.seek(pos).map_err(TaleError::from)?;
self.current_position = new_pos;
self.pending_data.clear();
Ok(new_pos)
}
pub fn reset(&mut self) -> Result<(), TaleError> {
self.seek(SeekFrom::Start(0))?;
Ok(())
}
pub fn memory_pressure(&self) -> Option<Result<MemoryPressure, TaleError>> {
self.memory_budget.as_ref().map(|budget| budget.current_pressure())
}
pub fn memory_stats(&self) -> Option<Result<crate::memory_budget::MemoryBudgetStats, TaleError>> {
self.memory_budget.as_ref().map(|budget| budget.usage_stats())
}
pub fn has_memory_budget(&self) -> bool {
self.memory_budget.is_some()
}
}
impl FileProcessor for ChunkedFileReader {
fn process_lines<F>(&mut self, mut line_processor: F) -> Result<(), TaleError>
where
F: FnMut(&str) -> Result<(), TaleError>,
{
while let Some(chunk) = self.read_chunk()? {
for line in chunk.lines() {
line_processor(line)?;
}
}
Ok(())
}
fn skip_lines(&mut self, count: u64) -> Result<(), TaleError> {
let mut lines_skipped = 0u64;
while lines_skipped < count {
if let Some(chunk) = self.read_chunk()? {
let mut lines_in_chunk = 0u64;
let mut last_newline_pos = None;
for (i, &byte) in chunk.data.iter().enumerate() {
if byte == b'\n' {
lines_in_chunk += 1;
last_newline_pos = Some(i);
if lines_skipped + lines_in_chunk == count {
let position_after_newline = i + 1;
if position_after_newline < chunk.data.len() {
self.pending_data = chunk.data[position_after_newline..].to_vec();
}
return Ok(());
}
}
}
lines_skipped += lines_in_chunk;
if !chunk.ends_at_line_boundary && lines_skipped < count {
if let Some(last_nl) = last_newline_pos {
let after_last_newline = last_nl + 1;
if after_last_newline < chunk.data.len() {
self.pending_data = chunk.data[after_last_newline..].to_vec();
}
} else {
self.pending_data = chunk.data;
}
}
} else {
break;
}
}
Ok(())
}
fn file_size(&self) -> u64 {
self.file_size
}
fn seek(&mut self, pos: SeekFrom) -> Result<u64, TaleError> {
self.seek(pos)
}
fn position(&self) -> u64 {
self.current_position
}
}