use std::io::{self, Read, Seek, Write, BufWriter};
use std::collections::HashMap;
use std::path::PathBuf;
use zip::ZipArchive;
use memmap2::{Mmap, MmapOptions};
use tempfile::NamedTempFile;
#[derive(Debug, Clone)]
pub struct MmapConfig {
pub threshold: u64,
pub max_maps: usize,
pub temp_dir: Option<PathBuf>,
pub huge_file_threshold: u64,
pub stream_chunk_size: usize,
pub enable_streaming: bool,
}
impl Default for MmapConfig {
fn default() -> Self {
Self {
threshold: 1024 * 1024, max_maps: 8,
temp_dir: None,
huge_file_threshold: 100 * 1024 * 1024, stream_chunk_size: 8 * 1024 * 1024, enable_streaming: true,
}
}
}
struct MmapEntry {
_temp_file: NamedTempFile,
mmap: Mmap,
}
#[derive(Debug, Clone, PartialEq)]
pub enum FileSizeCategory {
Small,
Large,
Huge,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ProcessingStrategy {
Standard,
MemoryMap,
Streaming,
}
#[derive(Debug, Clone)]
pub struct FileInfo {
pub name: String,
pub size: u64,
pub compressed_size: u64,
pub compression_ratio: f64,
pub category: FileSizeCategory,
pub recommended_strategy: ProcessingStrategy,
}
pub struct ZipDocumentReader<R: Read + Seek> {
pub(crate) zip: ZipArchive<R>,
mmap_config: MmapConfig,
mmap_cache: HashMap<String, MmapEntry>,
access_count: HashMap<String, u64>,
}
impl<R: Read + Seek> ZipDocumentReader<R> {
pub fn new(r: R) -> io::Result<Self> {
Self::with_mmap_config(r, MmapConfig::default())
}
pub fn with_mmap_config(
r: R,
config: MmapConfig,
) -> io::Result<Self> {
Ok(Self {
zip: ZipArchive::new(r)?,
mmap_config: config,
mmap_cache: HashMap::new(),
access_count: HashMap::new(),
})
}
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(self), fields(
crate_name = "file",
file_name = %name
)))]
pub fn read_all(
&mut self,
name: &str,
) -> io::Result<Vec<u8>> {
self.read_smart(name)
}
pub fn read_smart(
&mut self,
name: &str,
) -> io::Result<Vec<u8>> {
let file_info = self.get_file_info(name)?;
match file_info.recommended_strategy {
ProcessingStrategy::Standard => {
*self.access_count.entry(name.to_string()).or_insert(0) += 1;
self.read_standard(name)
},
ProcessingStrategy::MemoryMap => {
match self.read_mmap(name) {
Ok(data) => Ok(data.to_vec()),
Err(_) => {
*self
.access_count
.entry(name.to_string())
.or_insert(0) += 1;
self.read_standard(name)
},
}
},
ProcessingStrategy::Streaming => {
if self.mmap_config.enable_streaming {
*self.access_count.entry(name.to_string()).or_insert(0) +=
1;
self.read_huge_file_streaming(name)
} else {
match self.read_mmap(name) {
Ok(data) => Ok(data.to_vec()),
Err(_) => {
*self
.access_count
.entry(name.to_string())
.or_insert(0) += 1;
self.read_standard(name)
},
}
}
},
}
}
pub fn read_mmap(
&mut self,
name: &str,
) -> io::Result<&[u8]> {
if self.mmap_cache.contains_key(name) {
*self.access_count.entry(name.to_string()).or_insert(0) += 1;
return Ok(&self.mmap_cache[name].mmap[..]);
}
if self.mmap_cache.len() >= self.mmap_config.max_maps {
self.evict_least_used();
}
self.create_mmap_entry(name)?;
self.access_count.insert(name.to_string(), 1);
Ok(&self.mmap_cache[name].mmap[..])
}
pub fn read_standard(
&mut self,
name: &str,
) -> io::Result<Vec<u8>> {
let mut f = self.zip.by_name(name)?;
let mut buf = Vec::with_capacity(f.size() as usize);
std::io::copy(&mut f, &mut buf)?;
Ok(buf)
}
pub fn read_plugin_state(
&mut self,
plugin_name: &str,
) -> io::Result<Option<Vec<u8>>> {
let plugin_file_path = format!("plugins/{plugin_name}");
match self.read_all(&plugin_file_path) {
Ok(data) => Ok(Some(data)),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e),
}
}
pub fn read_all_plugin_states(
&mut self
) -> io::Result<std::collections::HashMap<String, Vec<u8>>> {
let mut plugin_states = std::collections::HashMap::new();
let mut plugin_files = Vec::new();
for i in 0..self.zip.len() {
let file = self.zip.by_index(i)?;
let file_name = file.name().to_string();
if file_name.starts_with("plugins/") && !file_name.ends_with('/') {
let plugin_name =
file_name.strip_prefix("plugins/").unwrap().to_string();
plugin_files.push((plugin_name, file_name));
}
}
for (plugin_name, file_name) in plugin_files {
let data = self.read_all(&file_name)?;
plugin_states.insert(plugin_name, data);
}
Ok(plugin_states)
}
pub fn list_plugins(&mut self) -> io::Result<Vec<String>> {
let mut plugins = Vec::new();
for i in 0..self.zip.len() {
let file = self.zip.by_index(i)?;
let file_name = file.name();
if file_name.starts_with("plugins/") && !file_name.ends_with('/') {
let plugin_name =
file_name.strip_prefix("plugins/").unwrap().to_string();
plugins.push(plugin_name);
}
}
Ok(plugins)
}
pub fn has_plugin_state(
&mut self,
plugin_name: &str,
) -> bool {
let plugin_file_path = format!("plugins/{plugin_name}");
self.zip.by_name(&plugin_file_path).is_ok()
}
fn create_mmap_entry(
&mut self,
name: &str,
) -> io::Result<()> {
let mut temp_file =
if let Some(ref temp_dir) = self.mmap_config.temp_dir {
NamedTempFile::new_in(temp_dir)?
} else {
NamedTempFile::new()?
};
{
let mut zip_file = self.zip.by_name(name)?;
let mut writer = BufWriter::new(&mut temp_file);
std::io::copy(&mut zip_file, &mut writer)?;
writer.flush()?;
}
temp_file.as_file().sync_all()?;
let mmap = unsafe { MmapOptions::new().map(temp_file.as_file())? };
self.mmap_cache.insert(
name.to_string(),
MmapEntry { _temp_file: temp_file, mmap },
);
Ok(())
}
fn evict_least_used(&mut self) {
if let Some((lru_name, _)) = self
.access_count
.iter()
.min_by_key(|(_, count)| **count)
.map(|(name, count)| (name.clone(), *count))
{
self.mmap_cache.remove(&lru_name);
self.access_count.remove(&lru_name);
}
}
pub fn mmap_config(&self) -> &MmapConfig {
&self.mmap_config
}
pub fn mmap_stats(&self) -> MmapStats {
let total_size: u64 =
self.mmap_cache.values().map(|entry| entry.mmap.len() as u64).sum();
MmapStats {
cached_entries: self.mmap_cache.len(),
total_cached_size: total_size,
max_entries: self.mmap_config.max_maps,
threshold_bytes: self.mmap_config.threshold,
}
}
pub fn clear_mmap_cache(&mut self) {
self.mmap_cache.clear();
self.access_count.clear();
}
pub fn get_file_size(
&mut self,
name: &str,
) -> io::Result<u64> {
let f = self.zip.by_name(name)?;
Ok(f.size())
}
pub fn get_compressed_size(
&mut self,
name: &str,
) -> io::Result<u64> {
let f = self.zip.by_name(name)?;
Ok(f.compressed_size())
}
pub fn classify_file_size(
&mut self,
name: &str,
) -> io::Result<FileSizeCategory> {
let file_size = self.get_file_size(name)?;
if file_size >= self.mmap_config.huge_file_threshold {
Ok(FileSizeCategory::Huge)
} else if file_size >= self.mmap_config.threshold {
Ok(FileSizeCategory::Large)
} else {
Ok(FileSizeCategory::Small)
}
}
pub fn get_file_info(
&mut self,
name: &str,
) -> io::Result<FileInfo> {
let (size, compressed_size) = {
let f = self.zip.by_name(name)?;
(f.size(), f.compressed_size())
};
let category = if size >= self.mmap_config.huge_file_threshold {
FileSizeCategory::Huge
} else if size >= self.mmap_config.threshold {
FileSizeCategory::Large
} else {
FileSizeCategory::Small
};
let recommended_strategy = self.recommend_processing_strategy(size);
Ok(FileInfo {
name: name.to_string(),
size,
compressed_size,
compression_ratio: if size > 0 {
compressed_size as f64 / size as f64
} else {
1.0
},
category,
recommended_strategy,
})
}
pub fn recommend_processing_strategy(
&self,
file_size: u64,
) -> ProcessingStrategy {
if file_size >= self.mmap_config.huge_file_threshold
&& self.mmap_config.enable_streaming
{
ProcessingStrategy::Streaming
} else if file_size >= self.mmap_config.threshold {
ProcessingStrategy::MemoryMap
} else {
ProcessingStrategy::Standard
}
}
pub fn preheat_mmap(
&mut self,
names: &[&str],
) -> io::Result<()> {
for &name in names {
if !self.mmap_cache.contains_key(name) {
let file_size = self.get_file_size(name)?;
if file_size >= self.mmap_config.threshold {
self.create_mmap_entry(name)?;
}
}
}
Ok(())
}
fn read_huge_file_streaming(
&mut self,
name: &str,
) -> io::Result<Vec<u8>> {
let mut file = self.zip.by_name(name)?;
let total_size = file.size() as usize;
let mut result = Vec::with_capacity(total_size);
let chunk_size = self.mmap_config.stream_chunk_size;
let mut buffer = vec![0u8; chunk_size];
loop {
let bytes_read = file.read(&mut buffer)?;
if bytes_read == 0 {
break;
}
result.extend_from_slice(&buffer[..bytes_read]);
}
Ok(result)
}
pub fn create_stream_reader(
&mut self,
name: &str,
) -> io::Result<ZipStreamReader> {
let mut file = self.zip.by_name(name)?;
let total_size = file.size();
let chunk_size = self.mmap_config.stream_chunk_size;
let mut chunks = Vec::new();
let mut buffer = vec![0u8; chunk_size];
loop {
let bytes_read = file.read(&mut buffer)?;
if bytes_read == 0 {
break;
}
chunks.push(buffer[..bytes_read].to_vec());
}
Ok(ZipStreamReader {
chunks: chunks.into_iter(),
total_size,
current_pos: 0,
})
}
pub fn process_smart<F>(
&mut self,
name: &str,
mut processor: F,
) -> io::Result<()>
where
F: FnMut(&[u8]) -> io::Result<()>,
{
let file_info = self.get_file_info(name)?;
match file_info.recommended_strategy {
ProcessingStrategy::Standard => {
*self.access_count.entry(name.to_string()).or_insert(0) += 1;
let data = self.read_standard(name)?;
processor(&data)
},
ProcessingStrategy::MemoryMap => {
match self.read_mmap(name) {
Ok(data) => {
processor(data)
},
Err(_) => {
*self
.access_count
.entry(name.to_string())
.or_insert(0) += 1;
let data = self.read_standard(name)?;
processor(&data)
},
}
},
ProcessingStrategy::Streaming => {
if self.mmap_config.enable_streaming {
*self.access_count.entry(name.to_string()).or_insert(0) +=
1;
self.process_huge_file(name, processor)
} else {
match self.read_mmap(name) {
Ok(data) => processor(data),
Err(_) => {
*self
.access_count
.entry(name.to_string())
.or_insert(0) += 1;
let data = self.read_standard(name)?;
processor(&data)
},
}
}
},
}
}
pub fn process_files_smart<F>(
&mut self,
file_names: &[&str],
mut processor: F,
) -> io::Result<()>
where
F: FnMut(&str, &[u8]) -> io::Result<()>,
{
for &name in file_names {
let file_info = self.get_file_info(name)?;
match file_info.recommended_strategy {
ProcessingStrategy::Standard
| ProcessingStrategy::MemoryMap => {
let data = self.read_smart(name)?;
processor(name, &data)?;
},
ProcessingStrategy::Streaming => {
let mut accumulated_data = Vec::new();
self.process_smart(name, |chunk| {
accumulated_data.extend_from_slice(chunk);
Ok(())
})?;
processor(name, &accumulated_data)?;
},
}
}
Ok(())
}
pub fn process_huge_file<F>(
&mut self,
name: &str,
mut processor: F,
) -> io::Result<()>
where
F: FnMut(&[u8]) -> io::Result<()>,
{
let mut file = self.zip.by_name(name)?;
let chunk_size = self.mmap_config.stream_chunk_size;
let mut buffer = vec![0u8; chunk_size];
loop {
let bytes_read = file.read(&mut buffer)?;
if bytes_read == 0 {
break;
}
processor(&buffer[..bytes_read])?;
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct MmapStats {
pub cached_entries: usize,
pub total_cached_size: u64,
pub max_entries: usize,
pub threshold_bytes: u64,
}
impl std::fmt::Display for MmapStats {
fn fmt(
&self,
f: &mut std::fmt::Formatter<'_>,
) -> std::fmt::Result {
write!(
f,
"mmap 缓存: {}/{} 条目, {:.2} MB 总大小, 阈值 {:.2} MB",
self.cached_entries,
self.max_entries,
self.total_cached_size as f64 / (1024.0 * 1024.0),
self.threshold_bytes as f64 / (1024.0 * 1024.0)
)
}
}
pub struct ZipStreamReader {
chunks: std::vec::IntoIter<Vec<u8>>,
total_size: u64,
current_pos: u64,
}
impl ZipStreamReader {
pub fn read_chunk(&mut self) -> io::Result<Option<Vec<u8>>> {
if let Some(chunk) = self.chunks.next() {
self.current_pos += chunk.len() as u64;
Ok(Some(chunk))
} else {
Ok(None)
}
}
pub fn total_size(&self) -> u64 {
self.total_size
}
pub fn position(&self) -> u64 {
self.current_pos
}
pub fn is_finished(&self) -> bool {
self.current_pos >= self.total_size
}
pub fn reset(&mut self) {
self.current_pos = 0;
}
pub fn process_chunks<F>(
&mut self,
mut processor: F,
) -> io::Result<()>
where
F: FnMut(&[u8]) -> io::Result<()>,
{
while let Some(chunk) = self.read_chunk()? {
processor(&chunk)?;
}
Ok(())
}
pub fn read_all_streaming(&mut self) -> io::Result<Vec<u8>> {
let mut result = Vec::with_capacity(self.total_size as usize);
while let Some(chunk) = self.read_chunk()? {
result.extend_from_slice(&chunk);
}
Ok(result)
}
pub fn compute_hash<H>(
&mut self,
mut hasher: H,
) -> io::Result<()>
where
H: FnMut(&[u8]),
{
self.reset();
while let Some(chunk) = self.read_chunk()? {
hasher(&chunk);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
use crate::zipdoc::ZipDocumentWriter;
#[test]
fn test_mmap_integration_basic() -> io::Result<()> {
let mut zip_data = Vec::new();
{
let cursor = Cursor::new(&mut zip_data);
let mut writer = ZipDocumentWriter::new(cursor)?;
writer.add_stored("small.txt", b"small content")?;
let large_content = vec![42u8; 2 * 1024 * 1024]; writer.add_stored("large.bin", &large_content)?;
writer.finalize()?;
}
let cursor = Cursor::new(zip_data);
let mut reader = ZipDocumentReader::new(cursor)?;
let small_data = reader.read_all("small.txt")?;
assert_eq!(small_data, b"small content");
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 0);
let large_data = reader.read_all("large.bin")?;
assert_eq!(large_data.len(), 2 * 1024 * 1024);
assert!(large_data.iter().all(|&b| b == 42));
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 1);
assert_eq!(stats.total_cached_size, 2 * 1024 * 1024);
let large_data2 = reader.read_all("large.bin")?;
assert_eq!(large_data2, large_data);
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 1);
Ok(())
}
#[test]
fn test_mmap_zero_copy_read() -> io::Result<()> {
let mut zip_data = Vec::new();
{
let cursor = Cursor::new(&mut zip_data);
let mut writer = ZipDocumentWriter::new(cursor)?;
let test_data = vec![123u8; 3 * 1024 * 1024]; writer.add_stored("test.bin", &test_data)?;
writer.finalize()?;
}
let cursor = Cursor::new(zip_data);
let mut reader = ZipDocumentReader::new(cursor)?;
let mmap_data = reader.read_mmap("test.bin")?;
assert_eq!(mmap_data.len(), 3 * 1024 * 1024);
assert!(mmap_data.iter().all(|&b| b == 123));
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 1);
assert_eq!(stats.total_cached_size, 3 * 1024 * 1024);
Ok(())
}
#[test]
fn test_mmap_cache_eviction() -> io::Result<()> {
let config = MmapConfig {
threshold: 1024, max_maps: 2, temp_dir: None,
huge_file_threshold: 100 * 1024 * 1024,
stream_chunk_size: 8 * 1024 * 1024,
enable_streaming: true,
};
let mut zip_data = Vec::new();
{
let cursor = Cursor::new(&mut zip_data);
let mut writer = ZipDocumentWriter::new(cursor)?;
for i in 1..=3 {
let content = vec![i as u8; 2048]; writer.add_stored(&format!("file{i}.bin"), &content)?;
}
writer.finalize()?;
}
let cursor = Cursor::new(zip_data);
let mut reader = ZipDocumentReader::with_mmap_config(cursor, config)?;
let _data1 = reader.read_all("file1.bin")?;
let _data2 = reader.read_all("file2.bin")?;
assert_eq!(reader.mmap_stats().cached_entries, 2);
let _data3 = reader.read_all("file3.bin")?;
assert_eq!(reader.mmap_stats().cached_entries, 2);
Ok(())
}
#[test]
fn test_mmap_config_threshold() -> io::Result<()> {
let config = MmapConfig {
threshold: 5 * 1024 * 1024, max_maps: 8,
temp_dir: None,
huge_file_threshold: 100 * 1024 * 1024,
stream_chunk_size: 8 * 1024 * 1024,
enable_streaming: true,
};
let mut zip_data = Vec::new();
{
let cursor = Cursor::new(&mut zip_data);
let mut writer = ZipDocumentWriter::new(cursor)?;
let small_content = vec![1u8; 1024 * 1024]; writer.add_stored("small.bin", &small_content)?;
let large_content = vec![2u8; 6 * 1024 * 1024]; writer.add_stored("large.bin", &large_content)?;
writer.finalize()?;
}
let cursor = Cursor::new(zip_data);
let mut reader = ZipDocumentReader::with_mmap_config(cursor, config)?;
let _small_data = reader.read_all("small.bin")?;
assert_eq!(reader.mmap_stats().cached_entries, 0);
let _large_data = reader.read_all("large.bin")?;
assert_eq!(reader.mmap_stats().cached_entries, 1);
Ok(())
}
#[test]
fn test_mmap_preheat() -> io::Result<()> {
let mut zip_data = Vec::new();
{
let cursor = Cursor::new(&mut zip_data);
let mut writer = ZipDocumentWriter::new(cursor)?;
for i in 1..=3 {
let content = vec![i as u8; 2 * 1024 * 1024]; writer.add_stored(&format!("data{i}.bin"), &content)?;
}
writer.finalize()?;
}
let cursor = Cursor::new(zip_data);
let mut reader = ZipDocumentReader::new(cursor)?;
reader.preheat_mmap(&["data1.bin", "data2.bin"])?;
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 2);
let _data1 = reader.read_mmap("data1.bin")?;
let _data2 = reader.read_mmap("data2.bin")?;
let stats = reader.mmap_stats();
assert_eq!(stats.cached_entries, 2);
Ok(())
}
#[test]
fn test_mmap_stats_display() {
let stats = MmapStats {
cached_entries: 3,
total_cached_size: 5 * 1024 * 1024, max_entries: 8,
threshold_bytes: 1024 * 1024, };
let display = format!("{stats}");
assert!(display.contains("3/8 条目"));
assert!(display.contains("5.00 MB"));
assert!(display.contains("1.00 MB"));
}
}