diff --git a/Cargo.lock b/Cargo.lock index 3fdd0295c..38e72797d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4664,6 +4664,7 @@ dependencies = [ "uuid", "vergen", "winapi", + "windows-sys 0.59.0", "zeroize", ] diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 30484d1fc..9b5bd4690 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -77,10 +77,11 @@ fancy-regex = "0.17.0" mime_guess = "2" crc32fast = "1.5.0" -[target.'cfg(target_os = "macos")'.dependencies] +[target.'cfg(unix)'.dependencies] libc = "0.2" [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["psapi", "minwindef", "processthreadsapi"] } +windows-sys = { version = "0.59", features = ["Win32_Foundation", "Win32_System_Threading"] } #[target.'cfg(target_os = "linux")'.dependencies] #mimalloc = "0.1" #[target.'cfg(not(target_env = "msvc"))'.dependencies] diff --git a/backend/src/processing/processor/playlist.rs b/backend/src/processing/processor/playlist.rs index 94e307321..6a4af8462 100644 --- a/backend/src/processing/processor/playlist.rs +++ b/backend/src/processing/processor/playlist.rs @@ -1,5 +1,5 @@ use crate::model::{AppConfig, ConfigFavourites, ConfigInput, ConfigRename, ReverseProxyDisabledHeaderConfig, TVGuide}; -use crate::utils::m3u; +use crate::utils::{log_memory_snapshot, m3u}; use crate::utils::xtream; use crate::utils::{epg, StepMeasureCallback}; use std::collections::{HashMap, HashSet}; @@ -11,7 +11,6 @@ use tokio::task::JoinSet; use crate::api::model::{EventManager, EventMessage, PlaylistStorageState, UpdateGuard}; use crate::messaging::send_message; use crate::model::messaging::MessageContent; -use crate::model::Epg; use crate::model::FetchedPlaylist; @@ -370,12 +369,13 @@ async fn playlist_download_from_input(client: &reqwest::Client, app_config: &Arc } async fn process_source(source_idx: usize, ctx: &PlaylistProcessingContext) -> (Vec, Vec, Vec) { + log_memory_snapshot(format!("source[{source_idx}] start").as_str()); let sources = ctx.config.sources.load(); let mut errors = vec![]; let mut input_stats = HashMap::, InputStats>::new(); let mut target_stats = Vec::::new(); if let Some(source) = sources.get_source_at(source_idx) { - let mut source_playlists = Vec::with_capacity(128); + let mut source_playlists = Vec::with_capacity(source.inputs.len()); let broadcast_step = create_broadcast_callback(ctx.event_manager.as_ref()); // Download the sources let mut source_downloaded = false; @@ -386,6 +386,7 @@ async fn process_source(source_idx: usize, ctx: &PlaylistProcessingContext) -> ( }; if is_input_enabled(input, &ctx.user_targets) { source_downloaded = true; + log_memory_snapshot(format!("source[{source_idx}] input '{}' before_download", input.name).as_str()); let start_time = Instant::now(); // Download the playlist for input @@ -401,12 +402,14 @@ async fn process_source(source_idx: usize, ctx: &PlaylistProcessingContext) -> ( } (playlist, download_err) }; + log_memory_snapshot(format!("source[{source_idx}] input '{}' after_download", input.name).as_str()); let tvguide = if input.input_type == InputType::Library { None } else { download_input_epg(ctx, input, &mut error_list).await }; + log_memory_snapshot(format!("source[{source_idx}] input '{}' after_epg_download", input.name).as_str()); errors.append(&mut error_list); let group_count = playlist_groups.get_group_count(); @@ -424,6 +427,7 @@ async fn process_source(source_idx: usize, ctx: &PlaylistProcessingContext) -> ( epg: tvguide, } ); + log_memory_snapshot(format!("source[{source_idx}] input '{}' after_source_push", input.name).as_str()); } let elapsed = start_time.elapsed().as_secs(); input_stats.insert(input_name.clone(), create_input_stat(group_count, channel_count, errors.len(), @@ -436,23 +440,34 @@ async fn process_source(source_idx: usize, ctx: &PlaylistProcessingContext) -> ( errors.push(notify_err!("Source at index {source_idx} is empty: {}", source.inputs.iter().map(Clone::clone).collect::>>().join(", "))); } else { debug_if_enabled!("Source has {} groups", source_playlists.iter_mut().map(FetchedPlaylist::get_channel_count).sum::()); - for target in &source.targets { - if is_target_enabled(target, &ctx.user_targets) { - match process_playlist_for_target(ctx, &mut source_playlists, target, - &mut input_stats, &mut errors).await { - Ok(()) => { - target_stats.push(TargetStats::success(&target.name)); - } - Err(mut err) => { - target_stats.push(TargetStats::failure(&target.name)); - errors.append(&mut err); - } + let enabled_targets: Vec<_> = source.targets + .iter() + .filter(|target| is_target_enabled(target, &ctx.user_targets)) + .collect(); + + for (idx, target) in enabled_targets.iter().enumerate() { + let consume_input_source = idx + 1 == enabled_targets.len(); + debug!( + "Processing target '{}' (use_memory_cache={}, consume_input_source={})", + target.name, target.use_memory_cache, consume_input_source + ); + log_memory_snapshot(format!("source[{source_idx}] target '{}' before_process", target.name).as_str()); + match process_playlist_for_target(ctx, &mut source_playlists, target, + &mut input_stats, &mut errors, consume_input_source).await { + Ok(()) => { + target_stats.push(TargetStats::success(&target.name)); + } + Err(mut err) => { + target_stats.push(TargetStats::failure(&target.name)); + errors.append(&mut err); } } + log_memory_snapshot(format!("source[{source_idx}] target '{}' after_process", target.name).as_str()); } } } } + log_memory_snapshot(format!("source[{source_idx}] end").as_str()); (input_stats.into_values().collect(), target_stats, errors) } @@ -655,11 +670,21 @@ fn get_processing_pipe(target: &ConfigTarget) -> ProcessingPipe { } } -fn execute_pipe<'a>(target: &ConfigTarget, pipe: &ProcessingPipe, fpl: &FetchedPlaylist<'a>, - duplicates: &mut HashSet) -> FetchedPlaylist<'a> { +fn execute_pipe<'a>(target: &ConfigTarget, pipe: &ProcessingPipe, fpl: &mut FetchedPlaylist<'a>, + duplicates: &mut HashSet, consume_source: bool) -> FetchedPlaylist<'a> { + let source = if consume_source { + if fpl.is_memory() { + MemoryPlaylistSource::new(fpl.source.take_groups()).boxed() + } else { + std::mem::replace(&mut fpl.source, MemoryPlaylistSource::default().boxed()) + } + } else { + fpl.clone_source() + }; + let mut new_fpl = FetchedPlaylist { input: fpl.input, - source: fpl.clone_source(), + source, epg: fpl.epg.clone(), }; if target.options.as_ref().is_some_and(|opt| opt.remove_duplicates) { @@ -709,11 +734,14 @@ async fn process_playlist_for_target(ctx: &PlaylistProcessingContext, target: &ConfigTarget, stats: &mut HashMap, InputStats>, errors: &mut Vec, + consume_input_source: bool, ) -> Result<(), Vec> { debug_if_enabled!("Processing order is {}", &target.processing_order); + log_memory_snapshot(format!("target '{}' start", target.name).as_str()); let mut duplicates: HashSet = HashSet::new(); - let mut processed_fetched_playlists: Vec = vec![]; + let mut new_epg = vec![]; + let mut new_playlist: Vec = vec![]; debug!("Executing processing pipes"); let broadcast_step = create_broadcast_callback(ctx.event_manager.as_ref()); @@ -721,11 +749,15 @@ async fn process_playlist_for_target(ctx: &PlaylistProcessingContext, let pipe = get_processing_pipe(target); let mut step = StepMeasure::new(&target.name, broadcast_step); for provider_fpl in playlists.iter_mut() { + log_memory_snapshot(format!("target '{}' input '{}' before_pipe", target.name, provider_fpl.input.name).as_str()); step.broadcast("Executing transformations on '{}' playlist", &target.name); - let mut processed_fpl = execute_pipe(target, &pipe, provider_fpl, &mut duplicates); + let mut processed_fpl = execute_pipe(target, &pipe, provider_fpl, &mut duplicates, consume_input_source); + log_memory_snapshot(format!("target '{}' input '{}' after_pipe", target.name, provider_fpl.input.name).as_str()); processed_fpl.sort_by_provider_ordinal(); playlist_resolve_series(&ctx.config, &ctx.client, target, errors, &pipe, provider_fpl, &mut processed_fpl).await; + log_memory_snapshot(format!("target '{}' input '{}' after_series_resolve", target.name, provider_fpl.input.name).as_str()); playlist_resolve_vod(&ctx.config, &ctx.client, target, errors, provider_fpl, &mut processed_fpl).await; + log_memory_snapshot(format!("target '{}' input '{}' after_vod_resolve", target.name, provider_fpl.input.name).as_str()); // stats let input_entry_name = processed_fpl.input.name.clone(); let group_count = processed_fpl.get_group_count(); @@ -734,15 +766,19 @@ async fn process_playlist_for_target(ctx: &PlaylistProcessingContext, stat.processed_stats.group_count = group_count; stat.processed_stats.channel_count = channel_count; } - processed_fetched_playlists.push(processed_fpl); + process_playlist_epg(&mut processed_fpl, &mut new_epg).await; + log_memory_snapshot(format!("target '{}' input '{}' after_epg_apply", target.name, processed_fpl.input.name).as_str()); + new_playlist.extend(processed_fpl.source.take_groups()); + log_memory_snapshot(format!("target '{}' input '{}' after_take_groups", target.name, processed_fpl.input.name).as_str()); + tokio::task::yield_now().await; } - step.tick("filter rename map"); - let (new_epg, mut new_playlist) = process_epg(&mut processed_fetched_playlists).await; - step.tick("epg"); + step.tick("filter rename map + epg"); + log_memory_snapshot(format!("target '{}' after_filter_rename_map_epg", target.name).as_str()); if target.favourites.is_some() { step.broadcast("Processing favourites for '{}' playlist", &target.name); process_favourites(&mut new_playlist, target.favourites.as_deref()); + log_memory_snapshot(format!("target '{}' after_favourites", target.name).as_str()); } if new_playlist.is_empty() { @@ -753,24 +789,31 @@ async fn process_playlist_for_target(ctx: &PlaylistProcessingContext, // Process Trakt categories if trakt_playlist(&ctx.client, target, errors, &mut new_playlist).await { step.tick("trakt categories"); + log_memory_snapshot(format!("target '{}' after_trakt", target.name).as_str()); } let mut flat_new_playlist = flatten_groups(new_playlist); step.tick("playlist merge"); + log_memory_snapshot(format!("target '{}' after_playlist_merge", target.name).as_str()); if sort_playlist(target, &mut flat_new_playlist) { step.tick("playlist sort"); + log_memory_snapshot(format!("target '{}' after_playlist_sort", target.name).as_str()); } assign_channel_no_playlist(&mut flat_new_playlist); step.tick("assigning channel numbers"); + log_memory_snapshot(format!("target '{}' after_assign_channel_numbers", target.name).as_str()); map_playlist_counter(target, &mut flat_new_playlist); step.tick("assigning channel counter"); + log_memory_snapshot(format!("target '{}' after_assign_channel_counter", target.name).as_str()); if process_watch(&ctx.config, &ctx.client, target, &flat_new_playlist).await { step.tick("group watches"); + log_memory_snapshot(format!("target '{}' after_group_watches", target.name).as_str()); } let result = persist_playlist(&ctx.config, &mut flat_new_playlist, flatten_tvguide(new_epg).as_ref(), target, ctx.playlist_state.as_ref()).await; step.stop("Persisting playlists"); + log_memory_snapshot(format!("target '{}' after_persist", target.name).as_str()); result } } @@ -832,19 +875,6 @@ async fn trakt_playlist(client: &reqwest::Client, target: &ConfigTarget, errors: true } -async fn process_epg(processed_fetched_playlists: &mut Vec>) -> (Vec, Vec) { - let mut new_playlist: Vec = vec![]; - let mut new_epg = vec![]; - - // each fetched playlist can have its own epgl url. - // we need to process each input epg. - for fp in processed_fetched_playlists { - process_playlist_epg(fp, &mut new_epg).await; - new_playlist.extend(fp.source.take_groups()); - } - (new_epg, new_playlist) -} - async fn process_watch(app_config: &Arc, client: &reqwest::Client, target: &ConfigTarget, new_playlist: &[PlaylistGroup]) -> bool { if let Some(watches) = &target.watch { if default_as_default().eq_ignore_ascii_case(&target.name) { @@ -883,6 +913,8 @@ pub async fn exec_processing(client: &reqwest::Client, app_config: Arc, client: &reqw group_series.into_values().collect() } -#[allow(clippy::too_many_arguments)] +#[allow(clippy::too_many_arguments, clippy::too_many_lines)] pub async fn playlist_resolve_series(cfg: &Arc, client: &reqwest::Client, target: &ConfigTarget, @@ -152,14 +153,106 @@ pub async fn playlist_resolve_series(cfg: &Arc, provider_fpl.source.release_resources(XtreamCluster::Series); let series_playlist = playlist_resolve_series_info(cfg, client, errors, processed_fpl, resolve_series, resolve_delay).await; provider_fpl.source.obtain_resources().await; - if series_playlist.is_empty() { return; } + // Keep input playlist as canonical source: + // - Store resolved series details back into provider_fpl items + // - Rebuild base (unmodified) episodes from provider_fpl metadata only if provider_fpl.is_memory() { - // original content saved into original list - for plg in &series_playlist { - provider_fpl.update_playlist(plg).await; + let mut resolved_series_by_provider_id: HashMap = HashMap::new(); + for pli in processed_fpl.items() { + if pli.header.xtream_cluster != XtreamCluster::Series + || pli.header.item_type != PlaylistItemType::SeriesInfo + { + continue; + } + let Some(provider_id) = pli.get_provider_id() else { continue; }; + if provider_id == 0 { + continue; + } + if let Some(StreamProperties::Series(properties)) = pli.header.additional_properties.as_ref() { + resolved_series_by_provider_id + .entry(provider_id) + .or_insert_with(|| properties.as_ref().clone()); + } + } + + if !resolved_series_by_provider_id.is_empty() { + let mut existing_uuids: HashSet = provider_fpl + .items() + .filter(|pli| { + pli.header.xtream_cluster == XtreamCluster::Series + && pli.header.item_type == PlaylistItemType::Series + }) + .map(|pli| pli.get_uuid()) + .collect(); + + let mut base_series_groups: IndexMap = IndexMap::new(); + let input = provider_fpl.input; + for pli in provider_fpl.items_mut() { + if pli.header.xtream_cluster != XtreamCluster::Series + || pli.header.item_type != PlaylistItemType::SeriesInfo + { + continue; + } + + let Some(provider_id) = pli.get_provider_id() else { continue; }; + if provider_id == 0 { + continue; + } + + if let Some(resolved) = resolved_series_by_provider_id.get(&provider_id) { + pli.header.additional_properties = Some(StreamProperties::Series(Box::new(resolved.clone()))); + } + + let Some(StreamProperties::Series(series_props)) = pli.header.additional_properties.as_ref() else { + continue; + }; + + let (category_id, group_title, series_name) = { + let header = &pli.header; + ( + header.category_id, + header.group.clone(), + if header.name.is_empty() { + header.title.clone() + } else { + header.name.clone() + }, + ) + }; + + if let Some(episodes) = parse_xtream_series_info( + &pli.get_uuid(), + series_props, + &group_title, + &series_name, + input, + ) { + let group = base_series_groups.entry(category_id).or_insert_with(|| PlaylistGroup { + id: category_id, + title: group_title, + channels: Vec::new(), + xtream_cluster: XtreamCluster::Series, + }); + for episode in episodes { + if existing_uuids.insert(episode.get_uuid()) { + group.channels.push(episode); + } + } + } + } + + for group in base_series_groups.values().filter(|g| !g.channels.is_empty()) { + provider_fpl.update_playlist(group).await; + } } } + + if series_playlist.is_empty() { + return; + } + + // run the target processing pipe over newly derived episodes // run the processing pipe over new items let mut new_playlist = series_playlist; for f in pipe { diff --git a/backend/src/processing/processor/xtream_vod.rs b/backend/src/processing/processor/xtream_vod.rs index ca08df03e..20e6d30f4 100644 --- a/backend/src/processing/processor/xtream_vod.rs +++ b/backend/src/processing/processor/xtream_vod.rs @@ -8,6 +8,7 @@ use log::{error, info, log_enabled, Level}; use shared::error::TuliproxError; use shared::model::{InputType, PlaylistEntry, StreamProperties, VideoStreamProperties, XtreamVideoInfo}; use shared::model::{PlaylistItemType, XtreamCluster}; +use std::collections::HashMap; use std::sync::Arc; use std::time::Instant; @@ -41,6 +42,9 @@ pub async fn playlist_resolve_vod(app_config: &Arc, let mut last_log_time = Instant::now(); let mut processed_vod_info_count = 0; let mut batch = Vec::with_capacity(BATCH_SIZE); + let sync_canonical_in_memory = provider_fpl.is_memory(); + let mut resolved_vod_by_provider_id: Option> = + sync_canonical_in_memory.then(HashMap::new); provider_fpl.source.release_resources(XtreamCluster::Video); @@ -72,6 +76,9 @@ pub async fn playlist_resolve_vod(app_config: &Arc, let video_stream_props = VideoStreamProperties::from_info(&info, pli); batch.push((provider_id, video_stream_props.clone())); + if let Some(resolved) = resolved_vod_by_provider_id.as_mut() { + resolved.insert(provider_id, video_stream_props.clone()); + } if batch.len() >= BATCH_SIZE { if let Err(err) = persist_input_vod_info_batch(app_config, &storage_path, XtreamCluster::Video, &input.name, std::mem::take(&mut batch)).await { error!("Failed to persist batch VOD info: {err}"); @@ -99,6 +106,24 @@ pub async fn playlist_resolve_vod(app_config: &Arc, } } + // Keep canonical input playlist updated for subsequent target processing in memory mode. + if let Some(resolved_vod_by_provider_id) = resolved_vod_by_provider_id { + for source_pli in provider_fpl.items_mut() { + if source_pli.header.xtream_cluster != XtreamCluster::Video + || source_pli.header.item_type != PlaylistItemType::Video + { + continue; + } + let Some(provider_id) = source_pli.get_provider_id() else { continue; }; + if provider_id == 0 { + continue; + } + if let Some(resolved) = resolved_vod_by_provider_id.get(&provider_id) { + source_pli.header.additional_properties = Some(StreamProperties::Video(Box::new(resolved.clone()))); + } + } + } + provider_fpl.source.obtain_resources().await; info!("resolved {processed_vod_info_count}/{vod_info_count} vod info"); } diff --git a/backend/src/repository/bplustree.rs b/backend/src/repository/bplustree.rs index 3d80c3b80..f396f095c 100644 --- a/backend/src/repository/bplustree.rs +++ b/backend/src/repository/bplustree.rs @@ -10,14 +10,19 @@ use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use shared::error::{string_to_io_error, to_io_error}; use std::ffi::OsString; -use std::fs::{File, OpenOptions}; +use std::fs::{File, Metadata, OpenOptions}; use std::io::{self, BufReader, Read, Seek, SeekFrom, Write}; use std::marker::PhantomData; use std::mem::size_of; #[cfg(unix)] -use std::os::unix::fs::FileExt; +use std::os::unix::fs::{FileExt, MetadataExt}; use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant, UNIX_EPOCH}; use tempfile::NamedTempFile; +#[cfg(windows)] +use windows_sys::Win32::Foundation::CloseHandle; +#[cfg(windows)] +use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION}; // Constants (Restored) const PAGE_SIZE: u16 = 4096; @@ -2721,10 +2726,41 @@ where /// at the cost of higher memory usage. type NodeCache = IndexMap, Option>)>; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct FileIdentity { + len: u64, + modified_ns: u128, + #[cfg(unix)] + dev: u64, + #[cfg(unix)] + ino: u64, +} + +impl FileIdentity { + fn from_metadata(metadata: &Metadata) -> Self { + let modified_ns = metadata + .modified() + .or_else(|_| metadata.created()) + .ok() + .and_then(|timestamp| timestamp.duration_since(UNIX_EPOCH).ok()) + .map_or(0, |duration| duration.as_nanos()); + + Self { + len: metadata.len(), + modified_ns, + #[cfg(unix)] + dev: metadata.dev(), + #[cfg(unix)] + ino: metadata.ino(), + } + } +} + pub struct BPlusTreeQuery { file: Option>, mmap: Option, filepath: PathBuf, + file_identity: Option, buffer: Vec, cache: IndexMap>, node_cache: NodeCache, @@ -2741,6 +2777,7 @@ where pub fn try_from_file(file: File) -> io::Result { let metadata = file.metadata()?; let file_len = metadata.len(); + let file_identity = Some(FileIdentity::from_metadata(&metadata)); if file_len < HEADER_SIZE { return Err(io::Error::new(io::ErrorKind::InvalidData, "File too small")); @@ -2773,6 +2810,7 @@ where file: if mmap.is_some() { None } else { Some(utils::file_reader(file)) }, mmap, filepath: PathBuf::new(), + file_identity, buffer: vec![0u8; PAGE_SIZE_USIZE], cache: IndexMap::with_capacity(CACHE_CAPACITY), node_cache: IndexMap::with_capacity(CACHE_CAPACITY), @@ -2794,7 +2832,73 @@ where &self.filepath } + fn read_root_offset_from_file(file: &File) -> io::Result { + let mut root_buf = [0u8; size_of::()]; + #[cfg(unix)] + file.read_exact_at(&mut root_buf, ROOT_OFFSET_POS)?; + #[cfg(not(unix))] + { + let mut f = file; + let current_pos = f.stream_position()?; + f.seek(SeekFrom::Start(ROOT_OFFSET_POS))?; + f.read_exact(&mut root_buf)?; + f.seek(SeekFrom::Start(current_pos))?; + } + Ok(u64::from_le_bytes(root_buf)) + } + + fn refresh_root_offset(&mut self) -> io::Result<()> { + let new_root_offset = if self.mmap.is_some() && !self.filepath.as_os_str().is_empty() { + let file = File::open(&self.filepath)?; + let metadata = file.metadata()?; + let current_identity = FileIdentity::from_metadata(&metadata); + let remap_required = self.file_identity != Some(current_identity); + let root_offset = Self::read_root_offset_from_file(&file)?; + + if remap_required { + if let Some(remapped) = unsafe { Mmap::map(&file).ok() } { + self.file = None; + self.mmap = Some(remapped); + } else { + self.mmap = None; + self.file = Some(utils::file_reader(file)); + } + self.cache.clear(); + self.node_cache.clear(); + } + + self.file_identity = Some(current_identity); + root_offset + } else if let Some(mmap) = &self.mmap { + let start = usize::try_from(ROOT_OFFSET_POS) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("Invalid root offset position: {e}")))?; + let end = start + .checked_add(size_of::()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset range"))?; + let root_bytes = mmap + .get(start..end) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Header too small"))?; + let bytes: [u8; size_of::()] = root_bytes + .try_into() + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset bytes"))?; + u64::from_le_bytes(bytes) + } else if let Some(file) = &mut self.file { + Self::read_root_offset_from_file(file.get_ref())? + } else { + return Err(io::Error::new(io::ErrorKind::InvalidData, "No data source available")); + }; + + if new_root_offset != self.root_offset { + self.root_offset = new_root_offset; + self.cache.clear(); + self.node_cache.clear(); + } + + Ok(()) + } + pub fn query(&mut self, key: &K) -> Result, BPlusTreeError> { + self.refresh_root_offset().map_err(BPlusTreeError::Io)?; if let Some(mmap) = &self.mmap { let mut cursor = io::Cursor::new(mmap.as_ref()); query_tree_mmap(mmap, &mut cursor, &mut self.node_cache, key, self.root_offset) @@ -2828,6 +2932,7 @@ where /// let value = query.query_zero_copy(&42)?; /// ``` pub fn query_zero_copy(&mut self, key: &K) -> Result, BPlusTreeError> { + self.refresh_root_offset().map_err(BPlusTreeError::Io)?; if let Some(mmap) = &self.mmap { let mut cursor = io::Cursor::new(mmap.as_ref()); query_tree_mmap_zero_copy(mmap, &mut cursor, key, self.root_offset) @@ -2847,6 +2952,7 @@ where V: Serialize + for<'de> Deserialize<'de> + Clone, { pub fn query_le(&mut self, key: &K) -> Result, BPlusTreeError> { + self.refresh_root_offset().map_err(BPlusTreeError::Io)?; if let Some(mmap) = &self.mmap { let mut cursor = io::Cursor::new(mmap.as_ref()); query_tree_le_mmap(mmap, &mut cursor, key, self.root_offset) @@ -2858,6 +2964,7 @@ where } pub fn len(&mut self) -> Result { + self.refresh_root_offset().map_err(BPlusTreeError::Io)?; if let Some(mmap) = &self.mmap { count_items_mmap::(mmap, self.root_offset) } else if let Some(file) = &mut self.file { @@ -2868,6 +2975,7 @@ where } pub fn is_empty(&mut self) -> Result { + self.refresh_root_offset().map_err(BPlusTreeError::Io)?; let (node, _) = if let Some(mmap) = &self.mmap { let mut cursor = io::Cursor::new(mmap.as_ref()); BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, self.root_offset, false)? @@ -3201,6 +3309,12 @@ pub enum FlushPolicy { None, } +#[derive(Debug, Clone, Copy)] +struct BatchRollbackState { + file_len: u64, + root_offset: u64, +} + pub struct BPlusTreeUpdate { file: BufReader, read_buffer: Vec, @@ -3312,6 +3426,98 @@ where }) } + /// Opens an update handle with lock-acquisition backoff. + /// + /// This function is **blocking** and should not run directly on a Tokio worker + /// thread. Internally it delegates to [`Self::try_new_with_backoff_stats`], + /// which uses `std::thread::sleep` while waiting for the file lock. + /// + /// In async contexts (Axum/Tokio), call this inside + /// `tokio::task::spawn_blocking(...)`. + pub fn try_new_with_backoff(filepath: &Path) -> io::Result { + Self::try_new_with_backoff_stats(filepath).map(|(tree, _)| tree) + } + + /// Opens an update handle with lock-acquisition backoff and returns + /// `(tree, retry_attempts)`. + /// + /// This function is **blocking**. It performs lock retries with + /// `std::thread::sleep`, so running it on a Tokio worker thread can stall + /// other async tasks scheduled on that thread. + /// + /// Use `tokio::task::spawn_blocking(...)` when calling from async code. + pub fn try_new_with_backoff_stats(filepath: &Path) -> io::Result<(Self, u64)> { + let mut attempts = 0u64; + let mut backoff = Duration::from_micros(10); + let max_backoff = Duration::from_millis(10); + let started_at = Instant::now(); + let timeout = Duration::from_secs(30); + + loop { + match Self::try_new(filepath) { + Ok(tree) => return Ok((tree, attempts)), + Err(err) if err.kind() == io::ErrorKind::WouldBlock => { + if started_at.elapsed() >= timeout { + return Err(io::Error::new(io::ErrorKind::TimedOut, "timeout acquiring lock")); + } + attempts = attempts.saturating_add(1); + std::thread::sleep(backoff); + backoff = (backoff * 2).min(max_backoff); + } + Err(err) => return Err(err), + } + } + } + + pub fn set_flush_policy(&mut self, policy: FlushPolicy) { + self.flush_policy = policy; + } + + /// Prepares a batch by cloning keys and serializing values eagerly. + /// + /// Keys are cloned because the returned prepared payload must outlive the + /// borrowed input slice. For large batches or expensive key types, this + /// can be a notable allocation/cloning cost. + pub fn prepare_upsert_batch(items: &[(&K, &V)]) -> io::Result)>> { + let mut prepared = Vec::with_capacity(items.len()); + for (key, value) in items { + prepared.push(((*key).clone(), binary_serialize(*value)?)); + } + Ok(prepared) + } + + pub fn upsert_batch_prepared_with_backoff(filepath: &Path, items: &[(&K, &V)]) -> io::Result { + let prepared = Self::prepare_upsert_batch(items)?; + let mut updater = Self::try_new_with_backoff(filepath)?; + updater.upsert_batch_encoded(prepared) + } + + fn capture_batch_rollback_state(&self) -> io::Result { + Ok(BatchRollbackState { + file_len: self.file.get_ref().metadata()?.len(), + root_offset: self.root_offset, + }) + } + + fn rollback_batch_state(&mut self, state: BatchRollbackState) -> io::Result<()> { + { + let file = self.file.get_mut(); + file.seek(SeekFrom::Start(ROOT_OFFSET_POS))?; + file.write_all(&state.root_offset.to_le_bytes())?; + file.flush()?; + file.set_len(state.file_len)?; + if self.flush_policy != FlushPolicy::None { + file.sync_all()?; + } + } + + // Seek via BufReader to invalidate any buffered state. + self.file.seek(SeekFrom::Start(state.file_len))?; + self.cache.clear(); + self.root_offset = state.root_offset; + Ok(()) + } + /// Helper to access metadata pub fn get_metadata(&mut self) -> io::Result { self.file.seek(SeekFrom::Start(METADATA_OFFSET_POS))?; @@ -3376,7 +3582,7 @@ where pub fn update(&mut self, key: &K, value: V) -> Result { let refs = [(key, &value)]; - let new_root_offset = self.update_batch_recursive(self.root_offset, &refs)?; + let new_root_offset = self.update_batch_recursive(self.root_offset, &refs, true)?; // Only update header if root offset actually changed // (if in-place update succeeded, offset stays the same) @@ -3398,10 +3604,8 @@ where /// Update multiple items in batch. This is more efficient than calling `update()` multiple times /// as it performs all updates and then commits the final root offset once. /// - /// **Optimization**: This method first attempts in-place value updates for values stored in - /// Single mode. If all values fit in their existing allocated space, no node rewrites are - /// needed at all - only the value data is updated. This eliminates the cascading rewrite - /// problem for same-size or smaller value updates. + /// For rollback safety in `update_batch`, callers may disable in-place updates and force + /// copy-on-write mode for all touched values. /// /// returns The final root offset after all updates, or an error if any update fails. /// Returns the original offset if all updates were performed in-place. @@ -3409,6 +3613,7 @@ where &mut self, offset: u64, items: &[(&K, &V)], + allow_in_place: bool, ) -> Result { let (mut node, pointers_opt) = BPlusTreeNode::::deserialize_from_block(&mut self.file, &mut self.read_buffer, offset, false)?; @@ -3421,22 +3626,27 @@ where for (key, value) in items { match node.keys.binary_search(key) { Ok(idx) => { - // Try in-place update first - let result = self.try_update_value_in_place(value, &node.value_info[idx]) - .map_err(BPlusTreeError::Io)?; + if allow_in_place { + // Try in-place update first + let result = self.try_update_value_in_place(value, &node.value_info[idx]) + .map_err(BPlusTreeError::Io)?; - match result { - InPlaceUpdateResult::Success => { - // In-place succeeded, value_info stays the same - } - InPlaceUpdateResult::PromotedToSingle(new_info) => { - // Packed value promoted to Single, need to update node - promoted.push((idx, new_info)); - } - InPlaceUpdateResult::NeedsCow => { - // Value doesn't fit in existing space, need full COW - needs_cow.push((idx, *key, *value)); + match result { + InPlaceUpdateResult::Success => { + // In-place succeeded, value_info stays the same + } + InPlaceUpdateResult::PromotedToSingle(new_info) => { + // Packed value promoted to Single, need to update node + promoted.push((idx, new_info)); + } + InPlaceUpdateResult::NeedsCow => { + // Value doesn't fit in existing space, need full COW + needs_cow.push((idx, *key, *value)); + } } + } else { + // Batch mode uses pure COW so rollback can safely truncate/revert root. + needs_cow.push((idx, *key, *value)); } } Err(_) => return Err(BPlusTreeError::KeyNotFound), @@ -3483,7 +3693,7 @@ where let sub_items = &items[current_idx..group_end]; let original_child_offset = pointers[child_idx]; - let new_child_offset = self.update_batch_recursive(original_child_offset, sub_items)?; + let new_child_offset = self.update_batch_recursive(original_child_offset, sub_items, allow_in_place)?; if new_child_offset != original_child_offset { pointers[child_idx] = new_child_offset; @@ -3511,24 +3721,41 @@ where return Ok(self.root_offset); } - let mut sorted_items = items.to_vec(); - sorted_items.sort_by(|a, b| a.0.cmp(b.0)); + let rollback_state = self.capture_batch_rollback_state().map_err(BPlusTreeError::Io)?; + let result = (|| -> Result { + let mut sorted_items = items.to_vec(); + sorted_items.sort_by(|a, b| a.0.cmp(b.0)); - let new_root_offset = self.update_batch_recursive(self.root_offset, &sorted_items)?; + // Disable in-place updates for batch rollback safety. + let new_root_offset = self.update_batch_recursive(self.root_offset, &sorted_items, false)?; - // Only update header if root offset actually changed - // (if all in-place updates succeeded, offset stays the same) - if new_root_offset != self.root_offset { - // Atomic Header Swap - only once at the end - self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS)).map_err(BPlusTreeError::Io)?; - self.file.get_mut().write_all(&new_root_offset.to_le_bytes()).map_err(BPlusTreeError::Io)?; - self.root_offset = new_root_offset; + // Only update header if root offset actually changed + // (if all in-place updates succeeded, offset stays the same) + if new_root_offset != self.root_offset { + // Atomic Header Swap - only once at the end + self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS)).map_err(BPlusTreeError::Io)?; + self.file.get_mut().write_all(&new_root_offset.to_le_bytes()).map_err(BPlusTreeError::Io)?; + self.root_offset = new_root_offset; + } + + self.file.get_mut().flush().map_err(BPlusTreeError::Io)?; + if self.flush_policy != FlushPolicy::None { + self.file.get_mut().sync_all().map_err(BPlusTreeError::Io)?; + } + Ok(new_root_offset) + })(); + + match result { + Ok(value) => Ok(value), + Err(err) => { + if let Err(rollback_err) = self.rollback_batch_state(rollback_state) { + return Err(BPlusTreeError::Io(io::Error::other(format!( + "update_batch failed: {err}; rollback failed: {rollback_err}" + )))); + } + Err(err) + } } - - self.file.get_mut().flush().map_err(BPlusTreeError::Io)?; - self.file.get_mut().sync_all().map_err(BPlusTreeError::Io)?; - - Ok(new_root_offset) } @@ -3845,23 +4072,42 @@ where return Ok(self.root_offset); } - let mut sorted_items = items.to_vec(); - sorted_items.sort_by(|a, b| a.0.cmp(b.0)); + let rollback_state = self.capture_batch_rollback_state()?; + let result = (|| -> io::Result { + let mut sorted_items = items.to_vec(); + sorted_items.sort_by(|a, b| a.0.cmp(b.0)); - let (mut current_root, promotions) = self.upsert_batch_recursive(self.root_offset, &sorted_items)?; + let (mut current_root, promotions) = self.upsert_batch_recursive(self.root_offset, &sorted_items)?; - // Handle promotions (splits) using balanced approach - current_root = self.build_higher_levels(current_root, promotions)?; + // Handle promotions (splits) using balanced approach + current_root = self.build_higher_levels(current_root, promotions)?; - self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS))?; - self.file.get_mut().write_all(¤t_root.to_le_bytes())?; - self.file.get_mut().flush()?; - if self.flush_policy != FlushPolicy::None { - self.file.get_mut().sync_all()?; + self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS))?; + self.file.get_mut().write_all(¤t_root.to_le_bytes())?; + self.file.get_mut().flush()?; + if self.flush_policy != FlushPolicy::None { + self.file.get_mut().sync_all()?; + } + + self.root_offset = current_root; + Ok(current_root) + })(); + + match result { + Ok(value) => Ok(value), + Err(err) => { + if let Err(rollback_err) = self.rollback_batch_state(rollback_state) { + return Err(io::Error::other(format!( + "upsert_batch failed: {err}; rollback failed: {rollback_err}" + ))); + } + Err(err) + } } + } - self.root_offset = current_root; - Ok(current_root) + pub fn upsert_batch_encoded(&mut self, items: Vec<(K, Vec)>) -> io::Result { + self.upsert_batch_preserialized(items) } /// Upsert multiple items using pre-serialized key-value data. @@ -3888,32 +4134,50 @@ where return Ok(self.root_offset); } - // Deserialize keys for sorting and tree traversal, keep values as bytes - let mut typed_items: Vec<(K, Vec)> = Vec::with_capacity(items.len()); - for (key, value_bytes) in items { - typed_items.push((key, value_bytes)); + let rollback_state = self.capture_batch_rollback_state()?; + let result = (|| -> io::Result { + // // Sort by key for efficient batch insertion + let mut sorted_items = items; + sorted_items.sort_by(|a, b| a.0.cmp(&b.0)); + + let (mut current_root, promotions) = self.upsert_batch_preserialized_recursive( + self.root_offset, + &sorted_items, + )?; + + // Handle promotions (splits) using a balanced approach + current_root = self.build_higher_levels(current_root, promotions)?; + + self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS))?; + self.file.get_mut().write_all(¤t_root.to_le_bytes())?; + self.file.get_mut().flush()?; + if self.flush_policy != FlushPolicy::None { + self.file.get_mut().sync_all()?; + } + + self.root_offset = current_root; + Ok(current_root) + })(); + + match result { + Ok(value) => Ok(value), + Err(err) => { + if let Err(rollback_err) = self.rollback_batch_state(rollback_state) { + return Err(io::Error::other(format!( + "upsert_batch_preserialized failed: {err}; rollback failed: {rollback_err}" + ))); + } + Err(err) + } } + } - // Sort by key for efficient batch insertion - typed_items.sort_by(|a, b| a.0.cmp(&b.0)); - - let (mut current_root, promotions) = self.upsert_batch_preserialized_recursive( - self.root_offset, - &typed_items, - )?; - - // Handle promotions (splits) using balanced approach - current_root = self.build_higher_levels(current_root, promotions)?; - - self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS))?; - self.file.get_mut().write_all(¤t_root.to_le_bytes())?; + pub fn commit(&mut self) -> io::Result<()> { self.file.get_mut().flush()?; if self.flush_policy != FlushPolicy::None { self.file.get_mut().sync_all()?; } - - self.root_offset = current_root; - Ok(current_root) + Ok(()) } /// Recursive helper for `upsert_batch_preserialized`. @@ -4135,7 +4399,7 @@ where temp_file.write_all(&root_offset.to_le_bytes())?; temp_file.flush()?; - // temp_file.as_file().sync_all()?; // Removed as requested; relying on persist or OS flush policy + // temp_file.as_file().sync_all()?; // Removed as requested; relying on persisted or OS flush policy // 5. Atomic Replace temp_file.persist(filepath).map_err(to_io_error)?; @@ -4150,6 +4414,49 @@ where } } +/// Single-writer wrapper that serializes prepared batch writes through one updater instance. +/// +/// This wrapper does not implement write-ahead logging (WAL) semantics or +/// crash-recovery replay. It only serializes access through a shared updater. +pub struct BPlusTreeSerialWriter { + updater: Mutex>, +} + +impl BPlusTreeSerialWriter +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, + V: Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, +{ + pub fn new(filepath: &Path, flush_policy: FlushPolicy) -> io::Result { + let mut updater = BPlusTreeUpdate::::try_new_with_backoff(filepath)?; + updater.set_flush_policy(flush_policy); + Ok(Self { + updater: Mutex::new(updater), + }) + } + + pub fn upsert_prepared(&self, items: Vec<(K, Vec)>) -> io::Result { + self.updater.lock().upsert_batch_encoded(items) + } + + pub fn upsert(&self, items: &[(&K, &V)]) -> io::Result { + let prepared = BPlusTreeUpdate::::prepare_upsert_batch(items)?; + self.upsert_prepared(prepared) + } + + pub fn commit(&self) -> io::Result<()> { + self.updater.lock().commit() + } + + /// Alias for `commit()`. + pub fn shutdown(&self) -> io::Result<()> { + self.commit() + } +} + +#[deprecated(note = "Not a real WAL implementation; use BPlusTreeSerialWriter instead")] +pub type BPlusTreeWalWriter = BPlusTreeSerialWriter; + pub struct BPlusTreeIterator<'a, K, V> { stack: Vec<&'a BPlusTreeNode>, current_keys: Option<&'a [K]>, @@ -4240,6 +4547,9 @@ where #[cfg(test)] mod tests { use super::*; + use serde::de::Deserializer; + use serde::ser::Error as SerError; + use serde::Serializer; use tempfile::tempdir; use parking_lot::Mutex; @@ -4250,6 +4560,59 @@ mod tests { use serde::{Deserialize, Serialize}; use shared::utils::generate_random_string; + + #[cfg(unix)] + #[allow(dead_code)] + fn process_is_alive(pid: u32) -> bool { + if pid == 0 { + return false; + } + + #[cfg(target_os = "linux")] + { + // Fast path for Linux/musl. + if Path::new("/proc").join(pid.to_string()).exists() { + return true; + } + } + + let pid_raw: libc::pid_t = match pid.try_into() { + Ok(value) => value, + Err(_) => return false, + }; + + // kill(pid, 0) probes process existence without sending a signal. + let rc = unsafe { libc::kill(pid_raw, 0) }; + if rc == 0 { + return true; + } + + match io::Error::last_os_error().raw_os_error() { + Some(code) if code == libc::ESRCH => false, + Some(code) if code == libc::EPERM => true, + _ => false, + } + } + + #[cfg(windows)] + #[allow(dead_code)] + fn process_is_alive(pid: u32) -> bool { + if pid == 0 { + return false; + } + + // PROCESS_QUERY_LIMITED_INFORMATION is sufficient to check process existence. + let handle = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid) }; + if handle == 0 { + return false; + } + + unsafe { + CloseHandle(handle); + } + true + } + // Example usage with a simple struct #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] struct Record { @@ -4257,6 +4620,57 @@ mod tests { data: String, } + #[derive(Debug, Clone, PartialEq, Eq)] + struct FailingValue { + payload: String, + fail_serialize: bool, + } + + impl FailingValue { + fn ok(payload: impl Into) -> Self { + Self { + payload: payload.into(), + fail_serialize: false, + } + } + + fn failing(payload: impl Into) -> Self { + Self { + payload: payload.into(), + fail_serialize: true, + } + } + } + + impl Serialize for FailingValue { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + if self.fail_serialize { + return Err(S::Error::custom("intentional serialize failure")); + } + self.payload.serialize(serializer) + } + } + + impl<'de> Deserialize<'de> for FailingValue { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let payload = String::deserialize(deserializer)?; + Ok(Self::ok(payload)) + } + } + + #[test] + fn test_process_is_alive_current_process() { + let current_pid = std::process::id(); + assert!(process_is_alive(current_pid)); + assert!(!process_is_alive(u32::MAX)); + } + #[test] fn insert_test() -> io::Result<()> { @@ -4790,6 +5204,107 @@ mod tests { Ok(()) } + #[test] + fn update_batch_error_rolls_back_file_and_data() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_update_batch_rollback.bin"); + let mut tree = BPlusTree::::new(); + let count = 2_000u32; + for i in 0..count { + tree.insert(i, format!("v{i}")); + } + tree.store(&filepath)?; + drop(tree); + + let size_before = std::fs::metadata(&filepath)?.len(); + let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; + let updates = [ + (0u32, "X".repeat(6_000)), + (u32::MAX, "missing".to_string()), + ]; + let refs: Vec<(&u32, &String)> = updates.iter().map(|(k, v)| (k, v)).collect(); + let result = updater.update_batch(&refs); + assert!(matches!(result, Err(BPlusTreeError::KeyNotFound))); + drop(updater); + + let size_after = std::fs::metadata(&filepath)?.len(); + assert_eq!(size_after, size_before, "failed batch must roll back file growth"); + + // Ensure header/root offset remains valid for update/load paths. + let mut reopened_updater = BPlusTreeUpdate::::try_new(&filepath)?; + assert_eq!(reopened_updater.query(&0).map_err(BPlusTreeError::to_io)?, Some("v0".to_string())); + drop(reopened_updater); + let loaded_tree = BPlusTree::::load(&filepath)?; + assert_eq!(loaded_tree.query(&0).cloned(), Some("v0".to_string())); + + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + assert_eq!(query.query(&0).map_err(BPlusTreeError::to_io)?, Some("v0".to_string())); + assert!(query.query(&u32::MAX).map_err(BPlusTreeError::to_io)?.is_none()); + + Ok(()) + } + + #[test] + fn update_batch_error_rolls_back_in_place_candidate_write() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_update_batch_in_place_rollback.bin"); + let original = "A".repeat(500); + let updated = "B".repeat(500); + + let mut tree = BPlusTree::::new(); + tree.insert(1, original.clone()); + tree.store(&filepath)?; + drop(tree); + + let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; + let updates = [ + (1u32, updated.clone()), + (u32::MAX, "missing".to_string()), + ]; + let refs: Vec<(&u32, &String)> = updates.iter().map(|(k, v)| (k, v)).collect(); + + let result = updater.update_batch(&refs); + assert!(matches!(result, Err(BPlusTreeError::KeyNotFound))); + drop(updater); + + // On failed batch, key 1 must still have original value. + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(original)); + assert_ne!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(updated)); + + Ok(()) + } + + #[test] + fn upsert_batch_error_rolls_back_file_and_data() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_upsert_batch_rollback.bin"); + let mut tree = BPlusTree::::new(); + tree.insert(1, FailingValue::ok("initial")); + tree.store(&filepath)?; + drop(tree); + + let size_before = std::fs::metadata(&filepath)?.len(); + let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; + let batch = [ + (1u32, FailingValue::ok("updated")), + (2u32, FailingValue::failing("boom")), + ]; + let refs: Vec<(&u32, &FailingValue)> = batch.iter().map(|(k, v)| (k, v)).collect(); + assert!(updater.upsert_batch(&refs).is_err()); + drop(updater); + + let size_after = std::fs::metadata(&filepath)?.len(); + assert_eq!(size_after, size_before, "failed upsert batch must roll back file growth"); + + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + let existing = query.query(&1).map_err(BPlusTreeError::to_io)?; + assert_eq!(existing.map(|v| v.payload), Some("initial".to_string())); + assert!(query.query(&2).map_err(BPlusTreeError::to_io)?.is_none()); + + Ok(()) + } + #[test] fn update_batch_large_test() -> io::Result<()> { let tempdir = tempfile::tempdir()?; @@ -4928,6 +5443,40 @@ mod tests { Ok(()) } + #[cfg(unix)] + #[test] + fn query_remaps_after_atomic_replace_with_same_length() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_query_replace_target.bin"); + let replacement_path = tempdir.path().join("tree_query_replace_source.bin"); + + let mut tree = BPlusTree::::new(); + tree.insert(1, Record { id: 1, data: "aaaa".to_string() }); + tree.insert(2, Record { id: 2, data: "bbbb".to_string() }); + tree.store(&filepath)?; + + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + assert!(query.mmap.is_some(), "Test requires mmap-backed query"); + let initial = query.query(&1).map_err(BPlusTreeError::to_io)?.expect("missing initial key"); + assert_eq!(initial.data, "aaaa"); + + let mut replacement = BPlusTree::::new(); + replacement.insert(1, Record { id: 1, data: "cccc".to_string() }); + replacement.insert(2, Record { id: 2, data: "dddd".to_string() }); + replacement.store(&replacement_path)?; + + let old_len = std::fs::metadata(&filepath)?.len(); + let replacement_len = std::fs::metadata(&replacement_path)?.len(); + assert_eq!(old_len, replacement_len, "test requires same-length replacement file"); + + std::fs::rename(&replacement_path, &filepath)?; + + let refreshed = query.query(&1).map_err(BPlusTreeError::to_io)?.expect("missing replaced key"); + assert_eq!(refreshed.data, "cccc"); + + Ok(()) + } + #[test] fn upsert_batch_mixed_test() -> io::Result<()> { let tempdir = tempfile::tempdir()?; @@ -5759,6 +6308,33 @@ mod tests { Ok(()) } + #[test] + fn test_query_refreshes_root_offset_after_external_upsert() -> io::Result<()> { + let tempdir = tempdir()?; + let filepath = tempdir.path().join("query_refresh_root.bin"); + + let mut tree = BPlusTree::::new(); + for i in 0..64 { + tree.insert(i, format!("value_{i}")); + } + tree.store(&filepath)?; + drop(tree); + + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + assert_eq!(query.query(&999).map_err(BPlusTreeError::to_io)?, None); + + let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; + let inserted = "fresh_value".to_string(); + updater.upsert_batch(&[(&999u32, &inserted)])?; + drop(updater); + + // Same query instance should observe the updated root/header. + assert_eq!(query.query(&999).map_err(BPlusTreeError::to_io)?, Some(inserted.clone())); + assert_eq!(query.query_zero_copy(&999).map_err(BPlusTreeError::to_io)?, Some(inserted)); + + Ok(()) + } + #[test] fn test_query_le_query_struct() -> io::Result<()> { let tempdir = tempdir()?; diff --git a/backend/src/repository/bplustree_stress.rs b/backend/src/repository/bplustree_stress.rs index c62cc5b8c..ad8d540f8 100644 --- a/backend/src/repository/bplustree_stress.rs +++ b/backend/src/repository/bplustree_stress.rs @@ -1,8 +1,10 @@ -use super::bplustree::{BPlusTree, BPlusTreeUpdate, BPlusTreeQuery}; +use super::bplustree::{BPlusTree, BPlusTreeQuery, BPlusTreeSerialWriter, BPlusTreeUpdate, FlushPolicy}; use rand::prelude::*; use rand::distr::Alphanumeric; use std::io::Write; -use std::time::Instant; +use std::sync::{Arc, Barrier}; +use std::thread; +use std::time::{Duration, Instant}; use tempfile::NamedTempFile; // Run with: `cargo test --release --package tuliprox -- stress_test_bplustree -- --nocapture` @@ -15,6 +17,24 @@ fn random_string(len: usize) -> String { .collect() } +#[inline] +fn lcg_next(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *state +} + +fn percentile_nearest_rank(sorted_values: &[f64], percentile: f64) -> f64 { + if sorted_values.is_empty() { + return 0.0; + } + let p = percentile.clamp(0.0, 1.0); + let n = sorted_values.len(); + let rank = ((p * n as f64).ceil() as usize).saturating_sub(1).min(n - 1); + sorted_values[rank] +} + #[test] #[allow( clippy::cast_precision_loss, @@ -29,35 +49,75 @@ fn stress_test_bplustree() { let mut log_file = std::fs::File::create(log_path).unwrap(); // Config - let num_items = 500_000; + let num_items = 1_000_000usize; let query_count = 50_000; + let insert_runs = 7usize; + let value_pool_size = 2_048usize; let small_val_len = 50; let large_val_len = 500; // Larger than packed limit (256) writeln!(log_file, "=== B+Tree Stress Test & Performance Analysis ===").unwrap(); writeln!(log_file, "Dataset: {num_items} items").unwrap(); + writeln!(log_file, "Insert benchmark runs: {insert_runs}").unwrap(); // ---------------------------------------------------------------- // Phase 1: Batch Insert (Sequential Keys) // ---------------------------------------------------------------- - writeln!(log_file, "\n[Phase 1] Batch Insert (Sequential)...").unwrap(); - let mut tree = BPlusTree::::new(); + writeln!(log_file, "\n[Phase 1] Batch Insert (Sequential, Multi-Run)...").unwrap(); let start_gen = Instant::now(); - - // Generate data - let mut initial_data = Vec::with_capacity(num_items); - for i in 0..num_items { - initial_data.push((i as u32, random_string(small_val_len))); - } - writeln!(log_file, "Generation Time: {:.2?}", start_gen.elapsed()).unwrap(); - let start = Instant::now(); - for (k, v) in initial_data { - tree.insert(k, v); + let mut value_pool = Vec::with_capacity(value_pool_size); + for _ in 0..value_pool_size { + value_pool.push(random_string(small_val_len)); } - let insert_duration = start.elapsed(); - writeln!(log_file, "Insert Time: {insert_duration:.2?}").unwrap(); - writeln!(log_file, "Throughput: {:.0} ops/sec", num_items as f64 / insert_duration.as_secs_f64()).unwrap(); + writeln!( + log_file, + "Value pool generation time ({} templates): {:.2?}", + value_pool_size, + start_gen.elapsed() + ) + .unwrap(); + + let mut insert_throughputs = Vec::with_capacity(insert_runs); + let mut final_tree: Option> = None; + for run_idx in 0..insert_runs { + let mut run_tree = BPlusTree::::new(); + let start = Instant::now(); + for i in 0..num_items { + let value = value_pool[i % value_pool_size].clone(); + run_tree.insert(i as u32, value); + } + let run_duration = start.elapsed(); + let run_throughput = num_items as f64 / run_duration.as_secs_f64(); + insert_throughputs.push(run_throughput); + writeln!( + log_file, + "Run {}: {:.2?} ({:.0} ops/sec)", + run_idx + 1, + run_duration, + run_throughput + ) + .unwrap(); + final_tree = Some(run_tree); + } + let mut sorted_insert_throughputs = insert_throughputs.clone(); + sorted_insert_throughputs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let insert_mean = insert_throughputs.iter().sum::() / insert_throughputs.len() as f64; + let insert_median = percentile_nearest_rank(&sorted_insert_throughputs, 0.5); + let insert_p95 = percentile_nearest_rank(&sorted_insert_throughputs, 0.95); + let insert_min = *sorted_insert_throughputs.first().unwrap_or(&0.0); + let insert_max = *sorted_insert_throughputs.last().unwrap_or(&0.0); + writeln!( + log_file, + "Insert Throughput Summary (ops/sec): mean={:.0}, median={:.0}, p95={:.0}, min={:.0}, max={:.0}", + insert_mean, + insert_median, + insert_p95, + insert_min, + insert_max + ) + .unwrap(); + let mut tree = final_tree.expect("final insert run should produce a tree"); // Prepare query keys let mut query_keys: Vec = (0..num_items as u32).collect(); @@ -172,4 +232,177 @@ fn stress_test_bplustree() { let query_duration = start.elapsed(); writeln!(log_file, "In-Memory Query Time: {query_duration:.2?}").unwrap(); writeln!(log_file, "Throughput: {:.0} ops/sec", query_count as f64 / query_duration.as_secs_f64()).unwrap(); + drop(tree_mem); + + // ---------------------------------------------------------------- + // Phase 7: Concurrent Readers + Writers (Disk-based) + // ---------------------------------------------------------------- + writeln!(log_file, "\n[Phase 7] Concurrent Readers + Writers (Disk-based)...").unwrap(); + let reader_threads = 8usize; + let writer_threads = 4usize; + let reader_ops_per_thread = 40_000u64; + let writer_batches_per_thread = 48u64; + let writer_batch_size = 128usize; + let num_items_u32 = num_items as u32; + let expected_writer_batches = writer_batches_per_thread * writer_threads as u64; + let expected_writer_updates = expected_writer_batches * writer_batch_size as u64; + + writeln!( + log_file, + "Readers: {reader_threads}, Writers: {writer_threads}, Reader Ops/Thread: {reader_ops_per_thread}, Writer Batches/Thread: {writer_batches_per_thread}, Batch Size: {writer_batch_size}" + ) + .unwrap(); + + let start_barrier = Arc::new(Barrier::new(reader_threads + writer_threads + 1)); + let shared_path = Arc::new(filepath.clone()); + let wal_writer = Arc::new( + BPlusTreeSerialWriter::::new(filepath.as_path(), FlushPolicy::Batch).unwrap(), + ); + let small_payload = "s".repeat(small_val_len.saturating_sub(24)); + let large_payload = "L".repeat(large_val_len.saturating_sub(24)); + + let mut reader_handles = Vec::with_capacity(reader_threads); + for reader_id in 0..reader_threads { + let barrier = Arc::clone(&start_barrier); + let path = Arc::clone(&shared_path); + reader_handles.push(thread::spawn(move || -> (u64, u64, Duration) { + let mut query = BPlusTreeQuery::::try_new(path.as_path()).unwrap(); + let mut prng_state = 0x9E37_79B9_7F4A_7C15u64 ^ ((reader_id as u64 + 1).wrapping_mul(0xBF58_476D_1CE4_E5B9)); + barrier.wait(); + let started = Instant::now(); + let mut ops = 0u64; + let mut misses = 0u64; + + for _ in 0..reader_ops_per_thread { + let key = (lcg_next(&mut prng_state) as u32) % num_items_u32; + if query.query_zero_copy(&key).unwrap().is_none() { + misses += 1; + } + ops += 1; + } + + (ops, misses, started.elapsed()) + })); + } + + let mut writer_handles = Vec::with_capacity(writer_threads); + for writer_id in 0..writer_threads { + let barrier = Arc::clone(&start_barrier); + let wal_writer = Arc::clone(&wal_writer); + let small_payload_local = small_payload.clone(); + let large_payload_local = large_payload.clone(); + writer_handles.push(thread::spawn(move || -> (u64, u64, u64, Duration) { + let mut prng_state = 0xD6E8_FEB8_6659_FD93u64 ^ ((writer_id as u64 + 1).wrapping_mul(0x94D0_49BB_1331_11EB)); + barrier.wait(); + let started = Instant::now(); + + let retries = 0u64; + let mut applied_batches = 0u64; + let mut applied_updates = 0u64; + + for batch_idx in 0..writer_batches_per_thread { + let mut owned_batch = Vec::with_capacity(writer_batch_size); + for _ in 0..writer_batch_size { + let key = (lcg_next(&mut prng_state) as u32) % num_items_u32; + let variant = lcg_next(&mut prng_state); + let payload = if (variant & 0b1111) == 0 { + &large_payload_local + } else { + &small_payload_local + }; + owned_batch.push((key, format!("w{writer_id}_b{batch_idx}_k{key}_{payload}"))); + } + let batch_refs: Vec<(&u32, &String)> = owned_batch.iter().map(|(k, v)| (k, v)).collect(); + // Pre-serialize + compress outside the file lock. + let prepared = BPlusTreeUpdate::::prepare_upsert_batch(&batch_refs).unwrap(); + wal_writer.upsert_prepared(prepared).unwrap(); + applied_batches += 1; + applied_updates += writer_batch_size as u64; + } + + (applied_updates, applied_batches, retries, started.elapsed()) + })); + } + + let concurrent_start = Instant::now(); + start_barrier.wait(); + + let mut total_reader_ops = 0u64; + let mut total_reader_misses = 0u64; + let mut total_reader_time = Duration::ZERO; + let mut max_reader_duration = Duration::ZERO; + for handle in reader_handles { + let (ops, misses, elapsed) = handle.join().unwrap(); + total_reader_ops += ops; + total_reader_misses += misses; + total_reader_time += elapsed; + if elapsed > max_reader_duration { + max_reader_duration = elapsed; + } + } + + let mut total_writer_updates = 0u64; + let mut total_writer_batches = 0u64; + let mut total_writer_retries = 0u64; + let mut total_writer_time = Duration::ZERO; + let mut max_writer_duration = Duration::ZERO; + for handle in writer_handles { + let (updates, batches, retries, elapsed) = handle.join().unwrap(); + total_writer_updates += updates; + total_writer_batches += batches; + total_writer_retries += retries; + total_writer_time += elapsed; + if elapsed > max_writer_duration { + max_writer_duration = elapsed; + } + } + wal_writer.commit().unwrap(); + wal_writer.shutdown().unwrap(); + + let concurrent_duration = concurrent_start.elapsed(); + let reader_avg_latency_us = total_reader_time.as_secs_f64() * 1_000_000.0 / total_reader_ops as f64; + let writer_avg_latency_us = total_writer_time.as_secs_f64() * 1_000_000.0 / total_writer_updates as f64; + + writeln!(log_file, "Wall Time: {concurrent_duration:.2?}").unwrap(); + writeln!( + log_file, + "Reader Ops: {total_reader_ops} (misses: {total_reader_misses}), Throughput: {:.0} ops/sec, Avg Latency: {:.2}us/op, Slowest Reader: {max_reader_duration:.2?}", + total_reader_ops as f64 / concurrent_duration.as_secs_f64(), + reader_avg_latency_us + ) + .unwrap(); + writeln!( + log_file, + "Writer Updates: {total_writer_updates} (batches: {total_writer_batches}), Throughput: {:.0} updates/sec, Avg Latency: {:.2}us/update, Lock Retries: {total_writer_retries} (single-writer queue), Slowest Writer: {max_writer_duration:.2?}", + total_writer_updates as f64 / concurrent_duration.as_secs_f64(), + writer_avg_latency_us + ) + .unwrap(); + + assert_eq!(total_writer_batches, expected_writer_batches, "all writer batches must be committed"); + assert_eq!(total_writer_updates, expected_writer_updates, "all writer updates must be committed"); + assert_eq!(total_reader_misses, 0, "all reader queries should resolve during concurrent load"); + + // ---------------------------------------------------------------- + // Phase 7b: Post-Concurrency Verification + // ---------------------------------------------------------------- + writeln!(log_file, "\n[Phase 7b] Post-Concurrency Verification...").unwrap(); + let mut verify_query = BPlusTreeQuery::::try_new(&filepath).unwrap(); + let verification_samples = 20_000u64; + let mut verification_misses = 0u64; + let mut verify_state = 0x243F_6A88_85A3_08D3u64; + let verify_start = Instant::now(); + for _ in 0..verification_samples { + let key = (lcg_next(&mut verify_state) as u32) % num_items_u32; + if verify_query.query(&key).unwrap().is_none() { + verification_misses += 1; + } + } + let verify_duration = verify_start.elapsed(); + writeln!( + log_file, + "Verification Time: {verify_duration:.2?}, misses: {verification_misses}/{verification_samples}" + ) + .unwrap(); + assert_eq!(verification_misses, 0, "post-concurrency verification failed"); } diff --git a/backend/src/repository/m3u_repository.rs b/backend/src/repository/m3u_repository.rs index ab6452c73..c59aabf0f 100644 --- a/backend/src/repository/m3u_repository.rs +++ b/backend/src/repository/m3u_repository.rs @@ -70,7 +70,7 @@ async fn persist_m3u_playlist_as_text( cfg: &Config, target: &ConfigTarget, target_output: &M3uTargetOutput, - m3u_playlist: Arc>, + m3u_playlist: &[M3uPlaylistItem], ) -> Result<(), TuliproxError> { let Some(filename) = target_output.filename.as_ref() else { return Ok(()); }; let Some(m3u_filename) = utils::get_file_path(&cfg.working_dir, Some(PathBuf::from(filename))) else { return Ok(()); }; @@ -82,7 +82,7 @@ async fn persist_m3u_playlist_as_text( let mut write_counter = 0usize; - for m3u in m3u_playlist.iter() { + for m3u in m3u_playlist { let line = m3u.to_m3u(target.options.as_ref(), false); let bytes = line.as_bytes(); await_playlist_write!(writer.write_all(bytes), "Failed to write entry to {} - {}", m3u_filename.display()); @@ -114,31 +114,27 @@ pub async fn m3u_write_playlist( let _m3u_path = ensure_m3u_storage_path(&config, target.name.as_str()).await?; let m3u_path = m3u_get_file_path_for_db(target_path); - let m3u_playlist = Arc::new( - new_playlist - .iter() - .flat_map(|pg| &pg.channels) - .filter(|&pli| !matches!(pli.header.item_type, PlaylistItemType::SeriesInfo | PlaylistItemType::LocalSeriesInfo)) - .map(M3uPlaylistItem::from) - .collect::>(), - ); + let m3u_playlist = new_playlist + .iter() + .flat_map(|pg| &pg.channels) + .filter(|&pli| !matches!(pli.header.item_type, PlaylistItemType::SeriesInfo | PlaylistItemType::LocalSeriesInfo)) + .map(M3uPlaylistItem::from) + .collect::>(); let file_lock = cfg.file_locks.write_lock(&m3u_path).await; - if let Err(err) = persist_m3u_playlist_as_text(&cfg.config.load(), target, target_output, Arc::clone(&m3u_playlist)).await { + if let Err(err) = persist_m3u_playlist_as_text(&cfg.config.load(), target, target_output, &m3u_playlist).await { error!("Persisting m3u playlist failed: {err}"); } - let playlist = Arc::clone(&m3u_playlist); let m3u_path_clone = m3u_path.clone(); // Move all B+Tree building and I/O to spawn_blocking - // We take ownership of `playlist` here (no cloning needed) task::spawn_blocking(move || -> Result<(), TuliproxError> { let _guard = file_lock; let mut tree = BPlusTree::new(); - for m3u in playlist.iter() { - tree.insert(m3u.virtual_id, m3u.clone()); + for m3u in m3u_playlist { + tree.insert(m3u.virtual_id, m3u); } tree.store_with_index(&m3u_path_clone, |pli| pli.source_ordinal).map_err(|err| cant_write_result!(&m3u_path_clone, err))?; Ok(()) @@ -202,7 +198,7 @@ pub async fn iter_raw_m3u_input_playlist(app_config: &AppConfig, input: &ConfigI async fn iter_raw_m3u_playlist(app_config: &AppConfig, m3u_path: &Path, cluster: Option) -> Option<(FileReadGuard, Box + Send>)> where - ItemKey: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, + ItemKey: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static, SortKey: for<'de> Deserialize<'de> + Send + 'static, { let file_lock = app_config.file_locks.read_lock(m3u_path).await; diff --git a/backend/src/repository/playlist_repository.rs b/backend/src/repository/playlist_repository.rs index 81b60351e..ae5445cd5 100644 --- a/backend/src/repository/playlist_repository.rs +++ b/backend/src/repository/playlist_repository.rs @@ -85,6 +85,8 @@ pub async fn persist_playlist(app_config: &Arc, playlist: &mut [Playl } rewrite_series_info_episode_virtual_id(playlist, &local_library_series, &provider_series); + drop(local_library_series); + drop(provider_series); for output in &target.output { let mut filtered: Option> = match output { @@ -123,8 +125,9 @@ pub async fn persist_playlist(app_config: &Arc, playlist: &mut [Playl if let Err(err) = target_id_mapping.persist() { errors.push(info_err!("{err}")); } - // We must drop target_id_mapping here to release the exclusive B+Tree lock - // otherwise the subsequent load_xtream_target_storage will deadlock waiting for a shared lock. + // Keep lock until all outputs are persisted to prevent concurrent writers + // from interleaving mapping and output state for the same target. + // We must release it before loading caches below (which may acquire read locks). drop(target_id_mapping); drop(file_lock); @@ -470,4 +473,4 @@ fn get_input_local_library_playlist_file_path(storage_path: &Path, input_name: & .map(|c| if c.is_alphanumeric() { c } else { '_' }) .collect(); storage_path.join(format!("lib_{sanitized_input_name}.{FILE_SUFFIX_DB}")) -} \ No newline at end of file +} diff --git a/backend/src/repository/target_id_mapping.rs b/backend/src/repository/target_id_mapping.rs index 7d89eae70..013e1287f 100644 --- a/backend/src/repository/target_id_mapping.rs +++ b/backend/src/repository/target_id_mapping.rs @@ -1,4 +1,4 @@ -use crate::repository::bplustree::{BPlusTree, BPlusTreeMetadata, BPlusTreeUpdate}; +use crate::repository::bplustree::{BPlusTree, BPlusTreeMetadata, BPlusTreeUpdate, FlushPolicy}; use chrono::Local; use log::error; use serde::{Deserialize, Serialize}; @@ -83,26 +83,29 @@ impl TargetIdMapping { .map_err(|e| info_err!("Failed to create UUID index at {}: {e}", uuid_index_path.display()))?; // Open disk-based update handles - let disk_by_virtual_id = match BPlusTreeUpdate::::try_new(path) { + let mut disk_by_virtual_id = match BPlusTreeUpdate::::try_new_with_backoff(path) { Ok(tree) => tree, Err(e) => { error!("Failed to open primary tree at {}: {e}", path.display()); // Create fresh and try again let _ = BPlusTree::::new().store(path); - BPlusTreeUpdate::try_new(path).map_err(|_| info_err!("Failed to create primary tree after retry"))? + BPlusTreeUpdate::try_new_with_backoff(path).map_err(|_| info_err!("Failed to create primary tree after retry"))? } }; - let disk_by_uuid = match BPlusTreeUpdate::::try_new(&uuid_index_path) { + let mut disk_by_uuid = match BPlusTreeUpdate::::try_new_with_backoff(&uuid_index_path) { Ok(tree) => tree, Err(e) => { error!("Failed to open UUID index at {}: {e}", uuid_index_path.display()); // Create fresh and try again let _ = BPlusTree::::new().store(&uuid_index_path); - BPlusTreeUpdate::try_new(&uuid_index_path).map_err(|_| info_err!("Failed to create UUID index after retry"))? + BPlusTreeUpdate::try_new_with_backoff(&uuid_index_path).map_err(|_| info_err!("Failed to create UUID index after retry"))? } }; + disk_by_virtual_id.set_flush_policy(FlushPolicy::Batch); + disk_by_uuid.set_flush_policy(FlushPolicy::Batch); + let mut virtual_id_counter: u32 = 0; // Load primary tree into memory @@ -213,6 +216,9 @@ impl TargetIdMapping { error!("Failed to write virtual_id_counter to tree header at {}: {e}", self.path.display()); e })?; + + self.disk_by_virtual_id.commit()?; + self.disk_by_uuid.commit()?; } Ok(()) } diff --git a/backend/src/repository/xtream_repository.rs b/backend/src/repository/xtream_repository.rs index 927eb5efc..dfddae4a3 100644 --- a/backend/src/repository/xtream_repository.rs +++ b/backend/src/repository/xtream_repository.rs @@ -132,26 +132,23 @@ pub async fn write_playlist_item_update( ensure_xtream_storage_path(&config, target_name).await? }; let xtream_path = xtream_get_file_path(&storage_path, pli.xtream_cluster); - - // Phase 1: Pre-serialize the item in async context (before spawn_blocking) - let value_bytes = crate::utils::binary_serialize(pli) - .map_err(|e| notify_err!("Failed to serialize value: {e}"))?; - let serialized_items = vec![(pli.virtual_id, value_bytes)]; if !file_exists_async(&xtream_path).await { return info_err_res!("BPlusTree file not found for update {}", xtream_path.display()); } + // Prepare encoded payload before opening the writer lock. + let prepared_items = BPlusTreeUpdate::::prepare_upsert_batch(&[(&pli.virtual_id, pli)]) + .map_err(|e| notify_err!("Failed to serialize value: {e}"))?; - // Phase 2: Acquire FileLockManager lock (async, in-process coordination) + // Keep FileLockManager lock for cross-operation coordination (e.g. swap + update). let file_lock = app_config.file_locks.write_lock(&xtream_path).await; - - // Phase 3: Execute all I/O in a single spawn_blocking call + let xtream_path_clone = xtream_path.clone(); tokio::task::spawn_blocking(move || -> Result<(), std::io::Error> { let _guard = file_lock; - let mut tree = BPlusTreeUpdate::::try_new(&xtream_path_clone)?; - tree.upsert_batch_preserialized(serialized_items)?; + let mut tree = BPlusTreeUpdate::::try_new_with_backoff(&xtream_path_clone)?; + tree.upsert_batch_encoded(prepared_items)?; Ok(()) }) .await @@ -176,28 +173,27 @@ pub async fn write_playlist_batch_item_upsert( ensure_xtream_storage_path(&config, target_name).await? }; let xtream_path = xtream_get_file_path(&storage_path, xtream_cluster); - - // Phase 1: Pre-serialize all items in async context (before spawn_blocking) - // This converts references to owned bytes, avoiding clones of the full struct - let mut serialized_items: Vec<(u32, Vec)> = Vec::with_capacity(pli_list.len()); - for pli in pli_list { - let value_bytes = crate::utils::binary_serialize(pli) - .map_err(|e| notify_err!("Failed to serialize value: {e}"))?; - serialized_items.push((pli.virtual_id, value_bytes)); - } - - // Phase 2: Acquire FileLockManager lock (async, in-process coordination) - let file_lock = app_config.file_locks.write_lock(&xtream_path).await; - - // Phase 3: Execute all I/O in a single spawn_blocking call + if !file_exists_async(&xtream_path).await { return info_err_res!("BPlusTree file not found for upsert {}", xtream_path.display()); } + + // Prepare encoded payload before opening the writer lock. + let batch_refs: Vec<(&u32, &XtreamPlaylistItem)> = pli_list + .iter() + .map(|pli| (&pli.virtual_id, pli)) + .collect(); + let prepared_items = BPlusTreeUpdate::::prepare_upsert_batch(&batch_refs) + .map_err(|e| notify_err!("Failed to serialize value: {e}"))?; + + // Keep FileLockManager lock for cross-operation coordination (e.g. swap + update). + let file_lock = app_config.file_locks.write_lock(&xtream_path).await; + let xtream_path_clone = xtream_path.clone(); tokio::task::spawn_blocking(move || -> Result<(), std::io::Error> { let _guard = file_lock; - let mut tree = BPlusTreeUpdate::::try_new(&xtream_path_clone)?; - tree.upsert_batch_preserialized(serialized_items)?; + let mut tree = BPlusTreeUpdate::::try_new_with_backoff(&xtream_path_clone)?; + tree.upsert_batch_encoded(prepared_items)?; Ok(()) }) .await @@ -344,18 +340,31 @@ pub async fn xtream_write_playlist( drop(lock); } - if let Err(err) = write_playlists_to_file( - app_cfg, - &path, - true, - StorageKey::VirtualId, - vec![ - (XtreamCluster::Live, live_col.iter().map(|item| XtreamPlaylistItem::from(&**item)).collect::>()), - (XtreamCluster::Video, vod_col.iter().map(|item| XtreamPlaylistItem::from(&**item)).collect::>()), - (XtreamCluster::Series, series_col.iter().map(|item| XtreamPlaylistItem::from(&**item)).collect::>()), - ], - ).await { - errors.push(format!("Persisting collection failed:{err}")); + // Process each cluster sequentially to avoid holding multiple fully + // materialized Xtream collections in memory at the same time. + for (cluster, col) in [ + (XtreamCluster::Live, &live_col), + (XtreamCluster::Video, &vod_col), + (XtreamCluster::Series, &series_col), + ] { + if col.is_empty() { + continue; + } + let data = col + .iter() + .map(|item| XtreamPlaylistItem::from(&**item)) + .collect::>(); + if let Err(err) = write_playlists_to_file( + app_cfg, + &path, + true, + StorageKey::VirtualId, + vec![(cluster, data)], + ) + .await + { + errors.push(format!("Persisting collection failed:{err}")); + } } if !errors.is_empty() { @@ -632,7 +641,7 @@ pub async fn iter_raw_xtream_input_playlist(app_config: &AppConfig, input: &Conf async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> Option<(FileReadGuard, Box + Send>)> where - ItemKey: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, + ItemKey: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static, SortKey: for<'de> Deserialize<'de> + Send + 'static, { let file_lock = app_config.file_locks.read_lock(xtream_path).await; @@ -849,22 +858,27 @@ async fn persist_input_info(app_config: &Arc, storage_path: &Path, cl input_name: &str, provider_id: u32, props: StreamProperties) -> Result<(), Error> { let xtream_path = xtream_get_file_path(storage_path, cluster); if xtream_path.exists() { - { - let _file_lock = app_config.file_locks.write_lock(&xtream_path).await; - let mut tree: BPlusTreeUpdate = BPlusTreeUpdate::try_new(&xtream_path).map_err(|err| Error::other(format!("failed to open BPlusTree for input {input_name}: {err}")))?; + let file_lock = app_config.file_locks.write_lock(&xtream_path).await; + let xtream_path_clone = xtream_path.clone(); + let input_name_owned = input_name.to_string(); + tokio::task::spawn_blocking(move || -> Result<(), Error> { + let _guard = file_lock; + let mut tree: BPlusTreeUpdate = BPlusTreeUpdate::try_new_with_backoff(&xtream_path_clone) + .map_err(|err| Error::other(format!("failed to open BPlusTree for input {input_name_owned}: {err}")))?; match tree.query(&provider_id) { Ok(Some(mut pli)) => { pli.additional_properties = Some(props); - tree.update(&provider_id, pli).map_err(|err| Error::other(format!("failed to write {cluster} info for input {input_name}: {err}")))?; + tree.update(&provider_id, pli).map_err(|err| Error::other(format!("failed to write {cluster} info for input {input_name_owned}: {err}")))?; } Ok(None) => { - error!("Could not find input entry for provider_id: {provider_id} and input: {input_name}"); + error!("Could not find input entry for provider_id: {provider_id} and input: {input_name_owned}"); } Err(err) => { - error!("Failed to query BPlusTree for provider_id: {provider_id} and input: {input_name}: {err}"); + error!("Failed to query BPlusTree for provider_id: {provider_id} and input: {input_name_owned}: {err}"); } } - } + Ok(()) + }).await.map_err(|err| Error::other(format!("failed to join blocking input info persist for {input_name}: {err}")))??; } Ok(()) } @@ -874,32 +888,38 @@ pub async fn persist_input_info_batch(app_config: &Arc, storage_path: if updates.is_empty() { return Ok(()); } let xtream_path = xtream_get_file_path(storage_path, cluster); if xtream_path.exists() { - let _file_lock = app_config.file_locks.write_lock(&xtream_path).await; - let mut tree: BPlusTreeUpdate = BPlusTreeUpdate::try_new(&xtream_path) - .map_err(|err| Error::other(format!("failed to open BPlusTree for input {input_name}: {err}")))?; + let file_lock = app_config.file_locks.write_lock(&xtream_path).await; + let xtream_path_clone = xtream_path.clone(); + let input_name_owned = input_name.to_string(); + tokio::task::spawn_blocking(move || -> Result<(), Error> { + let _guard = file_lock; + let mut tree: BPlusTreeUpdate = BPlusTreeUpdate::try_new_with_backoff(&xtream_path_clone) + .map_err(|err| Error::other(format!("failed to open BPlusTree for input {input_name_owned}: {err}")))?; - let mut updated_plis = Vec::with_capacity(updates.len()); - for (provider_id, props) in updates { - match tree.query(&provider_id) { - Ok(Some(mut pli)) => { - pli.additional_properties = Some(props); - updated_plis.push((provider_id, pli)); - } - Ok(None) => { - error!("Could not find input entry for provider_id: {provider_id} and input: {input_name}"); - } - Err(err) => { - error!("Failed to query BPlusTree for provider_id: {provider_id} and input: {input_name}: {err}"); + let mut updated_plis = Vec::with_capacity(updates.len()); + for (provider_id, props) in updates { + match tree.query(&provider_id) { + Ok(Some(mut pli)) => { + pli.additional_properties = Some(props); + updated_plis.push((provider_id, pli)); + } + Ok(None) => { + error!("Could not find input entry for provider_id: {provider_id} and input: {input_name_owned}"); + } + Err(err) => { + error!("Failed to query BPlusTree for provider_id: {provider_id} and input: {input_name_owned}: {err}"); + } } } - } - if !updated_plis.is_empty() { - let refs: Vec<(&u32, &XtreamPlaylistItem)> = updated_plis.iter() - .map(|(id, pli)| (id, pli)) - .collect(); - tree.update_batch(&refs).map_err(|err| Error::other(format!("failed to write batch {cluster} info for input {input_name}: {err}")))?; - } + if !updated_plis.is_empty() { + let refs: Vec<(&u32, &XtreamPlaylistItem)> = updated_plis.iter() + .map(|(id, pli)| (id, pli)) + .collect(); + tree.update_batch(&refs).map_err(|err| Error::other(format!("failed to write batch {cluster} info for input {input_name_owned}: {err}")))?; + } + Ok(()) + }).await.map_err(|err| Error::other(format!("failed to join blocking input info batch persist for {input_name}: {err}")))??; } Ok(()) } @@ -983,4 +1003,3 @@ pub async fn load_input_xtream_playlist(app_config: &Arc, storage_pat Ok(groups.into_values().collect()) } - diff --git a/backend/src/utils/network/xtream.rs b/backend/src/utils/network/xtream.rs index 11540f927..f64d65cfb 100644 --- a/backend/src/utils/network/xtream.rs +++ b/backend/src/utils/network/xtream.rs @@ -5,7 +5,7 @@ use crate::model::{is_input_expired, xtream_mapping_option_from_target_options, use crate::model::{InputSource, ProxyUserCredentials}; use crate::processing::parser::xtream; use crate::processing::parser::xtream::parse_xtream_series_info; -use crate::repository::BPlusTreeUpdate; +use crate::repository::{BPlusTreeUpdate, FlushPolicy}; use crate::repository::VirtualIdRecord; use crate::repository::{ensure_input_storage_path, get_input_storage_path, get_target_storage_path}; use crate::repository::{get_live_cat_collection_path, get_series_cat_collection_path, get_vod_cat_collection_path, xtream_get_file_path, CategoryEntry}; @@ -364,11 +364,9 @@ pub async fn download_xtream_playlist(app_config: &Arc, client: &reqw } }; - let base_url = if is_provider_url { - base_input_url.clone() - } else { - get_xtream_stream_url_base(&base_input_url, username, password) - }; + // Build a canonical Xtream player_api URL for both plain URLs and provider:// URLs. + // For provider URLs, this keeps the scheme unresolved while still adding the required path/query. + let base_url = get_xtream_stream_url_base(&base_input_url, username, password); let input_source_login = input_source.with_url(base_url.clone()); check_alias_user_state(app_config, client, input).await; @@ -518,6 +516,7 @@ pub fn create_vod_info_from_item(target: &ConfigTarget, user: &ProxyUserCredenti const BATCH_SIZE: usize = 1000; +#[allow(clippy::too_many_lines)] async fn process_xtream_cluster_to_disk( app_config: &Arc, input: &ConfigInput, @@ -568,11 +567,12 @@ async fn process_xtream_cluster_to_disk( notify_err!("Init tree error {e}") })?; - let mut tree = BPlusTreeUpdate::try_new(&tmp_xtream_path) + let mut tree: BPlusTreeUpdate = BPlusTreeUpdate::try_new_with_backoff(&tmp_xtream_path) .map_err(|e| { error!("Failed to open ghost tree at {}: {e}", tmp_xtream_path.display()); notify_err!("Failed to open tree {e}") })?; + tree.set_flush_policy(FlushPolicy::Batch); let mut buffer = Vec::with_capacity(BATCH_SIZE); @@ -581,7 +581,11 @@ async fn process_xtream_cluster_to_disk( buffer.push(item); if buffer.len() >= BATCH_SIZE { let batch: Vec<(&u32, &XtreamPlaylistItem)> = buffer.iter().map(|i| (&i.provider_id, i)).collect(); - tree.upsert_batch(&batch).map_err(|e| { + let prepared = BPlusTreeUpdate::::prepare_upsert_batch(&batch).map_err(|e| { + error!("Batch prepare failed for cluster {cluster}: {e}"); + notify_err!("Prepare failed {e}") + })?; + tree.upsert_batch_encoded(prepared).map_err(|e| { error!("Batch upsert failed for cluster {cluster}: {e}"); notify_err!("Upsert failed {e}") })?; @@ -592,11 +596,20 @@ async fn process_xtream_cluster_to_disk( // Final batch processing if !buffer.is_empty() { let batch: Vec<(&u32, &XtreamPlaylistItem)> = buffer.iter().map(|i| (&i.provider_id, i)).collect(); - tree.upsert_batch(&batch).map_err(|e| { + let prepared = BPlusTreeUpdate::::prepare_upsert_batch(&batch).map_err(|e| { + error!("Final batch prepare failed for cluster {cluster}: {e}"); + notify_err!("Prepare failed {e}") + })?; + tree.upsert_batch_encoded(prepared).map_err(|e| { error!("Final batch upsert failed for cluster {cluster}: {e}"); notify_err!("Upsert failed {e}") })?; } + + tree.commit().map_err(|e| { + error!("Commit failed for cluster {cluster}: {e}"); + notify_err!("Commit failed {e}") + })?; Ok::<(), TuliproxError>(()) }); @@ -627,11 +640,22 @@ async fn process_xtream_cluster_to_disk( // Acquire file lock to ensure atomic operations let swap_lock = app_config.file_locks.write_lock(&xtream_path).await; - // Optional compaction to optimize the newly created database file - if let Ok(mut tree_update) = BPlusTreeUpdate::::try_new(&tmp_xtream_path) { - if let Err(e) = tree_update.compact(&tmp_xtream_path) { + // Optional compaction to optimize the newly created database file. + // Run on blocking pool because B+Tree lock backoff uses std::thread::sleep. + let compact_path = tmp_xtream_path.clone(); + match tokio::task::spawn_blocking(move || -> Result<(), std::io::Error> { + if let Ok(mut tree_update) = BPlusTreeUpdate::::try_new_with_backoff(&compact_path) { + tree_update.compact(&compact_path)?; + } + Ok(()) + }).await { + Ok(Ok(())) => {} + Ok(Err(e)) => { error!("Failed to compact temporary database for {cluster} at {:?}: {e}", tmp_xtream_path.display()); } + Err(e) => { + error!("Compaction task join failed for {cluster} at {:?}: {e}", tmp_xtream_path.display()); + } } // Atomic Swap: Replace old database with new one diff --git a/backend/src/utils/sys_utils.rs b/backend/src/utils/sys_utils.rs index a370d5a88..09bd371c8 100644 --- a/backend/src/utils/sys_utils.rs +++ b/backend/src/utils/sys_utils.rs @@ -1,3 +1,8 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use log::{log_enabled, trace, Level}; +use sysinfo::{MemoryRefreshKind, Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System}; +use shared::utils::human_readable_byte_size; + #[macro_export] macro_rules! exit { ($($arg:tt)*) => {{ @@ -5,4 +10,55 @@ macro_rules! exit { std::process::exit(1); }}; } + pub use exit; + +static PROCESS_PEAK_RSS_BYTES: AtomicU64 = AtomicU64::new(0); +static PROCESS_PEAK_VMEM_BYTES: AtomicU64 = AtomicU64::new(0); + +// fn trim_allocator_after_update() { +// #[cfg(all(target_os = "linux", target_env = "gnu"))] +// { +// log_memory_snapshot("exec_processing before_malloc_trim"); +// // SAFETY: libc::malloc_trim is thread-safe for the process allocator +// // and does not require any additional invariants on the Rust side. +// let result = unsafe { libc::malloc_trim(0) }; +// debug!("allocator[malloc_trim] result={result}"); +// log_memory_snapshot("exec_processing after_malloc_trim"); +// } +// } + +pub fn log_memory_snapshot(stage: &str) { + if !log_enabled!(Level::Trace) { + return; + } + + let pid = Pid::from_u32(std::process::id()); + let refresh_kind = RefreshKind::nothing() + .with_memory(MemoryRefreshKind::nothing().with_ram()) + .with_processes(ProcessRefreshKind::nothing().with_memory()); + let mut sys = System::new_with_specifics(refresh_kind); + sys.refresh_processes(ProcessesToUpdate::Some(&[pid]), true); + sys.refresh_memory(); + + if let Some(process) = sys.processes().get(&pid) { + let rss_bytes = process.memory(); + let vmem_bytes = process.virtual_memory(); + let peak_rss = PROCESS_PEAK_RSS_BYTES.fetch_max(rss_bytes, Ordering::Relaxed).max(rss_bytes); + let peak_vmem = PROCESS_PEAK_VMEM_BYTES.fetch_max(vmem_bytes, Ordering::Relaxed).max(vmem_bytes); + + trace!( + "memory[{stage}] rss={} ({}) vmem={} ({}) peak_rss={} ({}) peak_vmem={} ({})", + rss_bytes, + human_readable_byte_size(rss_bytes), + vmem_bytes, + human_readable_byte_size(vmem_bytes), + peak_rss, + human_readable_byte_size(peak_rss), + peak_vmem, + human_readable_byte_size(peak_vmem), + ); + } else { + trace!("memory[{stage}] process info unavailable"); + } +} \ No newline at end of file diff --git a/shared/src/model/config/input.rs b/shared/src/model/config/input.rs index 179c34e21..b9fb5922f 100644 --- a/shared/src/model/config/input.rs +++ b/shared/src/model/config/input.rs @@ -385,7 +385,33 @@ impl ConfigInputDto { self.persist = get_trimmed_string(self.persist.as_deref()); if self.url.starts_with(PROVIDER_SCHEME_PREFIX) { - if let Ok((host, _path)) = parse_provider_scheme_url_parts(&self.url) { + let (host, _path) = match parse_provider_scheme_url_parts(&self.url) { + Ok(parts) => parts, + Err(err) => { + return info_err_res!( + "Malformed provider URL {}: {}", + sanitize_sensitive_info(&self.url), + sanitize_sensitive_info(err.to_string().as_str()) + ); + } + }; + if !provider_names.contains(host) { + return info_err_res!("Provider name {host} is not defined"); + } + } + + if let Some(staged_input) = self.staged.as_ref() { + if staged_input.url.starts_with(PROVIDER_SCHEME_PREFIX) { + let (host, _path) = match parse_provider_scheme_url_parts(&staged_input.url) { + Ok(parts) => parts, + Err(err) => { + return info_err_res!( + "Malformed provider URL {}: {}", + sanitize_sensitive_info(&staged_input.url), + sanitize_sensitive_info(err.to_string().as_str()) + ); + } + }; if !provider_names.contains(host) { return info_err_res!("Provider name {host} is not defined"); } @@ -399,10 +425,18 @@ impl ConfigInputDto { for alias in aliases { current_index = alias.prepare(current_index, input_type)?; if alias.url.starts_with(PROVIDER_SCHEME_PREFIX) { - if let Ok((host, _path)) = parse_provider_scheme_url_parts(&alias.url) { - if !provider_names.contains(host) { - return info_err_res!("Provider name {host} is not defined"); + let (host, _path) = match parse_provider_scheme_url_parts(&alias.url) { + Ok(parts) => parts, + Err(err) => { + return info_err_res!( + "Malformed provider URL {}: {}", + sanitize_sensitive_info(&alias.url), + sanitize_sensitive_info(err.to_string().as_str()) + ); } + }; + if !provider_names.contains(host) { + return info_err_res!("Provider name {host} is not defined"); } } } @@ -567,4 +601,4 @@ impl ConfigProviderDto { } Ok(()) } -} \ No newline at end of file +}