From 5b49ea58d8f7f9905ab9741d8e4828f88324a969 Mon Sep 17 00:00:00 2001 From: euzu <33094714+euzu@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:56:29 +0200 Subject: [PATCH] BPlusTree refactoring / optimization (#792) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * BPlusTree refactoring / optimization New Features - Added streaming JSON/CBOR “try” array responses that stop at the first failure, with coalesced HTTP chunking for smoother delivery. - Introduced B+Tree v3 persistence (with v2→v3 migration) including verified atomic updates, snapshots, compaction, and batch operations. - Updated playlist/index streaming to use sorted-index acceleration with corruption-safe fallback. Bug Fixes - Corrupt/unreadable Library databases now surface repository errors instead of appearing empty. - Playlist/lineup/EPG/Xtream/M3U streaming now propagates per-entry failures correctly (including improved counting/early-stop behavior). Documentation - Updated README and changelog to reflect the B+Tree v3 format and migration behavior. --- .gitignore | 1 + CHANGELOG.md | 15 + README.md | 19 +- backend/src/api/api_utils.rs | 171 +- .../src/api/endpoints/api_playlist_utils.rs | 40 +- backend/src/api/endpoints/hdhomerun_api.rs | 21 +- backend/src/api/endpoints/m3u_api.rs | 11 +- backend/src/api/endpoints/xmltv_api.rs | 79 +- backend/src/api/endpoints/xtream_api.rs | 20 +- .../src/api/model/hls_cache/manifest_fetch.rs | 30 +- backend/src/api/model/hls_cache/refresh.rs | 6 +- .../api/model/hls_cache/segment_fetcher.rs | 5 +- .../api/model/hls_cache/transient_fetcher.rs | 6 +- .../src/api/model/metadata_update_manager.rs | 3 +- .../model/streams/provider_stream_factory.rs | 8 +- backend/src/repository/bplustree.rs | 8642 ----------------- backend/src/repository/bplustree/common.rs | 103 + .../migration.rs} | 522 +- backend/src/repository/bplustree/mod.rs | 11 + .../src/repository/bplustree/sorted_index.rs | 477 + .../stress.rs} | 216 +- backend/src/repository/bplustree/v2.rs | 1891 ++++ backend/src/repository/bplustree/v3/format.rs | 1524 +++ .../src/repository/bplustree/v3/migration.rs | 409 + backend/src/repository/bplustree/v3/mod.rs | 140 + backend/src/repository/bplustree/v3/page.rs | 689 ++ backend/src/repository/bplustree/v3/tree.rs | 5720 +++++++++++ backend/src/repository/bplustree/v3/wal.rs | 1992 ++++ backend/src/repository/epg_repository.rs | 8 +- backend/src/repository/library_repository.rs | 54 +- .../src/repository/m3u_playlist_iterator.rs | 46 +- backend/src/repository/m3u_repository.rs | 43 +- backend/src/repository/mod.rs | 7 - backend/src/repository/playlist_repository.rs | 8 +- backend/src/repository/playlist_source.rs | 160 +- backend/src/repository/playlist_stream.rs | 124 +- .../src/repository/qos_snapshot_repository.rs | 16 +- backend/src/repository/sorted_index.rs | 788 -- backend/src/repository/target_id_mapping.rs | 14 +- backend/src/repository/user_repository.rs | 2 +- .../repository/xtream_playlist_iterator.rs | 76 +- backend/src/repository/xtream_repository.rs | 57 +- backend/src/utils/db_viewer.rs | 10 +- backend/src/utils/network/content_coding.rs | 4 +- backend/src/utils/network/request.rs | 5 +- 45 files changed, 14312 insertions(+), 9881 deletions(-) delete mode 100644 backend/src/repository/bplustree.rs create mode 100644 backend/src/repository/bplustree/common.rs rename backend/src/repository/{bplustree_migration.rs => bplustree/migration.rs} (72%) create mode 100644 backend/src/repository/bplustree/mod.rs create mode 100644 backend/src/repository/bplustree/sorted_index.rs rename backend/src/repository/{bplustree_stress.rs => bplustree/stress.rs} (73%) create mode 100644 backend/src/repository/bplustree/v2.rs create mode 100644 backend/src/repository/bplustree/v3/format.rs create mode 100644 backend/src/repository/bplustree/v3/migration.rs create mode 100644 backend/src/repository/bplustree/v3/mod.rs create mode 100644 backend/src/repository/bplustree/v3/page.rs create mode 100644 backend/src/repository/bplustree/v3/tree.rs create mode 100644 backend/src/repository/bplustree/v3/wal.rs delete mode 100644 backend/src/repository/sorted_index.rs diff --git a/.gitignore b/.gitignore index df0e3dfd5..c71d447bd 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .tools +.worktrees dist /resources/*.ts /release diff --git a/CHANGELOG.md b/CHANGELOG.md index e9cfd71be..c1d9d8484 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -570,6 +570,21 @@ ## 🛠 Maintenance +- **B+Tree v3 persistence engine**: + - Consolidated the facade, v2 compatibility reader, v3 engine, migration, WAL, sorted index, and stress tests under + `backend/src/repository/bplustree/`. + - Added checksummed 4 KiB Slotted Pages, WAL-before-data in-place updates, verified atomic full replacement, typed + v1/v2 startup migration, identity-bound sorted indexes, and corruption-reporting iterators. + - Reused mmap mappings, page validation, and decoded internal routes across cheap `BPlusTreeQuery` clones while + retaining request-local scratch buffers. + - Kept stored values up to 512 bytes inline, avoiding one mostly empty 4 KiB overflow page per typical Xtream/M3U + playlist entry and restoring compact full-scan behavior. + - Playlist APIs now coalesce small M3U, Xtream, HDHomeRun, XMLTV, JSON, and CBOR fragments into bounded 64 KiB + response chunks instead of emitting one HTTP body frame per entry. + - Corrupt individual B+Tree values are logged and skipped when the iterator can safely continue; database-open and + worker failures remain visible, and failed input-cache opens no longer replace existing Xtream persistence. + - Classified B+Tree read failures as repository errors + - Opening a corrupt existing Library database no longer silently produces an empty Library. - **Shared `FieldWrapper` For Form Inputs**: - Extracted the repeated label / field-id / `tp__input-wrapper` scaffolding from the `Input`, `NumberInput`, and `TextArea` primitives into a single shared `FieldWrapper` component, reducing duplication while keeping the diff --git a/README.md b/README.md index 1607ab46f..2b47c3e19 100644 --- a/README.md +++ b/README.md @@ -29,17 +29,16 @@ See [`LICENSE`](https://github.com/euzu/tuliprox/blob/develop/LICENSE). ### 2. Custom B+Tree Storage Engine — No External Database Needed -- Purpose-built B+Tree with Slotted Page architecture -- Adaptive LZ4 compression for minimal disk footprint -- Zero-copy scans at up to 96,000 ops/sec -- Batch upsert for massive throughput during playlist updates -- Atomic I/O with file locking — no corrupt data, ever -- Configurable flush policy (Immediate, Batch, None) +- Versioned 4 KiB Slotted Page format with checksummed headers, pages, cells, and overflow chains +- Adaptive LZ4 compression and overflow pages for compact variable-size values +- Mmap-backed streaming scans without loading the complete database into memory +- Shared query snapshots reuse mmap mappings, validated pages, and decoded internal routes +- WAL-protected in-place updates with `Immediate` and `Batch` flush policies +- Fully built replacements and compaction results are verified before atomic publication +- Typed startup migration upgrades recognized legacy v1/v2 databases to v3 +- Identity- and generation-bound sorted indexes accelerate playlist-order streaming - String interning (`Arc`) for playlist entries reduces memory footprint -- B+Tree compaction to reclaim disk space -- Persistent value caching with thread-safe access -- Packed block update optimization — direct disk writes for same-size updates, bypassing expensive - read-scan-modify-write cycles +- Page-local updates, tombstone reuse, free-page reuse, and explicit compaction reclaim disk space ### 3. Four Output Formats — One Tool to Rule Them All diff --git a/backend/src/api/api_utils.rs b/backend/src/api/api_utils.rs index a4b81d1d3..9a52c7789 100644 --- a/backend/src/api/api_utils.rs +++ b/backend/src/api/api_utils.rs @@ -35,7 +35,7 @@ use axum::{ http::{header, Extensions, HeaderMap, HeaderName, HeaderValue, Response, StatusCode}, response::IntoResponse, }; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use chrono::{DateTime, Utc}; use futures::{stream, Stream, StreamExt, TryStreamExt}; use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation}; @@ -3795,6 +3795,21 @@ where stream_json_array_stream(data) } +pub fn stream_json_or_bin_response_try_stream( + accept: Option<&str>, + data: S, +) -> axum::response::Response +where + P: serde::Serialize + Send + 'static, + S: Stream> + Send + Unpin + 'static, + E: std::fmt::Display + Send + 'static, +{ + if accept.is_some_and(|value| value.contains(CONTENT_TYPE_CBOR)) { + return stream_bin_array_try_stream(data); + } + stream_json_array_try_stream(data) +} + pub fn create_session_fingerprint( fingerprint: &Fingerprint, username: &str, @@ -3900,11 +3915,11 @@ where } }); - let body = Body::from_stream( + let body = Body::from_stream(coalesce_byte_stream( stream::once(async { Ok::<_, Infallible>(Bytes::from_static(b"[")) }) .chain(stream) .chain(stream::once(async { Ok::<_, Infallible>(Bytes::from_static(b"]")) })), - ); + )); try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_JSON).body(body)) } @@ -3928,7 +3943,7 @@ where } }); - let body = Body::from_stream( + let body = Body::from_stream(coalesce_byte_stream( stream::once(async { // CBOR: start indefinite-length array Ok::<_, Infallible>(Bytes::from_static(&[0x9f])) @@ -3938,7 +3953,7 @@ where // CBOR: end indefinite-length array Ok::<_, Infallible>(Bytes::from_static(&[0xff])) })), - ); + )); try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_CBOR).body(body)) } @@ -3963,15 +3978,54 @@ where } }); - let body = Body::from_stream( + let body = Body::from_stream(coalesce_byte_stream( stream::once(async { Ok::<_, Infallible>(Bytes::from_static(b"[")) }) .chain(stream) .chain(stream::once(async { Ok::<_, Infallible>(Bytes::from_static(b"]")) })), - ); + )); try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_JSON).body(body)) } +fn stream_json_array_try_stream(stream: S) -> axum::response::Response +where + P: serde::Serialize + Send + 'static, + S: Stream> + Send + Unpin + 'static, + E: std::fmt::Display + Send + 'static, +{ + let stream = stream::unfold((stream, true, false), |(mut stream, first, failed)| async move { + if failed { + return None; + } + match stream.next().await { + Some(Ok(item)) => { + let serialized = serde_json::to_vec(&item).map_err(|error| error.to_string()); + let bytes = serialized.map(|serialized| { + if first { + Bytes::from(serialized) + } else { + let mut framed = Vec::with_capacity(serialized.len() + 1); + framed.push(b','); + framed.extend_from_slice(&serialized); + Bytes::from(framed) + } + }); + let failed = bytes.is_err(); + Some((bytes, (stream, false, failed))) + } + Some(Err(error)) => Some((Err(error.to_string()), (stream, first, true))), + None => None, + } + }); + + let body = Body::from_stream(coalesce_byte_stream( + stream::once(async { Ok::<_, String>(Bytes::from_static(b"[")) }) + .chain(stream) + .chain(stream::once(async { Ok::<_, String>(Bytes::from_static(b"]")) })), + )); + try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_JSON).body(body)) +} + pub fn stream_bin_array_stream(stream: S) -> axum::response::Response where P: serde::Serialize + Send + 'static, @@ -3990,15 +4044,85 @@ where } }); - let body = Body::from_stream( + let body = Body::from_stream(coalesce_byte_stream( stream::once(async { Ok::<_, Infallible>(Bytes::from_static(&[0x9f])) }) .chain(stream) .chain(stream::once(async { Ok::<_, Infallible>(Bytes::from_static(&[0xff])) })), - ); + )); try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_CBOR).body(body)) } +fn stream_bin_array_try_stream(stream: S) -> axum::response::Response +where + P: serde::Serialize + Send + 'static, + S: Stream> + Send + Unpin + 'static, + E: std::fmt::Display + Send + 'static, +{ + let stream = stream::unfold((stream, false), |(mut stream, failed)| async move { + if failed { + return None; + } + match stream.next().await { + Some(Ok(item)) => { + let bytes = bin_serialize(&item).map(Bytes::from).map_err(|error| error.to_string()); + let failed = bytes.is_err(); + Some((bytes, (stream, failed))) + } + Some(Err(error)) => Some((Err(error.to_string()), (stream, true))), + None => None, + } + }); + + let body = Body::from_stream(coalesce_byte_stream( + stream::once(async { Ok::<_, String>(Bytes::from_static(&[0x9f])) }) + .chain(stream) + .chain(stream::once(async { Ok::<_, String>(Bytes::from_static(&[0xff])) })), + )); + try_unwrap_body!(Response::builder().header(header::CONTENT_TYPE, CONTENT_TYPE_CBOR).body(body)) +} + +const API_STREAM_CHUNK_SIZE: usize = 64 * 1024; + +pub(crate) fn coalesce_byte_stream(stream: S) -> impl Stream> +where + S: Stream> + Send + 'static, + E: Send + 'static, +{ + stream::unfold((Box::pin(stream), None, false), |(mut stream, pending_error, finished)| async move { + if let Some(error) = pending_error { + return Some((Err(error), (stream, None, true))); + } + if finished { + return None; + } + + let mut chunk = BytesMut::with_capacity(API_STREAM_CHUNK_SIZE); + loop { + match stream.next().await { + Some(Ok(bytes)) if chunk.is_empty() && bytes.len() >= API_STREAM_CHUNK_SIZE => { + return Some((Ok(bytes), (stream, None, false))); + } + Some(Ok(bytes)) => { + chunk.extend_from_slice(&bytes); + if chunk.len() >= API_STREAM_CHUNK_SIZE { + return Some((Ok(chunk.freeze()), (stream, None, false))); + } + } + Some(Err(error)) if chunk.is_empty() => { + return Some((Err(error), (stream, None, true))); + } + Some(Err(error)) => { + return Some((Ok(chunk.freeze()), (stream, Some(error), false))); + } + None if chunk.is_empty() => return None, + None => return Some((Ok(chunk.freeze()), (stream, None, true))), + } + } + }) + .fuse() +} + pub fn create_api_proxy_user(app_state: &Arc) -> ProxyUserCredentials { let config = app_state.app_config.config.load(); @@ -4087,6 +4211,35 @@ mod tests { test_runtime_provider_with_type(url, username, password, InputType::Xtream) } + #[tokio::test] + async fn streamed_json_array_coalesces_small_entries() { + let response = stream_json_array_stream(stream::iter(0..4_096u32)); + let mut body = response.into_body(); + let mut frames = 0usize; + let mut bytes = Vec::new(); + while let Some(frame) = body.frame().await { + let Ok(frame) = frame else { + return; + }; + if let Ok(data) = frame.into_data() { + frames += 1; + bytes.extend_from_slice(&data); + } + } + assert!(frames <= 2, "small JSON entries should be coalesced, got {frames} frames"); + let decoded = serde_json::from_slice::>(&bytes); + assert!(decoded.is_ok_and(|values| values.len() == 4_096)); + } + + #[tokio::test] + async fn coalesced_stream_remains_finished_when_polled_again() { + let stream = coalesce_byte_stream(stream::empty::>()); + futures::pin_mut!(stream); + + assert!(stream.next().await.is_none()); + assert!(stream.next().await.is_none()); + } + fn test_runtime_provider_with_type( url: &str, username: &str, diff --git a/backend/src/api/endpoints/api_playlist_utils.rs b/backend/src/api/endpoints/api_playlist_utils.rs index 5d1e2fe46..e4096938b 100644 --- a/backend/src/api/endpoints/api_playlist_utils.rs +++ b/backend/src/api/endpoints/api_playlist_utils.rs @@ -1,6 +1,8 @@ use crate::api::model::AppState; use crate::{ - api::api_utils::{empty_json_list_response, stream_json_or_bin_response_stream}, + api::api_utils::{ + empty_json_list_response, stream_json_or_bin_response_stream, stream_json_or_bin_response_try_stream, + }, model::{ConfigInput, ConfigTarget}, repository::{ iter_raw_m3u_input_playlist, iter_raw_m3u_target_playlist, iter_raw_xtream_input_playlist, @@ -9,7 +11,6 @@ use crate::{ utils::{m3u, xtream}, }; use axum::response::IntoResponse; -use log::warn; use serde_json::json; use shared::utils::{concat_path, concat_path_leading_slash, obfuscate_text, Internable}; use shared::model::{ @@ -44,8 +45,16 @@ pub(in crate::api::endpoints) async fn get_playlist_for_target( } else { |_pli: &XtreamPlaylistItem| true }; - let converted_stream = channel_iterator.filter(item_filter).map(UiPlaylistItem::from).map(move |uiu| rewrite_resource_url(&encrypt_secret, &resource_url, uiu)); - return stream_json_or_bin_response_stream(accept, converted_stream).into_response(); + let converted_stream = channel_iterator.filter_map(move |entry| match entry { + Ok(item) if item_filter(&item) => Some(Ok(rewrite_resource_url( + &encrypt_secret, + &resource_url, + UiPlaylistItem::from(item), + ))), + Ok(_) => None, + Err(error) => Some(Err(error)), + }); + return stream_json_or_bin_response_try_stream(accept, converted_stream).into_response(); } else if target.has_output(TargetType::M3u) { let Some(channel_iterator) = iter_raw_m3u_target_playlist(&app_state.app_config, target, Some(cluster)).await else { return empty_json_list_response(); @@ -61,17 +70,14 @@ pub(in crate::api::endpoints) async fn get_playlist_for_target( let converted_stream = channel_iterator.filter_map(move |res| match res { Ok(pli) => { if item_filter(&pli) { - Some(rewrite_resource_url(&encrypt_secret, &resource_url, UiPlaylistItem::from(pli))) + Some(Ok(rewrite_resource_url(&encrypt_secret, &resource_url, UiPlaylistItem::from(pli)))) } else { None } } - Err(err) => { - warn!("Skipping unreadable M3U target playlist entry: {err}"); - None - } + Err(error) => Some(Err(error)), }); - return stream_json_or_bin_response_stream(accept, converted_stream).into_response(); + return stream_json_or_bin_response_try_stream(accept, converted_stream).into_response(); } } (axum::http::StatusCode::BAD_REQUEST, axum::Json(json!({"error": "Invalid Arguments"}))).into_response() @@ -100,20 +106,14 @@ pub(in crate::api::endpoints) async fn get_playlist_for_input( let Some(channel_iterator) = iter_raw_xtream_input_playlist(&app_state.app_config, input, cluster).await else { return empty_json_list_response(); }; - let converted_stream = channel_iterator.map(UiPlaylistItem::from); - return stream_json_or_bin_response_stream(accept, converted_stream).into_response(); + let converted_stream = channel_iterator.map(|entry| entry.map(UiPlaylistItem::from)); + return stream_json_or_bin_response_try_stream(accept, converted_stream).into_response(); } else if input.input_type.is_m3u() { let Some(channels) = iter_raw_m3u_input_playlist(&app_state.app_config, input, Some(cluster)).await else { return empty_json_list_response(); }; - let converted_stream = channels.filter_map(|res| match res { - Ok(pli) => Some(UiPlaylistItem::from(pli)), - Err(err) => { - warn!("Skipping unreadable M3U input playlist entry: {err}"); - None - } - }); - return stream_json_or_bin_response_stream(accept, converted_stream).into_response(); + let converted_stream = channels.map(|entry| entry.map(UiPlaylistItem::from)); + return stream_json_or_bin_response_try_stream(accept, converted_stream).into_response(); } } (axum::http::StatusCode::BAD_REQUEST, axum::Json(json!({"error": "Invalid Arguments"}))).into_response() diff --git a/backend/src/api/endpoints/hdhomerun_api.rs b/backend/src/api/endpoints/hdhomerun_api.rs index 5de818cc9..de5042f43 100644 --- a/backend/src/api/endpoints/hdhomerun_api.rs +++ b/backend/src/api/endpoints/hdhomerun_api.rs @@ -1,6 +1,6 @@ use crate::{ api::{ - api_utils::{internal_server_error, try_unwrap_body}, + api_utils::{coalesce_byte_stream, internal_server_error, try_unwrap_body}, model::HdHomerunAppState, }, auth::{try_check_network_access_only, AuthBasic, Fingerprint}, @@ -16,6 +16,7 @@ use log::{error, warn}; use serde::{Deserialize, Serialize}; use serde_json::json; use shared::{ + error::TuliproxError, model::{M3uPlaylistItem, PlaylistItemType, TargetType, XtreamCluster, XtreamPlaylistItem}, utils::concat_path, }; @@ -115,11 +116,12 @@ fn xtream_item_to_lineup_stream( channels: Option, ) -> impl Stream> where - I: Stream + Send + Unpin + 'static, + I: Stream> + Send + Unpin + 'static, { match channels { Some(chans) => { - let mapped = chans.map(move |(item, has_next)| { + let mapped = chans.map(move |entry| { + let (item, has_next) = entry.map_err(|error| error.to_string())?; let input = cfg.get_input_by_name(&item.input_name); let (live_stream_use_prefix, live_stream_without_extension) = input.as_ref().map_or((true, false), |i| { @@ -159,11 +161,12 @@ where fn m3u_item_to_lineup_stream(channels: Option) -> impl Stream> where - I: Stream + Send + Unpin + 'static, + I: Stream> + Send + Unpin + 'static, { match channels { Some(chans) => { - let mapped = chans.map(move |(item, has_next)| { + let mapped = chans.map(move |entry| { + let (item, has_next) = entry.map_err(|error| error.to_string())?; let lineup = Lineup { guide_number: item.epg_channel_id.clone().unwrap_or(item.name.clone()), guide_name: item.title.clone(), @@ -284,7 +287,7 @@ async fn lineup_status( let cred = Arc::clone(&credentials); let live = if cred.allows_cluster(XtreamCluster::Live) { match XtreamPlaylistIterator::new(XtreamCluster::Live, &cfg, &target, None, &cred).await { - Ok(stream) => stream.count().await, + Ok(stream) => stream.filter_map(|entry| async move { entry.ok() }).count().await, Err(_) => 0, } } else { @@ -292,7 +295,7 @@ async fn lineup_status( }; let vod = if cred.allows_cluster(XtreamCluster::Video) { match XtreamPlaylistIterator::new(XtreamCluster::Video, &cfg, &target, None, &cred).await { - Ok(stream) => stream.count().await, + Ok(stream) => stream.filter_map(|entry| async move { entry.ok() }).count().await, Err(_) => 0, } } else { @@ -370,7 +373,7 @@ async fn lineup( return try_unwrap_body!(axum::response::Response::builder() .status(axum::http::StatusCode::OK) .header(axum::http::header::CONTENT_TYPE, mime::APPLICATION_JSON.to_string()) - .body(axum::body::Body::from_stream(body_stream))); + .body(axum::body::Body::from_stream(coalesce_byte_stream(body_stream)))); } else if (use_all || use_xtream) && target.has_output(TargetType::Xtream) { let Some(server_info) = app_state.app_state.app_config.get_user_server_info(credentials) else { return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response(); @@ -436,7 +439,7 @@ async fn lineup( return try_unwrap_body!(axum::response::Response::builder() .status(axum::http::StatusCode::OK) .header(axum::http::header::CONTENT_TYPE, mime::APPLICATION_JSON.to_string()) - .body(axum::body::Body::from_stream(body_stream))); + .body(axum::body::Body::from_stream(coalesce_byte_stream(body_stream)))); } axum::http::StatusCode::NOT_FOUND.into_response() } diff --git a/backend/src/api/endpoints/m3u_api.rs b/backend/src/api/endpoints/m3u_api.rs index f19f5d062..9f120689f 100644 --- a/backend/src/api/endpoints/m3u_api.rs +++ b/backend/src/api/endpoints/m3u_api.rs @@ -2,6 +2,7 @@ use crate::{ api::{ api_utils::{ admission_failure_response, create_m3u_catchup_session_key, + coalesce_byte_stream, create_playback_session_fingerprint, create_session_fingerprint, force_provider_stream_response, get_session_reservation_ttl_secs, get_user_target, get_user_target_by_credentials, is_seek_request, is_session_based_playback, is_stream_share_enabled, local_stream_response, redirect, redirect_response, resource_response, @@ -48,9 +49,11 @@ async fn m3u_api( match m3u_load_rewrite_playlist(&app_state.app_config, &target, &user).await { Ok(m3u_iter) => { - let content_stream = m3u_iter.map(|mut line| { - line.push('\n'); - Ok::(Bytes::from(line)) + let content_stream = m3u_iter.map(|line| { + line.map(|mut line| { + line.push('\n'); + Bytes::from(line) + }) }); let mut builder = axum::response::Response::builder() @@ -60,7 +63,7 @@ async fn m3u_api( builder = builder.header(axum::http::header::CONTENT_DISPOSITION, "attachment; filename=\"playlist.m3u\""); } - try_unwrap_body!(builder.body(axum::body::Body::from_stream(content_stream))) + try_unwrap_body!(builder.body(axum::body::Body::from_stream(coalesce_byte_stream(content_stream)))) } Err(err) => { error!("{}", sanitize_sensitive_info(&err.to_string())); diff --git a/backend/src/api/endpoints/xmltv_api.rs b/backend/src/api/endpoints/xmltv_api.rs index eb9fbb043..174a4bab6 100644 --- a/backend/src/api/endpoints/xmltv_api.rs +++ b/backend/src/api/endpoints/xmltv_api.rs @@ -1,10 +1,10 @@ use crate::{ api::{ api_utils::{ - create_api_proxy_user, empty_json_response_as_array, get_user_target, + coalesce_byte_stream, create_api_proxy_user, empty_json_response_as_array, get_user_target, get_user_target_by_credentials, internal_server_error, resource_response, - stream_json_or_bin_response_stream, try_unwrap_body, + stream_json_or_bin_response_try_stream, try_unwrap_body, }, model::{AppState, UserApiRequest, UserApiRequestQueryOrBody}, }, @@ -129,18 +129,33 @@ pub async fn serve_epg_web_ui( let bg_lock = app_state.app_config.file_locks.read_lock(epg_path).await; let epg_path = epg_path.to_path_buf(); let target_name = target.name.clone(); - let (tx, rx) = mpsc::channel::(64); + let (tx, rx) = mpsc::channel::>(64); let epg_path_for_log = epg_path.clone(); let target_name_for_log = target_name.clone(); + let join_error_tx = tx.clone(); let handle = task::spawn_blocking(move || { let _guard = bg_lock; - let Ok(query) = BPlusTreeQuery::, EpgChannel>::try_new(&epg_path) else { - error!("Failed to open epg db for target {} {}", target_name, epg_path.display()); - return; + let query = match BPlusTreeQuery::, EpgChannel>::try_new(&epg_path) { + Ok(query) => query, + Err(error) => { + let message = format!("Failed to open epg db for target {target_name} {}: {error}", epg_path.display()); + error!("{message}"); + let _ = tx.blocking_send(Err(message)); + return; + } }; - for (_, channel) in query.disk_iter() { - if tx.blocking_send(channel).is_err() { + for entry in query.disk_iter() { + let (_, channel) = match entry { + Ok(entry) => entry, + Err(error) => { + let message = format!("EPG stream failed for {}: {error}", epg_path.display()); + error!("{message}"); + let _ = tx.blocking_send(Err(message)); + break; + } + }; + if tx.blocking_send(Ok(channel)).is_err() { break; } } @@ -152,12 +167,15 @@ pub async fn serve_epg_web_ui( target_name_for_log, epg_path_for_log.display() ); + let _ = join_error_tx + .send(Err(format!("EPG web UI producer task failed for target {} {}: {err}", target_name_for_log, epg_path_for_log.display()))) + .await; } }); let stream = LockedReceiverStream::new(rx, iter_lock) - .map(move |channel| rewrite_epg_channel_resource_url(&encrypt_secret, &resource_url, channel)); - return stream_json_or_bin_response_stream(accept, stream); + .map(move |result| result.map(|channel| rewrite_epg_channel_resource_url(&encrypt_secret, &resource_url, channel))); + return stream_json_or_bin_response_try_stream(accept, stream); } try_unwrap_body!(empty_json_response_as_array()) } @@ -221,18 +239,33 @@ async fn serve_epg_with_rewrites( let bg_lock = app_state.app_config.file_locks.read_lock(epg_path).await; let epg_path = epg_path.to_path_buf(); - let (channel_tx, mut channel_rx) = mpsc::channel::(256); + let (channel_tx, mut channel_rx) = mpsc::channel::>(256); let epg_path_for_log = epg_path.clone(); + let join_error_tx = channel_tx.clone(); let spawn_handle = task::spawn_blocking(move || { let _guard = bg_lock; - let Ok(mut query) = BPlusTreeQuery::, EpgChannel>::try_new(&epg_path) else { - error!("Failed to open BPlusTreeQuery {}", epg_path.display()); - return; + let mut query = match BPlusTreeQuery::, EpgChannel>::try_new(&epg_path) { + Ok(query) => query, + Err(error) => { + let message = format!("Failed to open BPlusTreeQuery {}: {error}", epg_path.display()); + error!("{message}"); + let _ = channel_tx.blocking_send(Err(message)); + return; + } }; - for (_, channel) in query.iter() { - if channel_tx.blocking_send(channel).is_err() { + for entry in query.iter() { + let (_, channel) = match entry { + Ok(entry) => entry, + Err(error) => { + let message = format!("EPG rewrite stream failed for {}: {error}", epg_path.display()); + error!("{message}"); + let _ = channel_tx.blocking_send(Err(message)); + break; + } + }; + if channel_tx.blocking_send(Ok(channel)).is_err() { break; } } @@ -240,6 +273,9 @@ async fn serve_epg_with_rewrites( tokio::spawn(async move { if let Err(err) = spawn_handle.await { error!("EPG rewrite producer task failed for {}: {err}", epg_path_for_log.display()); + let _ = join_error_tx + .send(Err(format!("EPG rewrite producer task failed for {}: {err}", epg_path_for_log.display()))) + .await; } }); @@ -257,7 +293,14 @@ async fn serve_epg_with_rewrites( } let mut writer = quick_xml::writer::Writer::new(tx); - while let Some(channel) = channel_rx.recv().await { + while let Some(result) = channel_rx.recv().await { + let channel = match result { + Ok(channel) => channel, + Err(error) => { + error!("{error}"); + return; + } + }; let programmes = if limit > 0 { channel.get_programme_with_limit(limit) } else { @@ -350,7 +393,7 @@ async fn serve_epg_with_rewrites( let body_stream = ReaderStream::new(rx); try_unwrap_body!(axum::response::Response::builder() .header(axum::http::header::CONTENT_TYPE, mime::TEXT_XML.to_string()) - .body(axum::body::Body::from_stream(body_stream))) + .body(axum::body::Body::from_stream(coalesce_byte_stream(body_stream)))) } /// Looks up an EPG channel by its exact, target-output-case storage key. diff --git a/backend/src/api/endpoints/xtream_api.rs b/backend/src/api/endpoints/xtream_api.rs index 291a8d8b9..0df24cd86 100644 --- a/backend/src/api/endpoints/xtream_api.rs +++ b/backend/src/api/endpoints/xtream_api.rs @@ -5,7 +5,7 @@ use crate::{ api::{ api_utils, api_utils::{ - admission_failure_response, create_api_proxy_user, create_catchup_session_key, + admission_failure_response, coalesce_byte_stream, create_api_proxy_user, create_catchup_session_key, create_playback_session_fingerprint, create_session_fingerprint, empty_json_response_as_array, empty_json_response_as_object, force_provider_stream_response, get_session_reservation_ttl_secs, get_user_target, get_user_target_by_credentials, internal_server_error, is_seek_request, @@ -1593,17 +1593,19 @@ async fn xtream_player_api( fn xtream_create_content_stream(xtream_iter: S) -> impl Stream> where - S: Stream + Send + Unpin + 'static, + S: Stream> + Send + Unpin + 'static, { - let mapped = xtream_iter.map(move |(mut line, has_next)| { - if has_next { - line.push(','); - } - Ok::(Bytes::from(line)) + let mapped = xtream_iter.map(move |entry| { + entry.map_err(|error| error.to_string()).map(|(mut line, has_next)| { + if has_next { + line.push(','); + } + Bytes::from(line) + }) }); - stream::once(async { Ok::(Bytes::from("[")) }) + coalesce_byte_stream(stream::once(async { Ok::(Bytes::from("[")) }) .chain(mapped) - .chain(stream::once(async { Ok::(Bytes::from("]")) })) + .chain(stream::once(async { Ok::(Bytes::from("]")) }))) } async fn xtream_player_api_get( diff --git a/backend/src/api/model/hls_cache/manifest_fetch.rs b/backend/src/api/model/hls_cache/manifest_fetch.rs index 10a572cb9..0747ad661 100644 --- a/backend/src/api/model/hls_cache/manifest_fetch.rs +++ b/backend/src/api/model/hls_cache/manifest_fetch.rs @@ -1865,6 +1865,7 @@ pub fn retry_after_delay_ms(headers: &HeaderMap) -> Option { fn current_time_millis() -> u64 { chrono::Utc::now().timestamp_millis().try_into().unwrap_or_default() } #[cfg(test)] +#[allow(clippy::too_many_lines)] pub(crate) async fn refresh_from_live_hls_entrypoint_with_retries( origin_entry: &LiveHlsOriginEntry, headers: &HeaderMap, @@ -1930,12 +1931,14 @@ pub(crate) async fn refresh_from_live_hls_entrypoint_with_retries( origin_entry, attempt_index, next_retry_delay_ms(retry_policy, attempt_index, None, 0), - "error=request".to_string(), + "error=request", ); } Ok(Err( err @ (OriginManifestFetchError::ContentDecoding { .. } - | OriginManifestFetchError::ContentCoding(ContentCodingError::PrefixRead(_))), + | OriginManifestFetchError::ContentCoding(ContentCodingError::PrefixRead(_)) + | OriginManifestFetchError::Redirect(_) + | OriginManifestFetchError::Timeout), )) => { if attempt_index + 1 == attempts { return Err(err); @@ -1947,26 +1950,10 @@ pub(crate) async fn refresh_from_live_hls_entrypoint_with_retries( format!("error={}", err.log_label()), ); } - Ok(Err(err @ (OriginManifestFetchError::Redirect(_) | OriginManifestFetchError::Timeout))) => { - if attempt_index + 1 == attempts { - return Err(err); - } - log_origin_refresh_retry_scheduled( - origin_entry, - attempt_index, - next_retry_delay_ms(retry_policy, attempt_index, None, 0), - format!("error={}", err.log_label()), - ); - } Ok(Err(OriginManifestFetchError::RetryExhausted)) => return Err(OriginManifestFetchError::RetryExhausted), Ok(Err(OriginManifestFetchError::ProviderUnavailable(kind))) => { return Err(OriginManifestFetchError::ProviderUnavailable(kind)); } - Ok(Err( - err @ (OriginManifestFetchError::ContentCoding(_) - | OriginManifestFetchError::DecodedBodyLimitExceeded { .. } - | OriginManifestFetchError::InvalidUtf8 { .. }), - )) => return Err(err), Err(OriginManifestFetchError::Timeout) => { if attempt_index + 1 == attempts { return Err(OriginManifestFetchError::Timeout); @@ -1978,7 +1965,12 @@ pub(crate) async fn refresh_from_live_hls_entrypoint_with_retries( "error=timeout", ); } - Err(err) => return Err(err), + Ok(Err( + err @ (OriginManifestFetchError::ContentCoding(_) + | OriginManifestFetchError::DecodedBodyLimitExceeded { .. } + | OriginManifestFetchError::InvalidUtf8 { .. }), + )) + | Err(err) => return Err(err), } } diff --git a/backend/src/api/model/hls_cache/refresh.rs b/backend/src/api/model/hls_cache/refresh.rs index 3a59459ee..1c5d0d55b 100644 --- a/backend/src/api/model/hls_cache/refresh.rs +++ b/backend/src/api/model/hls_cache/refresh.rs @@ -2252,7 +2252,7 @@ mod tests { request.origin_entry = entry; request.now_ms = now_ms; - assert!(trigger_origin_refresh_sync(request).await); + assert!(Box::pin(trigger_origin_refresh_sync(request)).await); let lease = hls_proxy .access_leases() @@ -2402,7 +2402,7 @@ mod tests { let mut request = test_origin_refresh_request(Arc::clone(&session)); request.origin_entry = entry; - assert!(trigger_origin_refresh_sync(request).await); + assert!(Box::pin(trigger_origin_refresh_sync(request)).await); let session = session.read().await; assert!(!session.origin_request_headers.contains_key(header::COOKIE)); @@ -3063,7 +3063,7 @@ mod tests { LiveHlsOriginEntry::parse(&format!("{}/live/user/pass/12345.m3u8", server.base_url)).expect("entry url"); let metrics = Arc::clone(request.segment_worker_pool.metrics()); - assert!(trigger_origin_refresh_sync(request).await); + assert!(Box::pin(trigger_origin_refresh_sync(request)).await); let snapshot = metrics.snapshot(); assert_eq!(snapshot.refresh_started, 1); diff --git a/backend/src/api/model/hls_cache/segment_fetcher.rs b/backend/src/api/model/hls_cache/segment_fetcher.rs index 5ecb74a92..d1092134f 100644 --- a/backend/src/api/model/hls_cache/segment_fetcher.rs +++ b/backend/src/api/model/hls_cache/segment_fetcher.rs @@ -1015,7 +1015,10 @@ mod tests { let mut head = format!("HTTP/1.1 {} {reason}\r\nContent-Length: {}\r\n", response.status, response.body.len()); for (name, value) in response.headers { - head.push_str(&format!("{name}: {value}\r\n")); + head.push_str(&name); + head.push_str(": "); + head.push_str(&value); + head.push_str("\r\n"); } head.push_str("Connection: close\r\n\r\n"); let _ = socket.write_all(head.as_bytes()).await; diff --git a/backend/src/api/model/hls_cache/transient_fetcher.rs b/backend/src/api/model/hls_cache/transient_fetcher.rs index df69d5b12..fe546093b 100644 --- a/backend/src/api/model/hls_cache/transient_fetcher.rs +++ b/backend/src/api/model/hls_cache/transient_fetcher.rs @@ -1118,7 +1118,7 @@ mod tests { } struct TestTransientCacheFixture { - _temp_dir: tempfile::TempDir, + temp_dir: tempfile::TempDir, segment_cache: Arc, segment_repair: Arc, session: HlsSessionHandle, @@ -1160,7 +1160,7 @@ mod tests { ..Default::default() })); Self { - _temp_dir: temp_dir, + temp_dir, segment_cache, segment_repair, session, @@ -1677,7 +1677,7 @@ mod tests { assert_eq!(dropped_guards.load(Ordering::Relaxed), 1); assert!(fixture.segment_cache.metadata(&fixture.cache_key).await.expect("cache metadata reads").is_none()); assert!(!fixture.segment_cache.has_active_temp_files().await); - assert_eq!(std::fs::read_dir(fixture._temp_dir.path()).expect("cache root reads").count(), 0); + assert_eq!(std::fs::read_dir(fixture.temp_dir.path()).expect("cache root reads").count(), 0); let session = fixture.session.read().await; let entry = session.transient.object_cache.get(&fixture.cache_key).expect("fetching cache entry remains"); assert!(!matches!(entry.status, TransientObjectCacheStatus::Ready { .. })); diff --git a/backend/src/api/model/metadata_update_manager.rs b/backend/src/api/model/metadata_update_manager.rs index 26e20fc9a..67f436dcf 100644 --- a/backend/src/api/model/metadata_update_manager.rs +++ b/backend/src/api/model/metadata_update_manager.rs @@ -496,7 +496,8 @@ fn load_metadata_retry_states_from_disk(path: &Path) -> io::Result = Vec::new(); let mut query = BPlusTreeQuery::::try_new(path)?; - for (key, value) in query.iter() { + for entry in query.iter() { + let (key, value) = entry?; if let Some(state) = value.clone().into_task_retry_state() { result.insert(key.into_task_key(), state); } else { diff --git a/backend/src/api/model/streams/provider_stream_factory.rs b/backend/src/api/model/streams/provider_stream_factory.rs index 0138b8bbc..2b0e4ed3f 100644 --- a/backend/src/api/model/streams/provider_stream_factory.rs +++ b/backend/src/api/model/streams/provider_stream_factory.rs @@ -1301,7 +1301,13 @@ mod tests { let addr = listener.local_addr().unwrap(); let requests = Arc::new(AtomicUsize::new(0)); let task_requests = Arc::clone(&requests); - let response_headers = headers.iter().map(|(name, value)| format!("{name}: {value}\r\n")).collect::(); + let response_headers = headers.iter().fold(String::new(), |mut response, (name, value)| { + response.push_str(name); + response.push_str(": "); + response.push_str(value); + response.push_str("\r\n"); + response + }); tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); task_requests.fetch_add(1, Ordering::SeqCst); diff --git a/backend/src/repository/bplustree.rs b/backend/src/repository/bplustree.rs deleted file mode 100644 index 0f3293167..000000000 --- a/backend/src/repository/bplustree.rs +++ /dev/null @@ -1,8642 +0,0 @@ -//! Legacy-compatible B+Tree storage v2. -//! -//! This module remains the active compatibility branch for existing v2 files. -//! Stabilization work must preserve `STORAGE_VERSION = 2` and avoid requiring a -//! rewrite of existing repositories. The future v3 storage line is expected to -//! live behind an explicit version boundary and a typed migration path. - -use crate::{ - repository::storage::get_file_path_for_db_index, - utils, - utils::{binary_deserialize, binary_serialize, binary_serialize_into}, -}; -use fs2::FileExt as _; -use log::{error, warn}; -use lru::LruCache; -use memmap2::{Advice, Mmap}; -use parking_lot::Mutex; -use serde::{Deserialize, Serialize}; -use shared::error::{string_to_io_error, to_io_error}; -use smallvec::{smallvec, SmallVec}; -#[cfg(unix)] -use std::os::unix::fs::{FileExt, MetadataExt}; -use std::{ - borrow::Cow, - ffi::OsString, - fs::{File, Metadata, OpenOptions}, - io::{self, BufReader, Read, Seek, SeekFrom, Write}, - marker::PhantomData, - mem::size_of, - num::NonZeroUsize, - ops::Bound, - path::{Path, PathBuf}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, - thread::JoinHandle, - 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}; - -const PAGE_SIZE: u16 = 4096; -pub const PAGE_SIZE_USIZE: usize = PAGE_SIZE as usize; -const ZERO_PADDING_CHUNK: [u8; PAGE_SIZE_USIZE] = [0u8; PAGE_SIZE_USIZE]; -const LEN_SIZE: usize = 4; -const FLAG_SIZE: usize = 1; -pub(crate) const MAGIC: &[u8; 4] = b"BTRE"; -pub(crate) const STORAGE_VERSION: u32 = 2; -const HEADER_SIZE: u64 = PAGE_SIZE as u64; -const ROOT_OFFSET_POS: u64 = 8; -const METADATA_OFFSET_POS: u64 = 16; -const METADATA_DATA_START_POS: usize = 20; -// Reserve space for metadata (e.g. 4096 - 16 = 4080 bytes max, but let's be safe) -const METADATA_MAX_SIZE: usize = 4000; -const HEADER_FLAG_HAS_METADATA_FLAGS: u32 = 1 << 31; -const HEADER_FLAG_HAS_TOMBSTONES: u32 = 1 << 30; -const HEADER_METADATA_LEN_MASK: u32 = !(HEADER_FLAG_HAS_METADATA_FLAGS | HEADER_FLAG_HAS_TOMBSTONES); - -#[inline] -const fn encode_metadata_len_with_flags(metadata_len: u32, has_tombstones: bool) -> u32 { - let mut encoded = metadata_len | HEADER_FLAG_HAS_METADATA_FLAGS; - if has_tombstones { - encoded |= HEADER_FLAG_HAS_TOMBSTONES; - } - encoded -} - -#[inline] -const fn decode_metadata_len_and_flags(raw: u32) -> (u32, bool) { - let metadata_len = raw & HEADER_METADATA_LEN_MASK; - let has_metadata_flags = (raw & HEADER_FLAG_HAS_METADATA_FLAGS) != 0; - let has_tombstones = if has_metadata_flags { - (raw & HEADER_FLAG_HAS_TOMBSTONES) != 0 - } else { - // Legacy v2 files (without header flags) are treated conservatively: - // assume tombstones may exist until a rewrite/compact writes proper flags. - true - }; - (metadata_len, has_tombstones) -} - -// Maximum number of blocks to cache in memory (~4MB at 4KB per block) -const CACHE_CAPACITY: usize = 1024; -const INTERNAL_CACHE_CAPACITY: usize = 256; -const QUERY_REFRESH_INTERVAL: Duration = Duration::from_millis(25); - -// v2 uses conservative runtime fanout instead of pretending that size_of::() -// predicts serialized key size. Multi-block nodes keep existing files compatible. -const DEFAULT_INNER_ORDER: usize = 64; -const DEFAULT_LEAF_ORDER: usize = 64; - -// Value packing configuration -const SMALL_VALUE_THRESHOLD: usize = 256; -const PACK_BLOCK_HEADER_SIZE: usize = 4; -const PACK_VALUE_HEADER_SIZE: usize = 4; - -// LZ4 compression configuration -const COMPRESSION_MIN_SIZE: usize = 64; -const COMPRESSION_THRESHOLD_PERCENT: usize = 85; -const COMPRESSION_FLAG_NONE: u8 = 0x00; -pub const COMPRESSION_FLAG_LZ4: u8 = 0x01; - -// Page Configuration -const PAGE_HEADER_SIZE: u16 = 16; -const PAGE_HEADER_SIZE_USIZE: usize = PAGE_HEADER_SIZE as usize; -const SLOT_SIZE: usize = 2; // u16 - -const MAGIC_METADATA_TARGET_ID_MAPPING: u8 = 0x01; - -type TraversalStack = SmallVec<[(u64, usize); 8]>; -type OffsetStack = SmallVec<[u64; 8]>; - -fn advise_mmap(mmap: &Mmap, advice: Advice, context: &str) { - if let Err(err) = mmap.advise(advice) { - warn!("Failed to apply mmap advice {advice:?} for {context}: {err}"); - } -} - -fn mmap_with_advice(file: &File, advice: Advice, context: &str) -> Option { - let mmap = unsafe { - match Mmap::map(file) { - Ok(mmap) => mmap, - Err(err) => { - warn!("Failed to mmap B+Tree for {context}; falling back to buffered file I/O: {err}"); - return None; - } - } - }; - advise_mmap(&mmap, advice, context); - Some(mmap) -} - -/// Read exact bytes from a file at a specific offset. -/// Uses `FileExt::read_exact_at` on Unix for thread-safe concurrent reads. -/// Falls back to `seek` + `read_exact` + `seek` back on other platforms. -pub(crate) fn read_exact_at_offset(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { - #[cfg(unix)] - file.read_exact_at(buf, offset)?; - #[cfg(not(unix))] - { - let mut f = file; - let current_pos = f.stream_position()?; - f.seek(SeekFrom::Start(offset))?; - f.read_exact(buf)?; - f.seek(SeekFrom::Start(current_pos))?; - } - Ok(()) -} - -/* - B+Tree File Layout - ================== - - ┌─────────────────────────────────────────────────────────────┐ - │ File Header (PAGE_SIZE bytes, currently 4096) │ - ├─────────────────────────────────────────────────────────────┤ - │ MAGIC [4B: "BTRE"] │ - │ VERSION [4B: u32] │ - │ ROOT_OFFSET [8B: u64] │ - │ METADATA_LEN_FLAGS [4B: u32] │ - │ bit31: metadata flags initialized │ - │ bit30: has_tombstones │ - │ bits0..29: metadata length │ - │ METADATA [variable, up to 4000B] │ - │ [padding to PAGE_SIZE] │ - └─────────────────────────────────────────────────────────────┘ - - Leaf Node Layout (single or multi-block) - ┌─────────────────────────────────────────────────────────────┐ - │ IS_LEAF [1B: 0x01] │ - │ KEYS_LEN [4B: u32] │ - │ KEYS [MessagePack serialized Vec] │ - │ VALUE_INFO_LEN [4B: u32] │ - │ VALUE_INFO [MessagePack serialized Vec] │ - │ [padding to block boundary] │ - └─────────────────────────────────────────────────────────────┘ - - Internal Node Layout (supports multi-block when content exceeds PAGE_SIZE) - ┌─────────────────────────────────────────────────────────────┐ - │ IS_LEAF [1B: 0x00] │ - │ KEYS_LEN [4B: u32] │ - │ KEYS [MessagePack serialized Vec] │ - │ POINTERS_LEN [4B: u32] │ - │ POINTERS [MessagePack serialized Vec] │ - │ [padding to block boundary] │ - └─────────────────────────────────────────────────────────────┘ - - Note: Internal nodes can span multiple PAGE_SIZE blocks when - keys + pointers exceed a single page. The order calculation - uses a 75% safety factor to minimize multi-block nodes. - - Value Storage Modes: - - Single: Large values stored at [offset] with optional LZ4 compression - Format: [FLAG:1B][payload...] where FLAG = 0x00 (raw) or 0x01 (LZ4) - - Packed: Small values (≤256B) packed into PAGE_SIZE blocks - Format: [COUNT:4B][LEN:4B][data...][LEN:4B][data...]... -*/ - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[repr(u8)] -pub enum PageType { - Leaf = 1, - Internal = 2, - Overflow = 3, -} - -#[derive(Debug, Clone, Copy)] -pub struct PageHeader { - pub page_type: PageType, // 0x01=Leaf, 0x02=Internal, 0x03=Overflow - pub cell_count: u16, // Number of active cells - pub free_start: u16, // Offset to start of free space (after slots) - pub free_end: u16, // Offset to end of free space (before cells) - pub right_sibling: u64, // 0 if none, pointer to next leaf (for range scans) - pub checksum: u32, // TODO Data integrity check, currently not neccessary, maybe in future -} - -impl PageHeader { - pub fn new(page_type: PageType) -> Self { - Self { - page_type, - cell_count: 0, - free_start: PAGE_HEADER_SIZE, - free_end: PAGE_SIZE, - right_sibling: 0, - checksum: 0, - } - } - - pub fn serialize(&self) -> [u8; PAGE_HEADER_SIZE_USIZE] { - let mut buf = [0u8; PAGE_HEADER_SIZE_USIZE]; - buf[0] = self.page_type as u8; - buf[1] = 0; // padding - buf[2..4].copy_from_slice(&self.cell_count.to_le_bytes()); - buf[4..6].copy_from_slice(&self.free_start.to_le_bytes()); - buf[6..8].copy_from_slice(&self.free_end.to_le_bytes()); - buf[8..16].copy_from_slice(&self.right_sibling.to_le_bytes()); - buf - } - - pub fn deserialize(buf: &[u8]) -> Result { - if buf.len() < PAGE_HEADER_SIZE_USIZE { - return Err(PageError::Corrupted); - } - let page_type = match buf[0] { - 1 => PageType::Leaf, - 2 => PageType::Internal, - 3 => PageType::Overflow, - _ => return Err(PageError::Corrupted), - }; - - // Use try_into to safely read bytes, although the length check above makes it safe. - // we can map err. - let cell_count = u16::from_le_bytes(buf[2..4].try_into().map_err(|_| PageError::Corrupted)?); - let free_start = u16::from_le_bytes(buf[4..6].try_into().map_err(|_| PageError::Corrupted)?); - let free_end = u16::from_le_bytes(buf[6..8].try_into().map_err(|_| PageError::Corrupted)?); - let right_sibling = u64::from_le_bytes(buf[8..16].try_into().map_err(|_| PageError::Corrupted)?); - - Ok(Self { page_type, cell_count, free_start, free_end, right_sibling, checksum: 0 }) - } -} - -pub struct SlottedPage<'a> { - pub header: PageHeader, - data: &'a mut [u8], -} - -#[derive(Debug)] -pub enum PageError { - NoSpace, - InvalidIndex, - Corrupted, - Io(io::Error), -} - -impl std::fmt::Display for PageError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - PageError::NoSpace => write!(f, "Page has no space for insertion"), - PageError::InvalidIndex => write!(f, "Invalid cell index"), - PageError::Corrupted => write!(f, "Page data is corrupted"), - PageError::Io(err) => write!(f, "I/O error: {err}"), - } - } -} - -impl std::error::Error for PageError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - match self { - PageError::Io(err) => Some(err), - _ => None, - } - } -} - -impl From for PageError { - fn from(err: io::Error) -> Self { PageError::Io(err) } -} - -/// Error types for B+Tree operations that distinguish between different failure modes. -/// This allows callers to handle "key not found" differently from actual errors like corruption. -#[derive(Debug)] -pub enum BPlusTreeError { - /// An I/O error occurred during file operations - Io(io::Error), - /// Data corruption detected during deserialization - Corrupted(String), - /// The tree structure is invalid (e.g., missing child pointers) - InvalidStructure(String), - /// The requested key was not found in the tree (used for update operations) - KeyNotFound, -} - -impl std::fmt::Display for BPlusTreeError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - BPlusTreeError::Io(err) => write!(f, "I/O error: {err}"), - BPlusTreeError::Corrupted(msg) => write!(f, "Data corrupted: {msg}"), - BPlusTreeError::InvalidStructure(msg) => write!(f, "Invalid structure: {msg}"), - BPlusTreeError::KeyNotFound => write!(f, "Key not found"), - } - } -} - -impl std::error::Error for BPlusTreeError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - match self { - BPlusTreeError::Io(err) => Some(err), - _ => None, - } - } -} - -impl From for BPlusTreeError { - fn from(err: io::Error) -> Self { BPlusTreeError::Io(err) } -} - -impl BPlusTreeError { - pub fn to_io(self) -> io::Error { - match self { - BPlusTreeError::Io(e) => e, - BPlusTreeError::KeyNotFound => io::Error::new(io::ErrorKind::NotFound, "Key not found"), - err => io::Error::new(io::ErrorKind::InvalidData, err), - } - } -} - -impl From for BPlusTreeError { - fn from(err: PageError) -> Self { - match err { - PageError::Io(e) => BPlusTreeError::Io(e), - PageError::Corrupted => BPlusTreeError::Corrupted("Page data corrupted".into()), - other => BPlusTreeError::InvalidStructure(format!("{other:?}")), - } - } -} - -impl<'a> SlottedPage<'a> { - pub fn new(data: &'a mut [u8], page_type: PageType) -> Result { - if data.len() < PAGE_HEADER_SIZE_USIZE { - return Err(PageError::NoSpace); - } - let header = PageHeader::new(page_type); - // Initialize header in buffer - let h_bytes = header.serialize(); - data[..PAGE_HEADER_SIZE_USIZE].copy_from_slice(&h_bytes); - Ok(Self { header, data }) - } - - pub fn open(data: &'a mut [u8]) -> Result { - if data.len() < PAGE_HEADER_SIZE_USIZE { - return Err(PageError::Corrupted); - } - let header = PageHeader::deserialize(&data[..PAGE_HEADER_SIZE_USIZE])?; - Ok(Self { header, data }) - } - - pub fn commit(&mut self) { - let h_bytes = self.header.serialize(); - if self.data.len() >= PAGE_HEADER_SIZE_USIZE { - self.data[..PAGE_HEADER_SIZE_USIZE].copy_from_slice(&h_bytes); - } - } - - pub fn free_space(&self) -> usize { - if self.header.free_end >= self.header.free_start { - (self.header.free_end - self.header.free_start) as usize - } else { - 0 - } - } - - /// Insert a cell directly. Caller must ensure specific order (e.g. invalidating current sort). - /// Typically used by `insert_at_index`. - fn append_cell(&mut self, cell_data: &[u8]) -> Result { - let required = cell_data.len(); - if self.free_space() < required + SLOT_SIZE { - return Err(PageError::NoSpace); - } - - let req_u16 = u16::try_from(required).map_err(|_| PageError::NoSpace)?; - // Data grows downwards. Safe cast due to page size check. - let offset = self.header.free_end.checked_sub(req_u16).ok_or(PageError::NoSpace)?; - - // Bounds check - if (offset as usize) + required > self.data.len() { - return Err(PageError::NoSpace); - } - - self.data[offset as usize..(offset as usize + required)].copy_from_slice(cell_data); - - self.header.free_end = offset; - Ok(offset) - } - - pub fn insert_at_index(&mut self, index: usize, val: &[u8]) -> Result<(), PageError> { - // 1. Append cell data - let offset = self.append_cell(val)?; - - // 2. Insert slot - let slot_area_start = PAGE_HEADER_SIZE_USIZE; - let count = self.header.cell_count as usize; - - if index > count { - return Err(PageError::InvalidIndex); - } - - // Shift slots if necessary - let insert_pos = slot_area_start + (index * SLOT_SIZE); - if self.data.len() < insert_pos + SLOT_SIZE { - return Err(PageError::NoSpace); // Should cover src_start..src_end too if valid - } - - if index < count { - let src_start = insert_pos; - let src_end = slot_area_start + (count * SLOT_SIZE); - let dest_start = insert_pos + SLOT_SIZE; - - if self.data.len() < dest_start + (src_end - src_start) { - return Err(PageError::NoSpace); - } - self.data.copy_within(src_start..src_end, dest_start); - } - - // Write new slot - if insert_pos + 2 > self.data.len() { - return Err(PageError::NoSpace); - } - self.data[insert_pos..insert_pos + 2].copy_from_slice(&offset.to_le_bytes()); - - // Update header - self.header.cell_count += 1; - self.header.free_start += u16::try_from(SLOT_SIZE).map_err(|_| PageError::NoSpace)?; - self.commit(); - - Ok(()) - } - - // assumes all cells start with a 4-byte length header - // This creates tight coupling between SlottedPage (a generic page structure) - // and the specific cell format used by BPlusTreeNode. - pub fn get_cell(&self, index: usize) -> Option<&[u8]> { - if index >= self.header.cell_count as usize { - return None; - } - let slot_pos = PAGE_HEADER_SIZE_USIZE + (index * SLOT_SIZE); - // Safe slice access - if slot_pos + 2 > self.data.len() { - return None; - } - let offset = u16::from_le_bytes(self.data[slot_pos..slot_pos + 2].try_into().ok()?); - let offset = usize::from(offset); - let slot_area_end = PAGE_HEADER_SIZE_USIZE + (self.header.cell_count as usize * SLOT_SIZE); - - // Cell payload must not point into the page header or slot directory. - if offset < slot_area_end || offset < usize::from(self.header.free_end) { - return None; - } - - // Bounds check for length header - if offset + 4 > self.data.len() { - return None; - } - let len = u32::from_le_bytes(self.data[offset..offset + 4].try_into().ok()?) as usize; - - if offset + 4 + len > self.data.len() { - return None; - } - Some(&self.data[offset..offset + 4 + len]) - } - - pub fn get_cell_offset(&self, index: usize) -> Option { - let slot_pos = PAGE_HEADER_SIZE_USIZE + (index * SLOT_SIZE); - if slot_pos + 2 > self.data.len() { - return None; - } - Some(u16::from_le_bytes(self.data[slot_pos..slot_pos + 2].try_into().ok()?)) - } - - pub fn compact(&mut self) -> Result<(), PageError> { - let mut temp = vec![0u8; PAGE_SIZE_USIZE]; - { - let mut new_page = SlottedPage::new(&mut temp, self.header.page_type)?; - for i in 0..self.header.cell_count as usize { - if let Some(cell) = self.get_cell(i) { - if let Err(e) = new_page.insert_at_index(i, cell) { - error!("Compact insert failed at index {i}: {e:?}"); - return Err(e); - } - } else { - error!("Compact get_cell failed at index {i}"); - return Err(PageError::Corrupted); - } - } - } - self.data.copy_from_slice(&temp); - self.header = PageHeader::deserialize(&self.data[..PAGE_HEADER_SIZE_USIZE])?; - Ok(()) - } - - pub fn split_off(&mut self) -> Result>, PageError> { - let count = self.header.cell_count as usize; - let mut total_bytes = 0; - let mut split_idx = count / 2; - - let mut sizes = Vec::with_capacity(count); - for i in 0..count { - if let Some(cell) = self.get_cell(i) { - sizes.push(cell.len()); - total_bytes += cell.len(); - } else { - sizes.push(0); - } - } - - let target = total_bytes / 2; - let mut current = 0; - for (i, &s) in sizes.iter().enumerate() { - current += s; - if current >= target { - split_idx = i + 1; - break; - } - } - - // Fix for split logic: - if count == 0 { - return Err(PageError::InvalidIndex); // Cannot split empty page - } - if count == 1 { - // Cannot split single item fundamentally. - // Return Ok(None) explicitly to indicate no-op. - return Ok(None); - } - - if split_idx >= count { - split_idx = count.saturating_sub(1); - } - if split_idx == 0 && count > 1 { - split_idx = 1; - } - - let mut new_buffer = vec![0u8; PAGE_SIZE_USIZE]; - { - let mut new_page = SlottedPage::new(&mut new_buffer, self.header.page_type)?; - for i in split_idx..count { - if let Some(cell) = self.get_cell(i) { - new_page.insert_at_index(i - split_idx, cell)?; - } - } - } - - self.header.cell_count = u16::try_from(split_idx).map_err(|_| PageError::InvalidIndex)?; - let new_free_start = PAGE_HEADER_SIZE_USIZE + split_idx * SLOT_SIZE; - self.header.free_start = u16::try_from(new_free_start).map_err(|_| PageError::NoSpace)?; - self.commit(); - - self.compact()?; - - Ok(Some(new_buffer)) - } -} - -#[inline] -fn u32_from_bytes(bytes: &[u8]) -> io::Result { Ok(u32::from_le_bytes(bytes.try_into().map_err(to_io_error)?)) } - -#[inline] -fn node_flag_to_is_leaf(flag: u8) -> io::Result { - match flag { - 0 => Ok(false), - 1 => Ok(true), - _ => Err(io::Error::new(io::ErrorKind::InvalidData, format!("Invalid B+Tree node flag: {flag}"))), - } -} - -#[inline] -fn checked_slice_range(start: usize, len: usize, total_len: usize) -> io::Result> { - let end = start - .checked_add(len) - .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "B+Tree node slice offset overflow"))?; - if end > total_len { - return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "B+Tree node slice out of bounds")); - } - Ok(start..end) -} - -#[inline] -fn u64_from_bytes(bytes: &[u8]) -> io::Result { Ok(u64::from_le_bytes(bytes.try_into().map_err(to_io_error)?)) } - -#[inline] -fn get_entry_index_upper_bound(keys: &[K], key: &K) -> usize -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut left = 0; - let mut right = keys.len(); - while left < right { - let mid = left + ((right - left) >> 1); - if &keys[mid] <= key { - left = mid + 1; - } else { - right = mid; - } - } - left -} - -// Provide zero-copy scanning of MessagePack-encoded keys in B+tree internal nodes. -// This avoids deserializing the entire keys vector (which allocates for each key in the node) -// by scanning the raw bytes directly to find the correct child pointer. - -/// Result of parsing a `MessagePack` array header -struct MsgPackArrayHeader { - /// Number of elements in the array - count: usize, - /// Number of bytes consumed by the header - header_size: usize, -} - -/// Parse `MessagePack` array header to get element count -/// Returns (count, `header_bytes_consumed`) -#[inline] -fn parse_msgpack_array_header(bytes: &[u8]) -> Option { - if bytes.is_empty() { - return None; - } - let first = bytes[0]; - - // fixarray (0x90-0x9f): count in low 4 bits - if (0x90..=0x9f).contains(&first) { - return Some(MsgPackArrayHeader { count: (first & 0x0f) as usize, header_size: 1 }); - } - - // array16 (0xdc): next 2 bytes are big-endian count - if first == 0xdc && bytes.len() >= 3 { - let count = u16::from_be_bytes([bytes[1], bytes[2]]) as usize; - return Some(MsgPackArrayHeader { count, header_size: 3 }); - } - - // array32 (0xdd): next 4 bytes are big-endian count - if first == 0xdd && bytes.len() >= 5 { - let count = u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize; - return Some(MsgPackArrayHeader { count, header_size: 5 }); - } - - None -} - -/// Trait for types that can be scanned directly from `MessagePack` bytes without allocation. -/// -/// This enables zero-copy key scanning in B+tree internal nodes, avoiding the need -/// to deserialize the entire keys vector just to find which child pointer to follow. -pub trait MsgPackScannable: Ord + Sized { - /// Compare self with a key encoded at the given position in the byte slice. - /// Returns: - /// - `Some((ordering, bytes_consumed))` if successfully parsed - /// - `None` if parsing failed - fn compare_at_position(&self, bytes: &[u8]) -> Option<(std::cmp::Ordering, usize)>; - - /// Skip over a key at the given position without comparing. - /// Returns the number of bytes consumed, or None if parsing failed. - fn skip_at_position(bytes: &[u8]) -> Option; -} - -#[inline] -fn compare_u32_with_msgpack_signed(lhs: u32, value: i64, consumed: usize) -> (std::cmp::Ordering, usize) { - if value < 0 { - return (std::cmp::Ordering::Greater, consumed); - } - let Ok(value) = u32::try_from(value) else { - return (std::cmp::Ordering::Less, consumed); - }; - (lhs.cmp(&value), consumed) -} - -impl MsgPackScannable for u32 { - #[inline] - fn compare_at_position(&self, bytes: &[u8]) -> Option<(std::cmp::Ordering, usize)> { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - // Positive fixint (0x00-0x7f): value is the byte itself - if first <= 0x7f { - let value = u32::from(first); - return Some((self.cmp(&value), 1)); - } - - // uint8 (0xcc): next byte is the value - if first == 0xcc && bytes.len() >= 2 { - let value = u32::from(bytes[1]); - return Some((self.cmp(&value), 2)); - } - - // uint16 (0xcd): next 2 bytes are big-endian value - if first == 0xcd && bytes.len() >= 3 { - let value = u32::from(u16::from_be_bytes([bytes[1], bytes[2]])); - return Some((self.cmp(&value), 3)); - } - - // uint32 (0xce): next 4 bytes are big-endian value - if first == 0xce && bytes.len() >= 5 { - let value = u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]); - return Some((self.cmp(&value), 5)); - } - - // int8 (0xd0) - if first == 0xd0 && bytes.len() >= 2 { - let value = i8::from_be_bytes([bytes[1]]); - return Some(compare_u32_with_msgpack_signed(*self, i64::from(value), 2)); - } - - // int16 (0xd1) - if first == 0xd1 && bytes.len() >= 3 { - let value = i16::from_be_bytes([bytes[1], bytes[2]]); - return Some(compare_u32_with_msgpack_signed(*self, i64::from(value), 3)); - } - - // int32 (0xd2) - if first == 0xd2 && bytes.len() >= 5 { - let value = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]); - return Some(compare_u32_with_msgpack_signed(*self, i64::from(value), 5)); - } - - // int64 (0xd3) - if first == 0xd3 && bytes.len() >= 9 { - let value = i64::from_be_bytes([ - bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], bytes[8], - ]); - return Some(compare_u32_with_msgpack_signed(*self, value, 9)); - } - - // Negative fixint (0xe0-0xff) - if first >= 0xe0 { - return Some((std::cmp::Ordering::Greater, 1)); - } - - None - } - - #[inline] - fn skip_at_position(bytes: &[u8]) -> Option { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - if first <= 0x7f { - return Some(1); - } // fixint - if first == 0xcc { - return Some(2); - } // uint8 - if first == 0xcd { - return Some(3); - } // uint16 - if first == 0xce { - return Some(5); - } // uint32 - if first == 0xcf { - return Some(9); - } // uint64 - - // Negative fixint (0xe0-0xff) - if first >= 0xe0 { - return Some(1); - } - - // int8 (0xd0) - if first == 0xd0 { - return Some(2); - } - // int16 (0xd1) - if first == 0xd1 { - return Some(3); - } - // int32 (0xd2) - if first == 0xd2 { - return Some(5); - } - // int64 (0xd3) - if first == 0xd3 { - return Some(9); - } - - None - } -} - -impl MsgPackScannable for String { - #[inline] - fn compare_at_position(&self, bytes: &[u8]) -> Option<(std::cmp::Ordering, usize)> { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - // fixstr (0xa0-0xbf): length in low 5 bits - if (0xa0..=0xbf).contains(&first) { - let len = (first & 0x1f) as usize; - if bytes.len() > len { - let str_bytes = &bytes[1..=len]; - // Compare as bytes (valid UTF-8 has same ordering as String) - let ordering = self.as_bytes().cmp(str_bytes); - return Some((ordering, 1 + len)); - } - return None; - } - - // str8 (0xd9): 1 byte length - if first == 0xd9 && bytes.len() >= 2 { - let len = bytes[1] as usize; - if bytes.len() >= 2 + len { - let str_bytes = &bytes[2..2 + len]; - let ordering = self.as_bytes().cmp(str_bytes); - return Some((ordering, 2 + len)); - } - return None; - } - - // str16 (0xda): 2 byte length (big-endian) - if first == 0xda && bytes.len() >= 3 { - let len = u16::from_be_bytes([bytes[1], bytes[2]]) as usize; - if bytes.len() >= 3 + len { - let str_bytes = &bytes[3..3 + len]; - let ordering = self.as_bytes().cmp(str_bytes); - return Some((ordering, 3 + len)); - } - return None; - } - - // str32 (0xdb): 4 byte length (big-endian) - if first == 0xdb && bytes.len() >= 5 { - let len = u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize; - if bytes.len() >= 5 + len { - let str_bytes = &bytes[5..5 + len]; - let ordering = self.as_bytes().cmp(str_bytes); - return Some((ordering, 5 + len)); - } - return None; - } - - None - } - - #[inline] - fn skip_at_position(bytes: &[u8]) -> Option { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - // fixstr (0xa0-0xbf) - if (0xa0..=0xbf).contains(&first) { - let len = (first & 0x1f) as usize; - return Some(1 + len); - } - - // str8 (0xd9) - if first == 0xd9 && bytes.len() >= 2 { - let len = bytes[1] as usize; - return Some(2 + len); - } - - // str16 (0xda) - if first == 0xda && bytes.len() >= 3 { - let len = u16::from_be_bytes([bytes[1], bytes[2]]) as usize; - return Some(3 + len); - } - - // str32 (0xdb) - if first == 0xdb && bytes.len() >= 5 { - let len = u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize; - return Some(5 + len); - } - - None - } -} - -/// Find the child index for internal node traversal using zero-copy key scanning. -/// -/// This performs a linear scan through the MessagePack-encoded keys array, -/// comparing each key without deserializing the entire vector. -/// -/// Returns the index of the child pointer to follow (upper bound). -/// -/// # Arguments -/// * `keys_bytes` - The raw `MessagePack` bytes of the keys array -/// * `search_key` - The key we're searching for -/// -/// # Returns -/// * `Some(index)` - The child index to follow -/// * `None` - If parsing failed (caller should fall back to full deserialization) -#[inline] -fn find_child_index_zero_copy(keys_bytes: &[u8], search_key: &K) -> Option { - let header = parse_msgpack_array_header(keys_bytes)?; - let mut pos = header.header_size; - - // Linear scan through keys, finding upper bound - for i in 0..header.count { - if pos >= keys_bytes.len() { - return None; - } - - let (ordering, consumed) = search_key.compare_at_position(&keys_bytes[pos..])?; - - // Upper bound: first key > search_key - if ordering == std::cmp::Ordering::Less { - return Some(i); - } - - pos += consumed; - } - - // All keys are <= search_key, return count (rightmost child) - Some(header.count) -} - -/// Zero-copy scan result for internal nodes -struct ZeroCopyScanResult { - /// Whether this is a leaf node - is_leaf: bool, - /// Child index to follow (for internal nodes) - child_idx: usize, - /// Pointers array bytes start position (offset from node start) - pointers_start: usize, -} - -/// Scan an internal node to find the correct child without full deserialization. -/// -/// Node layout: -/// - `is_leaf`: 1 byte (`FLAG_SIZE`) -/// - `keys_len`: 4 bytes (`LEN_SIZE`) -/// - keys: `keys_len` bytes (`MessagePack` array) -/// - `pointers_len`: 4 bytes (`LEN_SIZE`) -/// - pointers: `pointers_len` bytes (`MessagePack` array of u64) -#[inline] -fn scan_internal_node_zero_copy(node_bytes: &[u8], search_key: &K) -> Option { - if node_bytes.len() < FLAG_SIZE + LEN_SIZE { - return None; - } - - let is_leaf = node_bytes[0] == 1; - - // Read keys_len - let keys_len_end = FLAG_SIZE + LEN_SIZE; - let keys_len = u32::from_le_bytes(node_bytes[FLAG_SIZE..keys_len_end].try_into().ok()?) as usize; - - let keys_start = FLAG_SIZE + LEN_SIZE; - let keys_end = keys_start + keys_len; - - if keys_end + LEN_SIZE > node_bytes.len() { - return None; - } - - if is_leaf { - // For leaf nodes, we can't use zero-copy for the full query - // (need to return keys for binary search). Just return that it's a leaf. - return Some(ZeroCopyScanResult { is_leaf: true, child_idx: 0, pointers_start: 0 }); - } - - // Scan keys to find child index - let keys_bytes = &node_bytes[keys_start..keys_end]; - let child_idx = find_child_index_zero_copy(keys_bytes, search_key)?; - - // Read pointers_len - let _pointers_len = u32::from_le_bytes(node_bytes[keys_end..keys_end + LEN_SIZE].try_into().ok()?) as usize; - - let pointers_start = keys_end + LEN_SIZE; - - // Parse pointer count from MessagePack array header - - Some(ZeroCopyScanResult { is_leaf: false, child_idx, pointers_start }) -} - -/// Read a specific child pointer from the pointers array without full deserialization. -#[inline] -fn read_pointer_at_index(pointers_bytes: &[u8], index: usize) -> Option { - let header = parse_msgpack_array_header(pointers_bytes)?; - - if index >= header.count { - return None; - } - - let mut pos = header.header_size; - - // Skip to the target pointer - for _ in 0..index { - if pos >= pointers_bytes.len() { - return None; - } - pos += skip_msgpack_u64(&pointers_bytes[pos..])?; - } - - // Read the target pointer - read_msgpack_u64(&pointers_bytes[pos..]) -} - -/// Skip a MessagePack-encoded u64 -#[inline] -fn skip_msgpack_u64(bytes: &[u8]) -> Option { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - if first <= 0x7f { - return Some(1); - } // fixint - if first == 0xcc { - return Some(2); - } // uint8 - if first == 0xcd { - return Some(3); - } // uint16 - if first == 0xce { - return Some(5); - } // uint32 - if first == 0xcf { - return Some(9); - } // uint64 - - None -} - -/// Read a MessagePack-encoded u64 -#[inline] -fn read_msgpack_u64(bytes: &[u8]) -> Option { - if bytes.is_empty() { - return None; - } - - let first = bytes[0]; - - // Positive fixint (0x00-0x7f) - if first <= 0x7f { - return Some(u64::from(first)); - } - - // uint8 (0xcc) - if first == 0xcc && bytes.len() >= 2 { - return Some(u64::from(bytes[1])); - } - - // uint16 (0xcd) - if first == 0xcd && bytes.len() >= 3 { - return Some(u64::from(u16::from_be_bytes([bytes[1], bytes[2]]))); - } - - // uint32 (0xce) - if first == 0xce && bytes.len() >= 5 { - return Some(u64::from(u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]))); - } - - // uint64 (0xcf) - if first == 0xcf && bytes.len() >= 9 { - return Some(u64::from_be_bytes([ - bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], bytes[8], - ])); - } - - None -} - -#[inline] -const fn msgpack_array_header_len(count: usize) -> usize { - if count <= 0x0f { - 1 - } else if count <= u16::MAX as usize { - 3 - } else { - 5 - } -} - -#[inline] -const fn msgpack_u64_array_upper_bound_len(count: usize) -> usize { - // Worst-case per u64: marker + 8 bytes payload. - msgpack_array_header_len(count) + count.saturating_mul(9) -} - -// Adaptively compress value bytes if beneficial. -// Returns borrowed raw bytes when compression is not useful to avoid an -// allocation on the common uncompressed path. -fn compress_if_beneficial(raw_bytes: &[u8]) -> (u8, Cow<'_, [u8]>) { - if raw_bytes.len() >= COMPRESSION_MIN_SIZE { - let compressed = lz4_flex::compress_prepend_size(raw_bytes); - let threshold = (raw_bytes.len() * COMPRESSION_THRESHOLD_PERCENT) / 100; - - if compressed.len() < threshold { - // Compression is effective - (COMPRESSION_FLAG_LZ4, Cow::Owned(compressed)) - } else { - // Compression not worth it - return borrowed raw bytes. - (COMPRESSION_FLAG_NONE, Cow::Borrowed(raw_bytes)) - } - } else { - // Too small to compress - return borrowed raw bytes. - (COMPRESSION_FLAG_NONE, Cow::Borrowed(raw_bytes)) - } -} - -/// Represents how a value is stored on disk -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] -enum ValueStorageMode { - /// Multiple small values packed in one block - /// (`block_offset`, `value_index_in_block`) - Packed(u64, u16), - - /// Single value in dedicated block(s) - /// (`block_offset`) - Single(u64), - - /// Entry is logically deleted. - Tombstone, -} - -#[derive(Debug, Clone)] -enum CacheData { - Compressed(u8, Vec), - PackedOffset(u16), -} - -/// Extended value info that includes storage mode and length -#[derive(Debug, Serialize, Deserialize)] -struct ValueInfo { - mode: ValueStorageMode, - length: u32, - #[serde(skip, default)] - cache: Mutex>, -} - -impl ValueInfo { - #[inline] - const fn tombstone() -> Self { Self { mode: ValueStorageMode::Tombstone, length: 0, cache: Mutex::new(None) } } - - #[inline] - const fn is_tombstone(&self) -> bool { matches!(self.mode, ValueStorageMode::Tombstone) } -} - -impl Clone for ValueInfo { - fn clone(&self) -> Self { - Self { - mode: self.mode, - length: self.length, - cache: Mutex::new(None), // Don't clone cache - } - } -} - -/// Result of attempting an in-place value update -enum InPlaceUpdateResult { - /// Update succeeded in-place, no node rewrite needed - Success, - /// Packed value was promoted to Single storage mode (node rewrite needed with new info) - PromotedToSingle(ValueInfo), - /// Value doesn't fit in existing space, need full COW - NeedsCow, -} - -#[derive(Debug, Clone)] -struct BPlusTreeNode { - keys: Vec, - children: Vec>, - is_leaf: bool, - value_info: Vec, - values: Vec, // only used in leaf nodes -} - -impl BPlusTreeNode -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - #[inline] - const fn new(is_leaf: bool) -> Self { - Self { is_leaf, keys: vec![], children: vec![], value_info: vec![], values: vec![] } - } - - #[inline] - fn is_overflow(&self, order: usize) -> bool { self.keys.len() > order } - - #[inline] - const fn get_median_index(order: usize) -> usize { order >> 1 } - - fn find_leaf_entry(node: &Self) -> Option<&K> { - if node.is_leaf { - node.keys.first() - } else if let Some(child) = node.children.first() { - Self::find_leaf_entry(child) - } else { - None - } - } - - fn query(&self, key: &K) -> Option<&V> { - if self.is_leaf { - return self.keys.binary_search(key).map_or(None, |idx| self.values.get(idx)); - } - self.children.get(self.get_entry_index_upper_bound(key))?.query(key) - } - - fn get_entry_index_upper_bound(&self, key: &K) -> usize { get_entry_index_upper_bound::(&self.keys, key) } - - fn insert(&mut self, key: K, v: V, inner_order: usize, leaf_order: usize) -> Option { - if self.is_leaf { - // Use single binary search instead of redundant searches - match self.keys.binary_search(&key) { - Ok(pos) => { - // Key exists, update value - self.values[pos] = v; - return None; - } - Err(pos) => { - // Key doesn't exist, insert at the correct position - self.keys.insert(pos, key); - self.values.insert(pos, v); - if self.is_overflow(leaf_order) { - return Some(self.split(leaf_order)); - } - } - } - } else { - let pos = self.get_entry_index_upper_bound(&key); - let child = self.children.get_mut(pos)?; - let node = child.insert(key.clone(), v, inner_order, leaf_order); - if let Some(tree_node) = node { - if let Some(leaf_key) = Self::find_leaf_entry(&tree_node) { - let idx = self.get_entry_index_upper_bound(leaf_key); - if self.keys.binary_search(leaf_key).is_err() { - self.keys.insert(idx, leaf_key.clone()); - self.children.insert(idx + 1, tree_node); - if self.is_overflow(inner_order) { - return Some(self.split(inner_order)); - } - } - } - } - } - None - } - - fn split(&mut self, order: usize) -> Self { - let median = Self::get_median_index(order); - if self.is_leaf { - let mut node = Self::new(true); - node.keys = self.keys.split_off(median); - node.values = self.values.split_off(median); - node - } else { - let mut node = Self::new(false); - node.keys = self.keys.split_off(median + 1); - node.children = self.children.split_off(median + 1); - // Internal keys are separators for children[1..]. The median key - // separates the two split nodes and is represented in the parent - // by the first leaf key of the returned right node. - let _separator = self.keys.pop(); - node - } - } - - /// Find the largest key <= `key` in this subtree. - /// Returns a reference to (key, value) if found (only valid for leaf entries). - fn find_le(&self, key: &K) -> Option<(&K, &V)> { - if self.is_leaf { - // find index of first key > key, then step one back - let idx = self.get_entry_index_upper_bound(key); - if idx == 0 { - None - } else { - let i = idx - 1; - // safe: leaf guarantees values.len() == keys.len() - Some((&self.keys[i], &self.values[i])) - } - } else { - // descend into the appropriate child (child index = upper_bound) - let child_idx = self.get_entry_index_upper_bound(key); - // child_idx can be equal to children.len() if key > all keys; children.get handles that - if let Some(child) = self.children.get(child_idx) { - child.find_le(key) - } else { - // fallback: if child_idx is out of bounds, try last child (defensive) - self.children.last().and_then(|c| c.find_le(key)) - } - } - } - - pub fn len(&self) -> usize { - if self.is_leaf { - self.keys.len() - } else { - self.children.iter().map(BPlusTreeNode::len).sum() - } - } - - pub fn traverse(&self, visit: &mut F) - where - F: FnMut(&Vec, &Vec), - { - if self.is_leaf { - visit(&self.keys, &self.values); - } - self.children.iter().for_each(|child| child.traverse(visit)); - } - - /// Write a packed value block to disk - fn write_packed_block( - file: &mut W, - buffer: &mut [u8], - offset: u64, - values: &[(u16, &[u8])], - ) -> io::Result<()> { - file.seek(SeekFrom::Start(offset))?; - - // Write count - let count = u32::try_from(values.len()).map_err(to_io_error)?; - buffer[0..4].copy_from_slice(&count.to_le_bytes()); - let mut pos = 4; - - // Write each value: length + data - for (_, value_bytes) in values { - let len = u32::try_from(value_bytes.len()).map_err(to_io_error)?; - buffer[pos..pos + 4].copy_from_slice(&len.to_le_bytes()); - pos += 4; - buffer[pos..pos + value_bytes.len()].copy_from_slice(value_bytes); - pos += value_bytes.len(); - } - - // Zero remaining space - if pos < PAGE_SIZE_USIZE { - buffer[pos..PAGE_SIZE_USIZE].fill(0u8); - } - - file.write_all(&buffer[..PAGE_SIZE_USIZE])?; - Ok(()) - } - - /// Calculate the serialized size of this node in bytes (rounded up to block size) - fn calculate_serialized_size(&self, serial_buf: &mut Vec) -> io::Result { - serial_buf.clear(); - - // Header: is_leaf flag - let mut size = FLAG_SIZE; - - // Keys: length + serialized data - binary_serialize_into(&mut *serial_buf, &self.keys)?; - size += LEN_SIZE + serial_buf.len(); - - if self.is_leaf { - // Leaf nodes now store value_info instead of values - // value_info: length + Vec<(u64, u32)> - // Reuse buf - serial_buf.clear(); - binary_serialize_into(&mut *serial_buf, &self.value_info)?; - size += LEN_SIZE + serial_buf.len(); - } else { - // Internal node: pointer length + pointers - // Pointer encoding is variable-length. Using small placeholder values - // can underestimate node size and cause offset overlap. - size += LEN_SIZE + msgpack_u64_array_upper_bound_len(self.children.len()); - } - - // Round up to block size - let blocks = size.div_ceil(PAGE_SIZE_USIZE); - Ok((blocks * PAGE_SIZE_USIZE) as u64) - } - - fn serialize_to_block( - &self, - file: &mut W, - buffer: &mut Vec, - serial_buf: &mut Vec, - offset: u64, - ) -> io::Result { - serial_buf.clear(); - binary_serialize_into(&mut *serial_buf, &self.keys)?; - let keys_len = u32::try_from(serial_buf.len()).map_err(to_io_error)?; - - if self.is_leaf { - let keys_end = serial_buf.len(); - // Append info_encoded to serial_buf to avoid second allocation - binary_serialize_into(&mut *serial_buf, &self.value_info)?; - let info_len = u32::try_from(serial_buf.len() - keys_end).map_err(to_io_error)?; - let info_slice = &serial_buf[keys_end..]; - - let content_size = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + info_len as usize; - let blocks = content_size.div_ceil(PAGE_SIZE_USIZE); - let capacity = blocks * PAGE_SIZE_USIZE; - debug_assert!( - content_size <= capacity, - "Leaf node content ({content_size}B) exceeds allocated capacity ({capacity}B)" - ); - - file.seek(SeekFrom::Start(offset))?; - - if buffer.len() < capacity { - buffer.resize(capacity, 0); - } - buffer[..capacity].fill(0); - - let mut pos = 0; - buffer[pos] = 1u8; - pos += FLAG_SIZE; - - buffer[pos..pos + LEN_SIZE].copy_from_slice(&keys_len.to_le_bytes()); - pos += LEN_SIZE; - - buffer[pos..pos + keys_len as usize].copy_from_slice(&serial_buf[0..keys_len as usize]); - pos += keys_len as usize; - - buffer[pos..pos + LEN_SIZE].copy_from_slice(&info_len.to_le_bytes()); - pos += LEN_SIZE; - - buffer[pos..pos + info_len as usize].copy_from_slice(info_slice); - - file.write_all(&buffer[..capacity])?; - - Ok(offset + (blocks as u64 * PAGE_SIZE_USIZE as u64)) - } else { - let ptr_count = self.children.len(); - // Conservative upper bound for MessagePack-encoded Vec. - // Must not underestimate, otherwise child blocks can overlap. - let ptr_encoded_size = msgpack_u64_array_upper_bound_len(ptr_count); - - let content_size = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + ptr_encoded_size; - let blocks_needed = content_size.div_ceil(PAGE_SIZE_USIZE); - - let parent_start = offset; - let mut current_offset = parent_start + (blocks_needed as u64 * PAGE_SIZE_USIZE as u64); - - let mut pointers = Vec::with_capacity(ptr_count); - for child in &self.children { - pointers.push(current_offset); - let mut child_scratch = Vec::new(); // Separate scratch for recursion to protect our serial_buf - current_offset = child.serialize_to_block(file, buffer, &mut child_scratch, current_offset)?; - } - - // Append pointers to serial_buf - let keys_end = serial_buf.len(); - binary_serialize_into(&mut *serial_buf, &pointers)?; - let pointers_len = u32::try_from(serial_buf.len() - keys_end).map_err(to_io_error)?; - let pointers_slice = &serial_buf[keys_end..]; - - file.seek(SeekFrom::Start(parent_start))?; - - let total_capacity = blocks_needed * PAGE_SIZE_USIZE; - let actual_content = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + pointers_len as usize; - debug_assert!( - actual_content <= total_capacity, - "Internal node content ({actual_content}B) exceeds allocated capacity ({total_capacity}B)" - ); - if buffer.len() < total_capacity { - buffer.resize(total_capacity, 0); - } - buffer[..total_capacity].fill(0); - - let mut pos = 0; - // Is_leaf=0 - buffer[pos] = 0u8; - pos += FLAG_SIZE; - - buffer[pos..pos + LEN_SIZE].copy_from_slice(&keys_len.to_le_bytes()); - pos += LEN_SIZE; - buffer[pos..pos + keys_len as usize].copy_from_slice(&serial_buf[0..keys_len as usize]); - pos += keys_len as usize; - - buffer[pos..pos + LEN_SIZE].copy_from_slice(&pointers_len.to_le_bytes()); - pos += LEN_SIZE; - buffer[pos..pos + pointers_len as usize].copy_from_slice(pointers_slice); - - file.write_all(&buffer[..total_capacity])?; - - Ok(current_offset) - } - } - - /// Serialize the tree in breadth-first order for better disk locality - /// This improves query performance by keeping nodes at the same level contiguous - fn serialize_breadth_first( - &mut self, - file: &mut W, - buffer: &mut Vec, - start_offset: u64, - ) -> io::Result { - let mut serial_buf = Vec::with_capacity(PAGE_SIZE_USIZE); - - self.serialize_bfs_pass1_populate_value_info(&mut serial_buf)?; - let (node_offsets, child_ids_by_node, current_offset) = - self.serialize_bfs_pass2_calculate_offsets(&mut serial_buf, start_offset)?; - self.serialize_bfs_pass3_assign_value_offsets(current_offset); - self.serialize_bfs_pass4_write_nodes(file, buffer, &mut serial_buf, &node_offsets, &child_ids_by_node)?; - self.serialize_bfs_pass5_write_values(file, buffer, &mut serial_buf)?; - - Ok(start_offset) - } - - fn serialize_bfs_pass1_populate_value_info(&mut self, serial_buf: &mut Vec) -> io::Result<()> { - let mut current_level_mut = vec![self]; - while !current_level_mut.is_empty() { - let mut next_level_mut = Vec::new(); - for node in current_level_mut { - if node.is_leaf { - node.value_info.clear(); - let mut serialized_values: Vec> = Vec::new(); - for value in &node.values { - serial_buf.clear(); - binary_serialize_into(serial_buf, value)?; - serialized_values.push(serial_buf.clone()); - } - - let mut current_pack_index: u16 = 0; - let mut current_pack_size = PACK_BLOCK_HEADER_SIZE; - let mut pack_count = 0u32; - - for value_bytes in serialized_values { - let size = value_bytes.len(); - - if size <= SMALL_VALUE_THRESHOLD { - let entry_size = PACK_VALUE_HEADER_SIZE + size; - - if current_pack_size + entry_size <= PAGE_SIZE_USIZE { - node.value_info.push(ValueInfo { - mode: ValueStorageMode::Packed(u64::from(pack_count), current_pack_index), - length: u32::try_from(size).map_err(to_io_error)?, - cache: Mutex::new(None), - }); - current_pack_index += 1; - current_pack_size += entry_size; - } else { - pack_count += 1; - current_pack_index = 1; - current_pack_size = PACK_BLOCK_HEADER_SIZE + entry_size; - - node.value_info.push(ValueInfo { - mode: ValueStorageMode::Packed(u64::from(pack_count), 0), - length: u32::try_from(size).map_err(to_io_error)?, - cache: Mutex::new(None), - }); - } - } else { - let (flag, payload) = compress_if_beneficial(&value_bytes); - let stored_size = 1 + payload.len(); - - let cache = if flag == COMPRESSION_FLAG_LZ4 { - Some(CacheData::Compressed(flag, payload.into_owned())) - } else { - None - }; - - node.value_info.push(ValueInfo { - mode: ValueStorageMode::Single(u64::MAX), - length: u32::try_from(stored_size).map_err(to_io_error)?, - cache: Mutex::new(cache), - }); - } - } - } else { - for child in &mut node.children { - next_level_mut.push(child); - } - } - } - current_level_mut = next_level_mut; - } - Ok(()) - } - - fn serialize_bfs_pass2_calculate_offsets( - &self, - serial_buf: &mut Vec, - start_offset: u64, - ) -> io::Result<(Vec, Vec>, u64)> { - let mut node_refs: Vec<&BPlusTreeNode> = vec![self]; - let mut node_offsets: Vec = vec![start_offset]; - let mut child_ids_by_node: Vec> = vec![Vec::new()]; - let mut current_offset = start_offset + self.calculate_serialized_size(serial_buf)?; - let mut current_level = vec![0usize]; - - while !current_level.is_empty() { - let mut next_level = Vec::new(); - for node_id in current_level { - let node = node_refs[node_id]; - if !node.is_leaf { - for child in &node.children { - let child_id = node_refs.len(); - node_refs.push(child); - node_offsets.push(current_offset); - child_ids_by_node.push(Vec::new()); - child_ids_by_node[node_id].push(child_id); - current_offset += child.calculate_serialized_size(serial_buf)?; - next_level.push(child_id); - } - } - } - current_level = next_level; - } - - Ok((node_offsets, child_ids_by_node, current_offset)) - } - - fn serialize_bfs_pass3_assign_value_offsets(&mut self, mut current_offset: u64) { - use std::collections::HashMap; - let mut current_level_mut = vec![self]; - while !current_level_mut.is_empty() { - let mut next_level_mut = Vec::new(); - for node in current_level_mut { - if node.is_leaf { - let mut pack_block_offsets: HashMap = HashMap::new(); - - for info in &mut node.value_info { - match &mut info.mode { - ValueStorageMode::Packed(pack_idx, _index) => { - if !pack_block_offsets.contains_key(pack_idx) { - pack_block_offsets.insert(*pack_idx, current_offset); - current_offset += PAGE_SIZE_USIZE as u64; - } - } - ValueStorageMode::Single(offset) if *offset == u64::MAX => { - *offset = current_offset; - current_offset += u64::from(info.length); - } - ValueStorageMode::Single(_) | ValueStorageMode::Tombstone => {} - } - } - - for info in &mut node.value_info { - if let ValueStorageMode::Packed(pack_idx, _index) = &mut info.mode { - let actual_offset = pack_block_offsets[pack_idx]; - *pack_idx = actual_offset; - } - } - } else { - for child in &mut node.children { - next_level_mut.push(child); - } - } - } - current_level_mut = next_level_mut; - } - } - - fn serialize_bfs_pass4_write_nodes( - &self, - file: &mut W, - buffer: &mut Vec, - serial_buf: &mut Vec, - node_offsets: &[u64], - child_ids_by_node: &[Vec], - ) -> io::Result<()> { - let mut node_refs: Vec<&BPlusTreeNode> = vec![self]; - let mut node_cursor = 0; - while node_cursor < node_refs.len() { - let node = node_refs[node_cursor]; - if !node.is_leaf { - node_refs.extend(node.children.iter()); - } - node_cursor += 1; - } - - if node_refs.len() != node_offsets.len() || node_refs.len() != child_ids_by_node.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "B+Tree serialization produced inconsistent node offset table", - )); - } - - for (node_id, node) in node_refs.iter().enumerate() { - let node_offset = node_offsets[node_id]; - - if node.is_leaf { - node.serialize_to_block(file, buffer, serial_buf, node_offset)?; - } else { - let node_child_ids = child_ids_by_node.get(node_id).ok_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidData, - "B+Tree serialization missing child id table entry", - ) - })?; - let mut child_offsets = Vec::with_capacity(node_child_ids.len()); - for child_id in node_child_ids { - let Some(child_offset) = node_offsets.get(*child_id) else { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "B+Tree serialization child id has no offset", - )); - }; - child_offsets.push(*child_offset); - } - - node.serialize_internal_with_offsets( - file, - buffer, - serial_buf, - node_offset, - &child_offsets, - )?; - } - } - Ok(()) - } - - fn serialize_bfs_pass5_write_values( - &self, - file: &mut W, - buffer: &mut [u8], - serial_buf: &mut Vec, - ) -> io::Result<()> { - use std::collections::HashMap; - let mut current_level_values = vec![self]; - while !current_level_values.is_empty() { - let mut next_level = Vec::new(); - for node in current_level_values { - if node.is_leaf { - let mut pack_blocks: HashMap)>> = HashMap::new(); - - for (value, info) in node.values.iter().zip(node.value_info.iter()) { - serial_buf.clear(); - binary_serialize_into(serial_buf, value)?; - - match info.mode { - ValueStorageMode::Packed(block_offset, index) => { - pack_blocks.entry(block_offset).or_default().push((index, serial_buf.clone())); - } - ValueStorageMode::Single(block_offset) => { - file.seek(SeekFrom::Start(block_offset))?; - - let cache_guard = info.cache.lock(); - let (flag, payload_ref) = - if let Some(CacheData::Compressed(c_flag, c_payload)) = cache_guard.as_ref() { - (*c_flag, c_payload.as_slice()) - } else { - (COMPRESSION_FLAG_NONE, serial_buf.as_slice()) - }; - - file.write_all(&[flag])?; - file.write_all(payload_ref)?; - } - ValueStorageMode::Tombstone => {} - } - } - - for (block_offset, mut values) in pack_blocks { - values.sort_by_key(|(idx, _)| *idx); - let value_refs: Vec<(u16, &[u8])> = - values.iter().map(|(idx, bytes)| (*idx, bytes.as_slice())).collect(); - Self::write_packed_block(file, buffer, block_offset, &value_refs)?; - } - } else { - for child in &node.children { - next_level.push(child); - } - } - } - current_level_values = next_level; - } - Ok(()) - } - - /// Serialize an internal node with pre-calculated child offsets - /// Supports multi-block internal nodes when keys + pointers exceed a single page - fn serialize_internal_with_offsets( - &self, - file: &mut W, - buffer: &mut Vec, - serial_buf: &mut Vec, - offset: u64, - child_offsets: &[u64], - ) -> io::Result { - // Similar to serialize_to_block but for internal nodes with known child offsets - serial_buf.clear(); - binary_serialize_into(&mut *serial_buf, &self.keys)?; - let keys_len = serial_buf.len(); - let keys_end = keys_len; - - binary_serialize_into(&mut *serial_buf, child_offsets)?; - let pointer_len = serial_buf.len() - keys_end; - - // Calculate total content size - let total_content_size = FLAG_SIZE + LEN_SIZE + keys_len + LEN_SIZE + pointer_len; - let blocks_needed = total_content_size.div_ceil(PAGE_SIZE_USIZE); - - let total_buffer_size = blocks_needed * PAGE_SIZE_USIZE; - if buffer.len() < total_buffer_size { - buffer.resize(total_buffer_size, 0); - } - buffer[..total_buffer_size].fill(0); - - let mut write_pos = 0; - - // Write is_leaf flag (0 for internal node) - buffer[write_pos] = u8::from(self.is_leaf); - write_pos += FLAG_SIZE; - - // Write keys length and data - buffer[write_pos..write_pos + LEN_SIZE] - .copy_from_slice(&u32::try_from(keys_len).map_err(to_io_error)?.to_le_bytes()); - write_pos += LEN_SIZE; - buffer[write_pos..write_pos + keys_len].copy_from_slice(&serial_buf[0..keys_end]); - write_pos += keys_len; - - // Write pointers length and data - buffer[write_pos..write_pos + LEN_SIZE] - .copy_from_slice(&u32::try_from(pointer_len).map_err(to_io_error)?.to_le_bytes()); - write_pos += LEN_SIZE; - buffer[write_pos..write_pos + pointer_len].copy_from_slice(&serial_buf[keys_end..]); - - // Write all blocks to file - file.seek(SeekFrom::Start(offset))?; - file.write_all(&buffer[..total_buffer_size])?; - - Ok(offset + total_buffer_size as u64) - } - - fn deserialize_from_block( - file: &mut R, - buffer: &mut Vec, - offset: u64, - nested: bool, - ) -> io::Result<(Self, Option>)> { - file.seek(SeekFrom::Start(offset))?; - - let header_required = FLAG_SIZE + LEN_SIZE; - if buffer.len() < header_required { - buffer.resize(header_required, 0); - } - - file.read_exact(&mut buffer[0..header_required])?; - - let is_leaf = node_flag_to_is_leaf(buffer[0])?; - #[allow(clippy::range_plus_one)] - let keys_len = u32_from_bytes(&buffer[FLAG_SIZE..FLAG_SIZE + LEN_SIZE])? as usize; - - let min_required = header_required + keys_len + LEN_SIZE; - if buffer.len() < min_required { - buffer.resize(min_required, 0); - } - - file.read_exact(&mut buffer[header_required..min_required])?; - - let mut read_pos = header_required; - let mut keys: Vec = binary_deserialize(&buffer[read_pos..read_pos + keys_len])?; - read_pos += keys_len; - - let payload_len = u32_from_bytes(&buffer[read_pos..read_pos + LEN_SIZE])? as usize; - read_pos += LEN_SIZE; - - let total_required = min_required + payload_len; - if buffer.len() < total_required { - buffer.resize(total_required, 0); - } - - file.read_exact(&mut buffer[min_required..total_required])?; - - let (value_info, values, children, children_pointer) = if is_leaf { - let mut info: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; - if info.len() != keys.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid leaf node: {} keys but {} value descriptors", keys.len(), info.len()), - )); - } - let vals = if nested { - let mut filtered_keys: Vec = Vec::with_capacity(keys.len()); - let mut filtered_info: Vec = Vec::with_capacity(info.len()); - let mut v = Vec::with_capacity(info.len()); - - let original_keys = std::mem::take(&mut keys); - for (entry_key, entry_info) in original_keys.into_iter().zip(info) { - if entry_info.is_tombstone() { - continue; - } - v.push(Self::load_value_from_info(file, &entry_info)?); - filtered_keys.push(entry_key); - filtered_info.push(entry_info); - } - - keys = filtered_keys; - info = filtered_info; - v - } else { - Vec::new() - }; - (info, vals, Vec::new(), None) - } else { - let pointers: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; - if pointers.len() != keys.len() + 1 { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid internal node: {} keys but {} child pointers", keys.len(), pointers.len()), - )); - } - let nodes = if nested { - let mut n = Vec::with_capacity(pointers.len()); - let mut child_buf = Vec::with_capacity(PAGE_SIZE_USIZE); - for &ptr in &pointers { - let (child, _) = Self::deserialize_from_block(file, &mut child_buf, ptr, nested)?; - n.push(child); - } - n - } else { - Vec::new() - }; - (Vec::new(), Vec::new(), nodes, Some(pointers)) - }; - - Ok((Self { keys, children, is_leaf, value_info, values }, children_pointer)) - } - - fn deserialize_from_mmap( - mmap: &[u8], - file: &mut R, - offset: u64, - nested: bool, - ) -> io::Result<(Self, Option>)> { - let start = usize::try_from(offset).map_err(to_io_error)?; - let header_end = start - .checked_add(FLAG_SIZE + LEN_SIZE) - .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; - // Basic safety check for mmap bounds - if header_end > mmap.len() { - return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); - } - - let keys_len = u32_from_bytes(&mmap[start + FLAG_SIZE..start + FLAG_SIZE + LEN_SIZE])? as usize; - let keys_start = header_end; - let len_pos = keys_start - .checked_add(keys_len) - .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; - - if len_pos + LEN_SIZE > mmap.len() { - return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); - } - let payload_len = u32_from_bytes(&mmap[len_pos..len_pos + LEN_SIZE])? as usize; - let total = len_pos - .checked_add(LEN_SIZE + payload_len) - .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; - - if total > mmap.len() { - return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); - } - - // We need to know the total size of the node to slice the mmap - // For simplicity, we can just slice a PAGE_SIZE or slightly more if we know it overflows. - // Actually, our serialize_to_block uses PAGE_SIZE blocks. - - //let slice = &mmap[start..]; - let slice = &mmap[start..total]; - Self::deserialize_from_block_slice(slice, Some(mmap), file, nested) - } - - fn deserialize_from_block_slice( - slice: &[u8], - mmap: Option<&[u8]>, - file: &mut R, - nested: bool, - ) -> io::Result<(Self, Option>)> { - let header_range = checked_slice_range(0, FLAG_SIZE + LEN_SIZE, slice.len())?; - // Node type - let is_leaf = node_flag_to_is_leaf(slice[0])?; - let mut read_pos = FLAG_SIZE; - - // ---- Keys ---- - let keys_length = u32_from_bytes(&slice[read_pos..header_range.end])? as usize; - read_pos += LEN_SIZE; - let keys_range = checked_slice_range(read_pos, keys_length, slice.len())?; - let mut keys: Vec = binary_deserialize(&slice[keys_range.clone()])?; - read_pos = keys_range.end; - - // ---- Value info (offset, length) for leaf nodes ---- - let (value_info, values): (Vec, Vec) = if is_leaf { - // Read value_info - let info_len_range = checked_slice_range(read_pos, LEN_SIZE, slice.len())?; - let info_length = u32_from_bytes(&slice[info_len_range.clone()])? as usize; - read_pos = info_len_range.end; - let info_range = checked_slice_range(read_pos, info_length, slice.len())?; - let mut info: Vec = binary_deserialize(&slice[info_range])?; - if info.len() != keys.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid leaf node: {} keys but {} value descriptors", keys.len(), info.len()), - )); - } - - // Values are loaded on-demand when nested=true - if nested { - let mut vals = Vec::with_capacity(info.len()); - let mut filtered_keys: Vec = Vec::with_capacity(keys.len()); - let mut filtered_info: Vec = Vec::with_capacity(info.len()); - let mut last_packed_block: Option<(u64, Vec)> = None; - let original_keys = std::mem::take(&mut keys); - for (entry_key, entry_info) in original_keys.into_iter().zip(info) { - if entry_info.is_tombstone() { - continue; - } - match entry_info.mode { - ValueStorageMode::Packed(block_offset, index) => { - // Packed loading optimization: reuse block if it's the same - if let Some((offset, ref block)) = last_packed_block { - if offset == block_offset { - vals.push(Self::extract_value_from_packed_block(block, index, &entry_info.cache)?); - filtered_keys.push(entry_key); - filtered_info.push(entry_info); - continue; - } - } - - // Load new block - let mut block = vec![0u8; PAGE_SIZE_USIZE]; - file.seek(SeekFrom::Start(block_offset))?; - file.read_exact(&mut block)?; - vals.push(Self::extract_value_from_packed_block(&block, index, &entry_info.cache)?); - last_packed_block = Some((block_offset, block)); - filtered_keys.push(entry_key); - filtered_info.push(entry_info); - } - ValueStorageMode::Single(_) => { - last_packed_block = None; - vals.push(Self::load_value_from_info(file, &entry_info)?); - filtered_keys.push(entry_key); - filtered_info.push(entry_info); - } - ValueStorageMode::Tombstone => {} - } - } - keys = filtered_keys; - info = filtered_info; - (info, vals) - } else { - (info, Vec::new()) - } - } else { - (Vec::new(), Vec::new()) - }; - - // ---- Pointers for internal nodes ---- - let (children, children_pointer): (Vec, Option>) = if is_leaf { - (Vec::new(), None) - } else { - let pointers_len_range = checked_slice_range(read_pos, LEN_SIZE, slice.len())?; - let pointers_length = u32_from_bytes(&slice[pointers_len_range.clone()])? as usize; - read_pos = pointers_len_range.end; - let pointers_range = checked_slice_range(read_pos, pointers_length, slice.len())?; - let pointers: Vec = binary_deserialize(&slice[pointers_range])?; - if pointers.len() != keys.len() + 1 { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid internal node: {} keys but {} child pointers", keys.len(), pointers.len()), - )); - } - if nested { - let mut nodes = Vec::with_capacity(pointers.len()); - let mut child_buffer = vec![0u8; PAGE_SIZE_USIZE]; - for &ptr in &pointers { - let (child, _) = if let Some(m) = mmap { - Self::deserialize_from_mmap(m, file, ptr, nested)? - } else { - Self::deserialize_from_block(file, &mut child_buffer, ptr, nested)? - }; - nodes.push(child); - } - (nodes, None) - } else { - (Vec::new(), Some(pointers)) - } - }; - - Ok((Self { keys, children, is_leaf, value_info, values }, children_pointer)) - } - - /// Load a value based on its storage info - fn load_value_from_info(file: &mut R, info: &ValueInfo) -> io::Result { - // Fast path: Check cache for Single mode - if let ValueStorageMode::Single(_) = info.mode { - let cache_guard = info.cache.lock(); - if let Some(CacheData::Compressed(flag, payload)) = cache_guard.as_ref() { - if *flag == COMPRESSION_FLAG_LZ4 { - let decompressed = lz4_flex::decompress_size_prepended(payload).map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 cache decompression failed: {e}")) - })?; - return binary_deserialize(&decompressed); - } - return binary_deserialize(payload); - } - } - - match info.mode { - ValueStorageMode::Single(offset) => { - let stored_len = info.length as usize; - if stored_len < 1 { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid value length")); - } - - // Read everything: flag + payload - file.seek(SeekFrom::Start(offset))?; - let mut buffer = vec![0u8; stored_len]; - file.read_exact(&mut buffer)?; - - let flag = buffer[0]; - // Split payload without re-allocating if possible? Vec::split_off allocates new vec for tail. - // We want payload as Vec for cache. - let payload = buffer[1..].to_vec(); - - // Decompress for result - let data = if flag == COMPRESSION_FLAG_LZ4 { - lz4_flex::decompress_size_prepended(&payload).map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 decompression failed: {e}")) - })? - } else { - payload.clone() - }; - - // Update cache - *info.cache.lock() = Some(CacheData::Compressed(flag, payload)); - - binary_deserialize(&data) - } - ValueStorageMode::Packed(block_offset, index) => { - Self::load_value_from_packed_block(file, block_offset, index, info.length, &info.cache) - } - ValueStorageMode::Tombstone => { - Err(io::Error::new(io::ErrorKind::NotFound, "value was deleted (tombstone)")) - } - } - } - - /// Load a value from a packed block - fn load_value_from_packed_block( - file: &mut R, - block_offset: u64, - value_index: u16, - _expected_length: u32, - cache: &Mutex>, - ) -> io::Result { - file.seek(SeekFrom::Start(block_offset))?; - - let mut block_buffer = vec![0u8; PAGE_SIZE_USIZE]; - file.read_exact(&mut block_buffer)?; - - Self::extract_value_from_packed_block(&block_buffer, value_index, cache) - } - - /// Helper to extract value from a packed block that is already in memory - fn extract_value_from_packed_block( - block_buffer: &[u8], - value_index: u16, - cache: &Mutex>, - ) -> io::Result { - // Read count - if block_buffer.len() < 4 { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block too small")); - } - let mut pos = 4; - - // Skip to target value - for i in 0..=value_index { - if pos + 4 > block_buffer.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Packed block corrupted: position {pos} exceeds block size"), - )); - } - - let len = u32::from_le_bytes(block_buffer[pos..pos + 4].try_into().map_err(to_io_error)?) as usize; - pos += 4; - - if i == value_index { - // Found target value - if pos + len > block_buffer.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Packed value corrupted: length {len} at position {pos} exceeds block size"), - )); - } - - // Cache the position in the block for future in-place updates - *cache.lock() = Some(CacheData::PackedOffset(u16::try_from(pos).map_err(to_io_error)?)); - - let value_data = &block_buffer[pos..pos + len]; - return binary_deserialize(value_data); - } - - pos += len; - } - - Err(io::Error::new(io::ErrorKind::InvalidData, format!("Value index {value_index} not found in packed block"))) - } -} - -// ----------------------------------------------------------------------------- -// Metadata Enum -// ----------------------------------------------------------------------------- -#[derive(Clone, Debug, PartialEq)] -pub enum BPlusTreeMetadata { - Empty, - TargetIdMapping(u32), -} - -impl BPlusTreeMetadata { - pub fn to_bytes(&self) -> Vec { - match self { - Self::Empty => Vec::new(), - Self::TargetIdMapping(val) => { - let mut bytes = vec![MAGIC_METADATA_TARGET_ID_MAPPING]; // Type tag - bytes.extend_from_slice(&val.to_le_bytes()); - bytes - } - } - } - - pub fn from_bytes(bytes: &[u8]) -> Self { - match bytes.len() { - 5 if bytes[0] == MAGIC_METADATA_TARGET_ID_MAPPING => { - let arr: [u8; 4] = bytes[1..5].try_into().unwrap_or([0; 4]); - Self::TargetIdMapping(u32::from_le_bytes(arr)) - } - _ => Self::Empty, // Unknown metadata treated as Empty for now - } - } -} - -#[derive(Debug, Clone)] -pub struct BPlusTree { - root: BPlusTreeNode, - inner_order: usize, - leaf_order: usize, - metadata: BPlusTreeMetadata, - dirty: bool, -} - -const fn sanitize_order(order: usize) -> usize { - if order < 2 { - 2 - } else { - order - } -} - -const fn default_orders() -> (usize, usize) { (DEFAULT_INNER_ORDER, DEFAULT_LEAF_ORDER) } - -impl Default for BPlusTree -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - fn default() -> Self { Self::new() } -} - -impl BPlusTree -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - pub const fn new() -> Self { - let (inner_order, leaf_order) = default_orders(); - Self::new_with_orders(inner_order, leaf_order) - } - - /// Create a v2 tree with explicit in-memory fanout. - /// - /// This does not change the on-disk v2 format. Orders below 2 are clamped - /// because B+Tree split logic requires at least two keys per node. - pub const fn new_with_orders(inner_order: usize, leaf_order: usize) -> Self { - Self { - root: BPlusTreeNode::::new(true), - inner_order: sanitize_order(inner_order), - leaf_order: sanitize_order(leaf_order), - metadata: BPlusTreeMetadata::Empty, - dirty: true, // an empty tree is stored! - } - } - - /// Helper to access metadata - pub fn get_metadata(&self) -> &BPlusTreeMetadata { &self.metadata } - - /// Helper to set metadata - pub fn set_metadata(&mut self, data: BPlusTreeMetadata) { - self.metadata = data; - self.dirty = true; - } - - pub fn is_empty(&self) -> bool { self.root.keys.is_empty() } - - pub fn len(&self) -> usize { self.root.len() } - - pub fn insert(&mut self, key: K, value: V) { - self.dirty = true; - if self.root.keys.is_empty() { - self.root.keys.push(key); - self.root.values.push(value); - return; - } - - if let Some(node) = self.root.insert(key, value, self.inner_order, self.leaf_order) { - let child_key_opt = - if node.is_leaf { node.keys.first() } else { BPlusTreeNode::::find_leaf_entry(&node) }; - - if let Some(child_key) = child_key_opt { - let mut new_root = BPlusTreeNode::::new(false); - new_root.keys.push(child_key.clone()); - new_root.children.push(std::mem::replace(&mut self.root, BPlusTreeNode::new(true))); - new_root.children.push(node); - - self.root = new_root; - } else { - error!("Failed to insert child key"); - } - } - } - - pub fn query(&self, key: &K) -> Option<&V> { self.root.query(key) } - - pub fn store(&mut self, filepath: &Path) -> io::Result { - if self.dirty { - // Advisory lock to prevent concurrent COW updates - let _lock = FileLock::try_lock(filepath)?; - self.store_internal(filepath) - } else { - Ok(0) - } - } - - /// Store the tree and build a sorted index file. - /// - /// # Arguments - /// * `filepath` - Path to store the `BPlusTree` - /// * `sort_key_extractor` - Closure that extracts the sort key from a value - /// - /// # Example - /// ```ignore - /// tree.store_with_index(&db_path, |v| v.name.clone())?; - /// ``` - pub fn store_with_index(&mut self, filepath: &Path, sort_key_extractor: F) -> io::Result - where - SortKey: Ord + Serialize, - F: Fn(&V) -> SortKey, - { - // Store the tree first - let result = self.store(filepath)?; - if result > 0 { - Self::store_index(filepath, sort_key_extractor)?; - } - - Ok(result) - } - - pub fn store_index(filepath: &Path, sort_key_extractor: F) -> io::Result<()> - where - SortKey: Ord + Serialize, - F: Fn(&V) -> SortKey, - { - let index_path = get_file_path_for_db_index(filepath); - - // Re-open the stored tree to get value locations - let mut query = BPlusTreeQuery::::try_new(filepath)?; - let entries_with_locations = query.collect_with_locations()?; - - // Collect (sort_key, primary_key, location) and sort - let mut sorted_entries: Vec<(SortKey, K, super::sorted_index::ValueLocation)> = - entries_with_locations.into_iter().map(|(k, v, loc)| (sort_key_extractor(&v), k, loc)).collect(); - - // Sort by sort key - sorted_entries.sort_by(|a, b| a.0.cmp(&b.0)); - - // Write index file - let mut writer = super::sorted_index::SortedIndexWriter::new(&index_path)?; - for (sort_key, primary_key, location) in &sorted_entries { - writer.push(sort_key, primary_key, *location)?; - } - writer.finish()?; - Ok(()) - } - - /// Internal store without locking, used for compaction or initial save. - fn store_internal(&mut self, filepath: &Path) -> io::Result { - let tempfile = if let Some(parent_dir) = filepath.parent() { - if let Ok(file) = NamedTempFile::new_in(parent_dir) { - file - } else { - let temp_dir = tempfile::env::temp_dir(); - NamedTempFile::new_in(&temp_dir)? - } - } else { - let temp_dir = tempfile::env::temp_dir(); - NamedTempFile::new_in(&temp_dir)? - }; - let mut file = utils::file_writer(&tempfile); - let mut buffer = vec![0u8; PAGE_SIZE_USIZE]; - - // Write header block 0 - let mut header = [0u8; PAGE_SIZE_USIZE]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&STORAGE_VERSION.to_le_bytes()); - // Placeholder for root offset, will be updated after serialization - header[8..16].copy_from_slice(&HEADER_SIZE.to_le_bytes()); - - let meta_bytes = self.metadata.to_bytes(); - if meta_bytes.len() > METADATA_MAX_SIZE || METADATA_DATA_START_POS + meta_bytes.len() > PAGE_SIZE_USIZE { - return Err(io::Error::new(io::ErrorKind::InvalidInput, "Metadata too large for header page")); - } - let metadata_len = - u32::try_from(meta_bytes.len()).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?; - let metadata_len_with_flags = encode_metadata_len_with_flags(metadata_len, false); - header[16..20].copy_from_slice(&metadata_len_with_flags.to_le_bytes()); - if !meta_bytes.is_empty() { - header[METADATA_DATA_START_POS..METADATA_DATA_START_POS + meta_bytes.len()].copy_from_slice(&meta_bytes); - } - - file.write_all(&header)?; - - // We need to ensure we pad to PAGE_SIZE before continuing - file.seek(SeekFrom::Start(HEADER_SIZE))?; - - // Use breadth-first serialization for better disk locality - match self.root.serialize_breadth_first(&mut file, &mut buffer, HEADER_SIZE) { - Ok(root_offset) => { - // Update root offset in header - file.seek(SeekFrom::Start(ROOT_OFFSET_POS))?; - file.write_all(&root_offset.to_le_bytes())?; - - file.flush()?; - drop(file); - if let Err(err) = utils::rename_or_copy(tempfile.path(), filepath, false) { - return Err(string_to_io_error(format!( - "Temp file rename/copy did not work {} {err}", - tempfile.path().to_string_lossy() - ))); - } - self.dirty = false; - Ok(root_offset) - } - Err(err) => Err(err), - } - } - - /// Bulk build a tree from pre-calculated `ValueInfos` (streaming compact helper). - /// Writes nodes to `file` starting at `start_offset`. - /// Returns the offset of the root node used to update the file header. - fn build_levels_from_pointers( - &self, - file: &mut W, - mut next_level_pointers: Vec<(K, u64)>, - mut current_offset: u64, - write_buffer: &mut Vec, - ) -> io::Result { - if next_level_pointers.is_empty() { - return Ok(current_offset); - } - - while next_level_pointers.len() > 1 { - let mut parent_level_pointers: Vec<(K, u64)> = Vec::new(); - let children = next_level_pointers; - - for chunk in children.chunks(self.inner_order) { - let mut node = BPlusTreeNode::::new(false); - let mut pointers = Vec::new(); - - if let Some((_, off)) = chunk.first() { - pointers.push(*off); - } - - for (k, off) in &chunk[1..] { - node.keys.push(k.clone()); - pointers.push(*off); - } - - let node_offset = current_offset; - let mut serial_buf = Vec::new(); - current_offset = - node.serialize_internal_with_offsets(file, write_buffer, &mut serial_buf, node_offset, &pointers)?; - - if let Some((k, _)) = chunk.first() { - parent_level_pointers.push((k.clone(), node_offset)); - } - } - next_level_pointers = parent_level_pointers; - } - - if let Some((_, root_off)) = next_level_pointers.first() { - Ok(*root_off) - } else { - Ok(current_offset) - } - } - - pub fn load(filepath: &Path) -> io::Result { - let file = File::open(filepath)?; - let mmap = mmap_with_advice(&file, Advice::Sequential, "full tree load") - .ok_or_else(|| io::Error::other("Failed to mmap B+Tree file"))?; - - if mmap.len() < PAGE_SIZE_USIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, "File too small")); - } - - // Verify Header - let header = &mmap[0..PAGE_SIZE_USIZE]; - if &header[0..4] != MAGIC { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid magic number")); - } - let version = u32_from_bytes(&header[4..8])?; - if version != STORAGE_VERSION { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Unsupported storage version: {version}"))); - } - let root_offset = u64_from_bytes(&header[8..16])?; - - // Read metadata - let metadata_len_raw = u32_from_bytes(&header[16..20])?; - let (metadata_len, _) = decode_metadata_len_and_flags(metadata_len_raw); - let metadata = if metadata_len > 0 { - if (METADATA_DATA_START_POS + metadata_len as usize) > PAGE_SIZE_USIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Metadata length exceeds header page size")); - } - header[METADATA_DATA_START_POS..(METADATA_DATA_START_POS + metadata_len as usize)].to_vec() - } else { - Vec::new() - }; - - let mut cursor = io::Cursor::new(mmap.as_ref()); - // Start after header block, with nested=true to deserialize all nodes - let (root, _) = BPlusTreeNode::::deserialize_from_mmap(&mmap, &mut cursor, root_offset, true)?; - - let (inner_order, leaf_order) = default_orders(); - Ok(Self { root, inner_order, leaf_order, metadata: BPlusTreeMetadata::from_bytes(&metadata), dirty: false }) - } - - /// Find the largest key <= `key` in the in-memory tree and return references to (key, value). - pub fn find_le(&self, key: &K) -> Option<(&K, &V)> { - // empty tree - if self.root.keys.is_empty() && self.root.is_leaf && self.root.values.is_empty() { - return None; - } - self.root.find_le(key) - } - - pub fn traverse(&self, mut visit: F) - where - F: FnMut(&Vec, &Vec), - { - self.root.traverse(&mut visit); - } -} - -type CachedNode = (BPlusTreeNode, Option>); -type NodeCache = LruCache>; - -struct SplitNodeCache { - internal: NodeCache, - leaf: NodeCache, -} - -impl SplitNodeCache { - fn new() -> Self { Self { internal: new_node_cache(INTERNAL_CACHE_CAPACITY), leaf: new_node_cache(CACHE_CAPACITY) } } - - fn clear(&mut self) { - self.internal.clear(); - self.leaf.clear(); - } - - fn put(&mut self, offset: u64, decoded: CachedNode) -> Option<&CachedNode> { - let is_leaf = decoded.0.is_leaf; - if is_leaf { - self.leaf.put(offset, decoded); - self.leaf.get(&offset) - } else { - self.internal.put(offset, decoded); - self.internal.get(&offset) - } - } -} - -fn read_node_cached<'a, K, V, R: Read + Seek>( - file: &mut R, - buffer: &mut Vec, - cache: &'a mut SplitNodeCache, - offset: u64, -) -> Result<&'a CachedNode, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - if cache.internal.contains(&offset) { - return cache - .internal - .get(&offset) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached internal node".into())); - } - if cache.leaf.contains(&offset) { - return cache - .leaf - .get(&offset) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached leaf node".into())); - } - - let decoded = BPlusTreeNode::::deserialize_from_block(file, buffer, offset, false)?; - cache - .put(offset, decoded) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached node".into())) -} - -fn query_tree( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - key: &K, - start_offset: u64, -) -> Result, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut offset = start_offset; - loop { - let (node, pointers) = read_node_cached(file, buffer, cache, offset)?; - - if node.is_leaf { - return match node.keys.binary_search(key) { - Ok(idx) => match node.value_info.get(idx) { - Some(info) => { - if info.is_tombstone() { - return Ok(None); - } - let value = BPlusTreeNode::::load_value_from_info(file, info)?; - Ok(Some(value)) - } - None => Ok(None), - }, - Err(_) => Ok(None), - }; - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - offset = *child_offset; - } else { - return Ok(None); - } - } else { - return Ok(None); - } - } -} - -fn query_tree_contains_live_key( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - key: &K, - start_offset: u64, - _has_tombstones: bool, -) -> Result -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut offset = start_offset; - loop { - let (node, pointers) = read_node_cached(file, buffer, cache, offset)?; - - if node.is_leaf { - return Ok(match node.keys.binary_search(key) { - // Always check tombstone — the has_tombstones flag is only a hint - // and may be incorrect for legacy files or after corruption. - Ok(idx) => node.value_info.get(idx).is_some_and(|info| !info.is_tombstone()), - Err(_) => false, - }); - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - offset = *child_offset; - } else { - return Ok(false); - } - } else { - return Ok(false); - } - } -} - -fn query_tree_le( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - key: &K, - start_offset: u64, -) -> Result, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - fn rightmost_live_value_le( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - key: &K, - mut offset: u64, - ) -> Result, BPlusTreeError> - where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - { - loop { - let (node, pointers) = read_node_cached(file, buffer, cache, offset)?; - if node.is_leaf { - let mut idx = get_entry_index_upper_bound::(&node.keys, key); - while idx > 0 { - idx -= 1; - let Some(info) = node.value_info.get(idx) else { - continue; - }; - if info.is_tombstone() { - continue; - } - return Ok(Some(BPlusTreeNode::::load_value_from_info(file, info)?)); - } - return Ok(None); - } else if let Some(ptrs) = pointers { - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - let Some(next_offset) = ptrs - .get(child_idx) - .or_else(|| child_idx.checked_sub(1).and_then(|previous| ptrs.get(previous))) - .or_else(|| ptrs.last()) - else { - return Ok(None); - }; - offset = *next_offset; - } else { - return Ok(None); - } - } - } - - fn previous_subtree_value( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - key: &K, - ancestors: &mut Vec<(Vec, usize)>, - ) -> Result, BPlusTreeError> - where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - { - while let Some((ptrs, child_idx)) = ancestors.pop() { - if let Some(previous_idx) = child_idx.checked_sub(1) { - if let Some(previous_offset) = ptrs.get(previous_idx) { - if let Some(value) = rightmost_live_value_le(file, buffer, cache, key, *previous_offset)? { - return Ok(Some(value)); - } - } - } - } - Ok(None) - } - - let mut offset = start_offset; - let mut ancestors: Vec<(Vec, usize)> = Vec::new(); - loop { - let (node, pointers) = read_node_cached(file, buffer, cache, offset)?; - - if node.is_leaf { - let mut idx = get_entry_index_upper_bound::(&node.keys, key); - while idx > 0 { - idx -= 1; - let Some(info) = node.value_info.get(idx) else { - continue; - }; - if info.is_tombstone() { - continue; - } - let value = BPlusTreeNode::::load_value_from_info(file, info)?; - return Ok(Some(value)); - } - return previous_subtree_value(file, buffer, cache, key, &mut ancestors); - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - ancestors.push((child_offsets.clone(), child_idx)); - offset = *child_offset; - } else if let Some(last) = child_offsets.last() { - ancestors.push((child_offsets.clone(), child_offsets.len().saturating_sub(1))); - offset = *last; - } else { - return Ok(None); - } - } else { - return Ok(None); - } - } -} - -fn count_items( - file: &mut R, - buffer: &mut Vec, - cache: &mut SplitNodeCache, - start_offset: u64, -) -> Result -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut count = 0; - let mut stack: OffsetStack = smallvec![start_offset]; - while let Some(offset) = stack.pop() { - let (node, pointers) = read_node_cached(file, buffer, cache, offset)?; - - if node.is_leaf { - count += node.value_info.iter().filter(|info| !info.is_tombstone()).count(); - } else if let Some(ptrs) = pointers { - stack.extend(ptrs.iter().copied()); - } - } - Ok(count) -} - -fn count_items_mmap(mmap: &[u8], start_offset: u64) -> Result -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut count = 0; - let mut stack: OffsetStack = smallvec![start_offset]; - let mut cursor = io::Cursor::new(mmap); - while let Some(offset) = stack.pop() { - let (node, pointers) = BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)?; - - if node.is_leaf { - count += node.value_info.iter().filter(|info| !info.is_tombstone()).count(); - } else if let Some(ptrs) = pointers { - stack.extend(ptrs); - } - } - Ok(count) -} - -fn read_node_mmap_cached<'a, K, V>( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - node_cache: &'a mut SplitNodeCache, - offset: u64, -) -> Result<&'a CachedNode, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - if node_cache.internal.contains(&offset) { - return node_cache - .internal - .get(&offset) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached internal node".into())); - } - if node_cache.leaf.contains(&offset) { - return node_cache - .leaf - .get(&offset) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached leaf node".into())); - } - - let decoded = BPlusTreeNode::::deserialize_from_mmap(mmap, cursor, offset, false)?; - node_cache - .put(offset, decoded) - .ok_or_else(|| BPlusTreeError::InvalidStructure("Missing cached node".into())) -} - -fn query_tree_mmap( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - node_cache: &mut SplitNodeCache, - key: &K, - start_offset: u64, -) -> Result, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut offset = start_offset; - loop { - let (node, pointers) = read_node_mmap_cached(mmap, cursor, node_cache, offset)?; - - if node.is_leaf { - return match node.keys.binary_search(key) { - Ok(idx) => match node.value_info.get(idx) { - Some(info) => { - if info.is_tombstone() { - return Ok(None); - } - let value = BPlusTreeNode::::load_value_from_info(cursor, info)?; - Ok(Some(value)) - } - None => Ok(None), - }, - Err(_) => Ok(None), - }; - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - offset = *child_offset; - } else { - return Ok(None); - } - } else { - return Ok(None); - } - } -} - -fn query_tree_mmap_contains_live_key( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - node_cache: &mut SplitNodeCache, - key: &K, - start_offset: u64, - _has_tombstones: bool, -) -> Result -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut offset = start_offset; - loop { - let (node, pointers) = read_node_mmap_cached(mmap, cursor, node_cache, offset)?; - - if node.is_leaf { - return Ok(match node.keys.binary_search(key) { - // Always check tombstone — the has_tombstones flag is only a hint - // and may be incorrect for legacy files or after corruption. - Ok(idx) => node.value_info.get(idx).is_some_and(|info| !info.is_tombstone()), - Err(_) => false, - }); - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - offset = *child_offset; - } else { - return Ok(false); - } - } else { - return Ok(false); - } - } -} - -/// Zero-copy optimized query for mmap'd data. -/// -/// This function uses zero-copy key scanning for internal nodes, avoiding -/// the need to deserialize the entire keys vector. It only falls back to -/// full deserialization for leaf nodes where we need to perform binary search -/// and access `value_info`. -/// -/// Performance improvement: For trees with many internal nodes, this eliminates -/// N heap allocations per internal node (where N is the number of keys). -fn query_tree_mmap_zero_copy( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - key: &K, - start_offset: u64, -) -> Result, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + MsgPackScannable, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - let mut offset = start_offset; - - loop { - let node_start = - usize::try_from(offset).map_err(|e| BPlusTreeError::Corrupted(format!("Invalid offset: {e}")))?; - - if node_start >= mmap.len() { - return Err(BPlusTreeError::Corrupted("Offset out of bounds".into())); - } - - let node_bytes = &mmap[node_start..]; - - // Try zero-copy scan first - if let Some(scan_result) = scan_internal_node_zero_copy(node_bytes, key) { - if scan_result.is_leaf { - // Fall back to full deserialization for leaf nodes - let (node, _) = BPlusTreeNode::::deserialize_from_mmap(mmap, cursor, offset, false)?; - - return match node.keys.binary_search(key) { - Ok(idx) => match node.value_info.get(idx) { - Some(info) => { - if info.is_tombstone() { - return Ok(None); - } - let value = BPlusTreeNode::::load_value_from_info(cursor, info)?; - Ok(Some(value)) - } - None => Ok(None), - }, - Err(_) => Ok(None), - }; - } - - // Internal node: read the child pointer using zero-copy - let pointers_bytes = &node_bytes[scan_result.pointers_start..]; - if let Some(child_offset) = read_pointer_at_index(pointers_bytes, scan_result.child_idx) { - offset = child_offset; - continue; - } - // Fall through to fallback if pointer read failed - } - - // Fallback: full deserialization (handles edge cases, unsupported key types, etc.) - let (node, pointers) = BPlusTreeNode::::deserialize_from_mmap(mmap, cursor, offset, false)?; - - if node.is_leaf { - return match node.keys.binary_search(key) { - Ok(idx) => match node.value_info.get(idx) { - Some(info) => { - if info.is_tombstone() { - return Ok(None); - } - let value = BPlusTreeNode::::load_value_from_info(cursor, info)?; - Ok(Some(value)) - } - None => Ok(None), - }, - Err(_) => Ok(None), - }; - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - offset = *child_offset; - } else { - return Ok(None); - } - } else { - return Ok(None); - } - } -} - -fn query_tree_le_mmap( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - key: &K, - start_offset: u64, -) -> Result, BPlusTreeError> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - fn rightmost_live_value_le( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - key: &K, - mut offset: u64, - ) -> Result, BPlusTreeError> - where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - { - loop { - let (node, pointers) = BPlusTreeNode::::deserialize_from_mmap(mmap, cursor, offset, false)?; - if node.is_leaf { - let mut idx = get_entry_index_upper_bound::(&node.keys, key); - while idx > 0 { - idx -= 1; - let Some(info) = node.value_info.get(idx) else { - continue; - }; - if info.is_tombstone() { - continue; - } - return Ok(Some(BPlusTreeNode::::load_value_from_info(cursor, info)?)); - } - return Ok(None); - } else if let Some(ptrs) = pointers { - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - let Some(next_offset) = ptrs - .get(child_idx) - .or_else(|| child_idx.checked_sub(1).and_then(|previous| ptrs.get(previous))) - .or_else(|| ptrs.last()) - else { - return Ok(None); - }; - offset = *next_offset; - } else { - return Ok(None); - } - } - } - - fn previous_subtree_value( - mmap: &[u8], - cursor: &mut io::Cursor<&[u8]>, - key: &K, - ancestors: &mut Vec<(Vec, usize)>, - ) -> Result, BPlusTreeError> - where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - { - while let Some((ptrs, child_idx)) = ancestors.pop() { - if let Some(previous_idx) = child_idx.checked_sub(1) { - if let Some(previous_offset) = ptrs.get(previous_idx) { - if let Some(value) = rightmost_live_value_le(mmap, cursor, key, *previous_offset)? { - return Ok(Some(value)); - } - } - } - } - Ok(None) - } - - let mut offset = start_offset; - let mut ancestors: Vec<(Vec, usize)> = Vec::new(); - loop { - let (node, pointers) = BPlusTreeNode::::deserialize_from_mmap(mmap, cursor, offset, false)?; - - if node.is_leaf { - let mut idx = get_entry_index_upper_bound::(&node.keys, key); - while idx > 0 { - idx -= 1; - let Some(info) = node.value_info.get(idx) else { - continue; - }; - if info.is_tombstone() { - continue; - } - let value = BPlusTreeNode::::load_value_from_info(cursor, info)?; - return Ok(Some(value)); - } - return previous_subtree_value(mmap, cursor, key, &mut ancestors); - } - - let child_idx = get_entry_index_upper_bound::(&node.keys, key); - if let Some(child_offsets) = pointers { - if let Some(child_offset) = child_offsets.get(child_idx) { - ancestors.push((child_offsets.clone(), child_idx)); - offset = *child_offset; - } else { - return Ok(None); - } - } else { - return Ok(None); - } - } -} - -fn lru_cache_capacity(capacity: usize) -> NonZeroUsize { - match NonZeroUsize::new(capacity) { - Some(capacity) => capacity, - None => NonZeroUsize::MIN, - } -} - -fn new_node_cache(capacity: usize) -> NodeCache { LruCache::new(lru_cache_capacity(capacity)) } - -#[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(), - } - } -} - -/// `BPlusTreeQuery` performs on-disk queries without loading the entire tree into memory. -/// For frequent queries, consider using `BPlusTree::load()` instead, which loads the full tree into memory -/// at the cost of higher memory usage. -pub struct BPlusTreeQuery { - file: Option>, - mmap: Option, - filepath: PathBuf, - file_identity: Option, - has_tombstones: bool, - buffer: Vec, - cache: SplitNodeCache, - node_cache: SplitNodeCache, - last_refresh_at: Instant, - refresh_interval: Duration, - root_offset: u64, - _marker_k: PhantomData, - _marker_v: PhantomData, -} - -impl BPlusTreeQuery -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - 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")); - } - - // Try Mmap - let mmap = mmap_with_advice(&file, Advice::Normal, "B+Tree query"); - - // Verify Header - let mut header = [0u8; METADATA_DATA_START_POS]; - read_exact_at_offset(&file, &mut header, 0)?; - - if &header[0..4] != MAGIC { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid magic number")); - } - let version = u32::from_le_bytes( - header[4..8].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid version slice"))?, - ); - if version != STORAGE_VERSION { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Unsupported storage version: {version}"))); - } - let root_offset = u64::from_le_bytes( - header[8..16] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset slice"))?, - ); - let metadata_len_raw = u32::from_le_bytes( - header[16..20] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length slice"))?, - ); - let (metadata_len, has_tombstones) = decode_metadata_len_and_flags(metadata_len_raw); - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - - Ok(Self { - file: if mmap.is_some() { None } else { Some(utils::file_reader(file)) }, - mmap, - filepath: PathBuf::new(), - file_identity, - has_tombstones, - buffer: vec![0u8; PAGE_SIZE_USIZE], - cache: SplitNodeCache::new(), - node_cache: SplitNodeCache::new(), - last_refresh_at: Instant::now(), - refresh_interval: QUERY_REFRESH_INTERVAL, - root_offset, - _marker_k: PhantomData, - _marker_v: PhantomData, - }) - } - - pub fn try_new(filepath: &Path) -> io::Result { - let file = File::open(filepath)?; - let mut query = Self::try_from_file(file)?; - query.filepath = filepath.to_path_buf(); - Ok(query) - } - - /// Clone an existing query without re-reading the file header. - /// This avoids synchronous disk initialization while still providing - /// an independent reader. - pub fn try_clone(&self) -> io::Result { - let (file, mmap) = if self.mmap.is_some() { - if self.filepath.as_os_str().is_empty() { - return Err(io::Error::new(io::ErrorKind::InvalidInput, "Missing filepath for mmap clone")); - } - let file = File::open(&self.filepath)?; - if let Some(mapped) = mmap_with_advice(&file, Advice::Normal, "B+Tree query clone") { - (None, Some(mapped)) - } else { - (Some(utils::file_reader(file)), None) - } - } else if let Some(file) = &self.file { - let cloned = file.get_ref().try_clone()?; - (Some(utils::file_reader(cloned)), None) - } else { - return Err(io::Error::new(io::ErrorKind::InvalidInput, "No data source available to clone")); - }; - - Ok(Self { - file, - mmap, - filepath: self.filepath.clone(), - file_identity: self.file_identity, - has_tombstones: self.has_tombstones, - buffer: vec![0u8; PAGE_SIZE_USIZE], - cache: SplitNodeCache::new(), - node_cache: SplitNodeCache::new(), - last_refresh_at: Instant::now(), - refresh_interval: self.refresh_interval, - root_offset: self.root_offset, - _marker_k: PhantomData, - _marker_v: PhantomData, - }) - } - - /// Returns the filepath this query was opened from. - pub fn filepath(&self) -> &Path { &self.filepath } - - #[cfg(test)] - pub(crate) fn clone_error_fixture() -> Self { - Self { - file: None, - mmap: None, - filepath: PathBuf::new(), - file_identity: None, - has_tombstones: false, - buffer: vec![0u8; PAGE_SIZE_USIZE], - cache: SplitNodeCache::new(), - node_cache: SplitNodeCache::new(), - last_refresh_at: Instant::now(), - refresh_interval: QUERY_REFRESH_INTERVAL, - root_offset: 0, - _marker_k: PhantomData, - _marker_v: PhantomData, - } - } - - /// Force a header/root refresh, bypassing the automatic refresh throttle. - pub fn refresh(&mut self) -> io::Result<()> { self.refresh_root_offset() } - - /// Configure how often hot query paths check the backing file for a new root. - /// - /// Set to `Duration::ZERO` to preserve the old always-refresh behavior. - pub fn set_refresh_interval(&mut self, interval: Duration) { self.refresh_interval = interval; } - - fn read_root_offset_and_tombstone_flag_from_file(file: &File) -> io::Result<(u64, bool)> { - let mut root_and_metadata_len = [0u8; size_of::() + size_of::()]; - read_exact_at_offset(file, &mut root_and_metadata_len, ROOT_OFFSET_POS)?; - let root_offset = u64::from_le_bytes( - root_and_metadata_len[0..8] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset bytes"))?, - ); - let metadata_len_raw = u32::from_le_bytes( - root_and_metadata_len[8..12] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length bytes"))?, - ); - let (metadata_len, has_tombstones) = decode_metadata_len_and_flags(metadata_len_raw); - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - - Ok((root_offset, has_tombstones)) - } - - fn refresh_root_offset(&mut self) -> io::Result<()> { - let (new_root_offset, new_has_tombstones) = 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, has_tombstones) = Self::read_root_offset_and_tombstone_flag_from_file(&file)?; - - if remap_required { - if let Some(remapped) = mmap_with_advice(&file, Advice::Normal, "B+Tree query remap") { - 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, has_tombstones) - } 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 root_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..root_end) - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Header too small"))?; - let root_offset_bytes: [u8; size_of::()] = root_bytes - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset bytes"))?; - let metadata_len_start = usize::try_from(METADATA_OFFSET_POS) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("Invalid metadata position: {e}")))?; - let metadata_len_end = metadata_len_start - .checked_add(size_of::()) - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length range"))?; - let metadata_len_bytes = mmap - .get(metadata_len_start..metadata_len_end) - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Header too small for metadata length"))?; - let metadata_len_raw = u32::from_le_bytes( - metadata_len_bytes - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length bytes"))?, - ); - let (metadata_len, has_tombstones) = decode_metadata_len_and_flags(metadata_len_raw); - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - (u64::from_le_bytes(root_offset_bytes), has_tombstones) - } else if self.file.is_some() { - if self.filepath.as_os_str().is_empty() { - if let Some(file) = &mut self.file { - Self::read_root_offset_and_tombstone_flag_from_file(file.get_ref())? - } else { - return Err(io::Error::new(io::ErrorKind::InvalidData, "No data source available")); - } - } else { - let file = File::open(&self.filepath)?; - let metadata = file.metadata()?; - let current_identity = FileIdentity::from_metadata(&metadata); - let replacement_required = self.file_identity != Some(current_identity); - let root_state = Self::read_root_offset_and_tombstone_flag_from_file(&file)?; - - if replacement_required { - self.file = Some(utils::file_reader(file)); - self.file_identity = Some(current_identity); - self.cache.clear(); - self.node_cache.clear(); - } - - root_state - } - } else { - return Err(io::Error::new(io::ErrorKind::InvalidData, "No data source available")); - }; - - if new_root_offset != self.root_offset || new_has_tombstones != self.has_tombstones { - self.root_offset = new_root_offset; - self.has_tombstones = new_has_tombstones; - self.cache.clear(); - self.node_cache.clear(); - } - - self.last_refresh_at = Instant::now(); - Ok(()) - } - - fn refresh_root_offset_if_due(&mut self) -> io::Result<()> { - if self.refresh_interval.is_zero() || self.last_refresh_at.elapsed() >= self.refresh_interval { - self.refresh_root_offset() - } else { - Ok(()) - } - } - - pub fn query(&mut self, key: &K) -> Result, BPlusTreeError> { - self.refresh_root_offset_if_due().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) - } else if let Some(file) = &mut self.file { - query_tree(file, &mut self.buffer, &mut self.cache, key, self.root_offset) - } else { - Err(BPlusTreeError::InvalidStructure("No data source available".into())) - } - } - - pub const fn has_tombstones(&self) -> bool { self.has_tombstones } - - pub fn contains_live_key(&mut self, key: &K) -> Result { - self.refresh_root_offset_if_due().map_err(BPlusTreeError::Io)?; - if let Some(mmap) = &self.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - query_tree_mmap_contains_live_key( - mmap, - &mut cursor, - &mut self.node_cache, - key, - self.root_offset, - self.has_tombstones, - ) - } else if let Some(file) = &mut self.file { - query_tree_contains_live_key::( - file, - &mut self.buffer, - &mut self.cache, - key, - self.root_offset, - self.has_tombstones, - ) - } else { - Err(BPlusTreeError::InvalidStructure("No data source available".into())) - } - } -} - -/// Additional methods for key types that support zero-copy scanning. -/// -/// These methods provide optimized query performance by avoiding heap allocations -/// when traversing internal nodes. For trees with many levels, this can significantly -/// reduce memory pressure and improve query latency. -impl BPlusTreeQuery -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + MsgPackScannable, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - /// Zero-copy optimized query that avoids deserializing internal node keys. - /// - /// This method scans through MessagePack-encoded keys directly, comparing - /// without allocating. It's significantly faster for trees with many internal - /// nodes, especially when keys are `u32` or `String`. - /// - /// # Example - /// ```ignore - /// let mut query = BPlusTreeQuery::::try_new(&path)?; - /// let value = query.query_zero_copy(&42)?; - /// ``` - pub fn query_zero_copy(&mut self, key: &K) -> Result, BPlusTreeError> { - self.refresh_root_offset_if_due().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) - } else if let Some(file) = &mut self.file { - // For file-based queries, fall back to regular query - // (zero-copy would require reading node data into a buffer first) - query_tree(file, &mut self.buffer, &mut self.cache, key, self.root_offset) - } else { - Err(BPlusTreeError::InvalidStructure("No data source available".into())) - } - } -} - -impl BPlusTreeQuery -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - pub fn query_le(&mut self, key: &K) -> Result, BPlusTreeError> { - self.refresh_root_offset_if_due().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) - } else if let Some(file) = &mut self.file { - query_tree_le(file, &mut self.buffer, &mut self.cache, key, self.root_offset) - } else { - Err(BPlusTreeError::InvalidStructure("No data source available".into())) - } - } - - pub fn len(&mut self) -> Result { - self.refresh_root_offset_if_due().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 { - count_items::(file, &mut self.buffer, &mut self.cache, self.root_offset) - } else { - Err(BPlusTreeError::InvalidStructure("No data source available".into())) - } - } - - pub fn is_empty(&mut self) -> Result { - self.refresh_root_offset_if_due().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)? - } else if let Some(file) = &mut self.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.buffer, self.root_offset, false)? - } else { - return Err(BPlusTreeError::InvalidStructure("No data source available".into())); - }; - Ok(node.is_leaf && node.keys.is_empty()) - } - - /// Provides a disk-backed iterator that traverses the entire tree in order. - pub fn iter(&mut self) -> BPlusTreeDiskIterator<'_, K, V> { BPlusTreeDiskIterator::new(self) } - - /// Iterates over key-value pairs within a given range using `right_sibling` pointers. - /// - /// This is more efficient than iterating the full tree and filtering when you only - /// need a subset of keys. - /// - /// Tombstones are skipped automatically. - pub fn range_iter( - &mut self, - start: Bound<&K>, - end: Bound<&K>, - ) -> impl Iterator> + '_ { - let start_cloned = match start { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let end_cloned = match end { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - RangeLeafIterator::new(self, start_cloned, end_cloned) - } - - /// Returns a page of key-value pairs within a given range using offset/limit. - /// - /// More efficient than `range_iter` when you only need a subset, as it stops - /// scanning once `limit` items have been collected. - /// - /// Returns `(items, has_more)` where `has_more` indicates whether additional - /// items exist beyond the returned page. - pub fn range_page( - &mut self, - start: Bound<&K>, - end: Bound<&K>, - offset: usize, - limit: usize, - ) -> Result<(Vec<(K, V)>, bool), BPlusTreeError> { - let start_cloned = match start { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let end_cloned = match end { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let mut iter = RangeLeafIterator::new(self, start_cloned, end_cloned); - iter.skip_items(offset)?; - - let mut collected: Vec<(K, V)> = Vec::with_capacity(limit); - let mut has_more = false; - - while collected.len() < limit { - match iter.next() { - Some(Ok(item)) => collected.push(item), - Some(Err(err)) => return Err(err), - None => return Ok((collected, false)), - } - } - - match iter.next() { - Some(Ok(_)) => has_more = true, - Some(Err(err)) => return Err(err), - None => {} - } - - Ok((collected, has_more)) - } - - pub(crate) fn into_sorted_parts(self) -> (PathBuf, Option>, Option) { - (self.filepath, self.file, self.mmap) - } - - /// Owned iterator that traverses the tree in order defined by a secondary sorted index. - /// - /// The index path is automatically derived from the tree filepath by changing - /// the extension to `.idx`. For example, if the tree is at `/data/items.bin`, - /// the index is expected at `/data/items.idx`. - /// - /// This iterator reads values directly from stored offsets in O(1) time, - /// avoiding O(log n) tree traversal per item. - pub fn disk_iter_sorted( - self, - ) -> io::Result> - where - SortKey: for<'de> Deserialize<'de>, - { - super::sorted_index::BPlusTreeSortedIteratorOwned::::new_hybrid( - self.filepath.clone(), - self.file, - self.mmap, - ) - } - - /// Owned iterator with explicit index path. - pub fn disk_iter_sorted_with_path( - self, - index_path: &Path, - ) -> io::Result> - where - SortKey: for<'de> Deserialize<'de>, - { - super::sorted_index::BPlusTreeSortedIteratorOwned::::with_index_path_hybrid( - self.filepath.clone(), - self.file, - self.mmap, - index_path, - ) - } - - /// Traverses the tree and calls the provided closure for each leaf's keys and values. - pub fn traverse(&mut self, mut f: F) -> io::Result<()> - where - F: FnMut(&[K], &[V]), - { - let mut it = self.iter(); - while let Some((keys, values)) = it.next_leaf()? { - f(&keys, &values); - } - Ok(()) - } - - /// Collects all entries with their value locations. - /// Used for building sorted indexes that need direct value access. - #[allow(clippy::too_many_lines)] - pub fn collect_with_locations(&mut self) -> io::Result> { - let mut result = Vec::new(); - let mut stack: OffsetStack = smallvec![self.root_offset]; - let has_tombstones = self.has_tombstones; - - while let Some(offset) = stack.pop() { - let (node, pointers) = if let Some(mmap) = &self.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? - } else if let Some(file) = &mut self.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.buffer, offset, false)? - } else { - return Err(io::Error::other("No data source available")); - }; - - if node.is_leaf { - if let Some(mmap) = &self.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - if has_tombstones { - for (key, info) in node.keys.into_iter().zip(node.value_info.iter()) { - if info.is_tombstone() { - continue; - } - let value = BPlusTreeNode::::load_value_from_info(&mut cursor, info)?; - let location = match info.mode { - ValueStorageMode::Single(offset) => { - super::sorted_index::ValueLocation::Single { offset, length: info.length } - } - ValueStorageMode::Packed(block_offset, index) => { - super::sorted_index::ValueLocation::Packed { - block_offset, - index, - length: info.length, - } - } - ValueStorageMode::Tombstone => continue, - }; - result.push((key, value, location)); - } - } else { - for (key, info) in node.keys.into_iter().zip(node.value_info.iter()) { - let value = BPlusTreeNode::::load_value_from_info(&mut cursor, info)?; - let location = match info.mode { - ValueStorageMode::Single(offset) => { - super::sorted_index::ValueLocation::Single { offset, length: info.length } - } - ValueStorageMode::Packed(block_offset, index) => { - super::sorted_index::ValueLocation::Packed { - block_offset, - index, - length: info.length, - } - } - ValueStorageMode::Tombstone => continue, - }; - result.push((key, value, location)); - } - } - } else if let Some(file) = &mut self.file { - if has_tombstones { - for (key, info) in node.keys.into_iter().zip(node.value_info.iter()) { - if info.is_tombstone() { - continue; - } - let value = BPlusTreeNode::::load_value_from_info(file, info)?; - let location = match info.mode { - ValueStorageMode::Single(offset) => { - super::sorted_index::ValueLocation::Single { offset, length: info.length } - } - ValueStorageMode::Packed(block_offset, index) => { - super::sorted_index::ValueLocation::Packed { - block_offset, - index, - length: info.length, - } - } - ValueStorageMode::Tombstone => continue, - }; - result.push((key, value, location)); - } - } else { - for (key, info) in node.keys.into_iter().zip(node.value_info.iter()) { - let value = BPlusTreeNode::::load_value_from_info(file, info)?; - let location = match info.mode { - ValueStorageMode::Single(offset) => { - super::sorted_index::ValueLocation::Single { offset, length: info.length } - } - ValueStorageMode::Packed(block_offset, index) => { - super::sorted_index::ValueLocation::Packed { - block_offset, - index, - length: info.length, - } - } - ValueStorageMode::Tombstone => continue, - }; - result.push((key, value, location)); - } - } - } - } else if let Some(ptrs) = pointers { - for ptr in ptrs.into_iter().rev() { - stack.push(ptr); - } - } - } - Ok(result) - } - - /// Provides an owned disk-backed iterator. - pub fn disk_iter(self) -> BPlusTreeDiskIteratorOwned { BPlusTreeDiskIteratorOwned::new(self) } -} - -/// Range scan iterator that seeks into the tree and then walks in-order -/// without scanning from the root for every entry. -struct RangeLeafIterator<'a, K, V> { - tree: &'a mut BPlusTreeQuery, - start_bound: Bound, - end_bound: Bound, - stack: TraversalStack, - current_leaf: Option>, - leaf_idx: usize, - initialized: bool, - exhausted: bool, -} - -/// Generate the `new`, `load_leaf_from_node`, `skip_items`, and -/// `key_past_end` methods for the `RangeLeafIterator` and -/// `RangeLeafIteratorUpdate` impls. These four methods have byte-identical -/// bodies across both iterators — they only touch shared fields (no -/// `self.tree.mmap` branching, no `refresh_root_offset_if_due`). -/// -/// The remaining methods (`initialize`, `descend_to_leaf`, `advance_leaf`) -/// stay impl-specific because they differ: `BPlusTreeQuery` carries an `mmap` -/// plus a `refresh_root_offset_if_due` hook that `BPlusTreeUpdate` lacks. -/// -/// `$lt` is the lifetime token (`'a`) from the enclosing `impl<'a, ...>` -/// block, threaded into the `new` signature so the function returns `Self` -/// borrowing from the same lifetime. -macro_rules! impl_range_leaf_common { - ($tree_ty:ty, $lt:tt) => { - fn new(tree: &$lt mut $tree_ty, start: Bound, end: Bound) -> Self { - Self { - tree, - start_bound: start, - end_bound: end, - stack: smallvec![], - current_leaf: None, - leaf_idx: 0, - initialized: false, - exhausted: false, - } - } - - fn load_leaf_from_node(&mut self, node: BPlusTreeNode, start_idx: usize) { - self.current_leaf = Some(node); - self.leaf_idx = start_idx; - } - - fn key_past_end(&self, key: &K) -> bool { - match &self.end_bound { - Bound::Included(end) => key > end, - Bound::Excluded(end) => key >= end, - Bound::Unbounded => false, - } - } - - fn skip_items(&mut self, mut remaining: usize) -> Result<(), BPlusTreeError> { - if remaining == 0 || self.exhausted { - return Ok(()); - } - if !self.initialized { - self.initialize().map_err(BPlusTreeError::Io)?; - } - - while remaining > 0 && !self.exhausted { - let Some(node) = self.current_leaf.as_ref() else { - self.advance_leaf().map_err(BPlusTreeError::Io)?; - continue; - }; - - while self.leaf_idx < node.keys.len() && remaining > 0 { - let idx = self.leaf_idx; - self.leaf_idx += 1; - let key = &node.keys[idx]; - if self.key_past_end(key) { - self.current_leaf = None; - self.exhausted = true; - return Ok(()); - } - if self.tree.has_tombstones && node.value_info[idx].is_tombstone() { - continue; - } - remaining -= 1; - } - - if self.leaf_idx >= node.keys.len() { - self.current_leaf = None; - } - } - - Ok(()) - } - }; -} - -impl<'a, K, V> RangeLeafIterator<'a, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - impl_range_leaf_common!(BPlusTreeQuery, 'a); - - fn descend_to_leaf(&mut self, mut offset: u64, mut start_key: Option<&K>) -> io::Result<()> { - loop { - let (node, pointers) = if let Some(mmap) = &self.tree.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? - } else if let Some(file) = &mut self.tree.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.tree.buffer, offset, false)? - } else { - return Err(io::Error::other("No data source available")); - }; - - if node.is_leaf { - let start_idx = if let Some(key) = start_key { - match self.start_bound { - Bound::Included(_) => node.keys.partition_point(|candidate| candidate < key), - Bound::Excluded(_) => node.keys.partition_point(|candidate| candidate <= key), - Bound::Unbounded => 0, - } - } else { - 0 - }; - self.load_leaf_from_node(node, start_idx); - return Ok(()); - } - - let child_idx = if let Some(key) = start_key { - get_entry_index_upper_bound(&node.keys, key) - } else { - 0 - }; - - let Some(ptrs) = pointers else { - self.exhausted = true; - return Ok(()); - }; - let Some(&next_offset) = ptrs.get(child_idx) else { - self.exhausted = true; - return Ok(()); - }; - self.stack.push((offset, child_idx.saturating_add(1))); - offset = next_offset; - start_key = None.or(start_key); - } - } - - fn initialize(&mut self) -> io::Result<()> { - if self.initialized { - return Ok(()); - } - self.initialized = true; - let start_key = match self.start_bound.clone() { - Bound::Included(key) | Bound::Excluded(key) => Some(key), - Bound::Unbounded => None, - }; - self.tree.refresh_root_offset_if_due()?; - self.descend_to_leaf(self.tree.root_offset, start_key.as_ref()) - } - - fn advance_leaf(&mut self) -> io::Result<()> { - while let Some((offset, child_idx)) = self.stack.pop() { - let (_node, pointers) = if let Some(mmap) = &self.tree.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? - } else if let Some(file) = &mut self.tree.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.tree.buffer, offset, false)? - } else { - return Err(io::Error::other("No data source available")); - }; - - let Some(ptrs) = pointers else { - continue; - }; - let Some(&next_offset) = ptrs.get(child_idx) else { - continue; - }; - if child_idx + 1 < ptrs.len() { - self.stack.push((offset, child_idx + 1)); - } - self.descend_to_leaf(next_offset, None)?; - return Ok(()); - } - - self.exhausted = true; - Ok(()) - } -} - -impl Iterator for RangeLeafIterator<'_, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = Result<(K, V), BPlusTreeError>; - - fn next(&mut self) -> Option { - if self.exhausted { - return None; - } - - // Lazy initialization - if !self.initialized { - if let Err(e) = self.initialize() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - } - - loop { - let Some(node) = self.current_leaf.as_ref() else { - if self.exhausted { - return None; - } - if let Err(e) = self.advance_leaf() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - if self.exhausted { - return None; - } - continue; - }; - - if self.leaf_idx >= node.keys.len() { - self.current_leaf = None; - if self.exhausted { - return None; - } - if let Err(e) = self.advance_leaf() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - if self.exhausted { - return None; - } - continue; - } - - let idx = self.leaf_idx; - self.leaf_idx += 1; - let key = node.keys[idx].clone(); - - if self.key_past_end(&key) { - self.current_leaf = None; - self.exhausted = true; - return None; - } - - let info = node.value_info[idx].clone(); - if self.tree.has_tombstones && info.is_tombstone() { - continue; - } - - let value = if let Some(mmap) = &self.tree.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - match BPlusTreeNode::::load_value_from_info(&mut cursor, &info) { - Ok(value) => value, - Err(err) => { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(err))); - } - } - } else if let Some(file) = &mut self.tree.file { - match BPlusTreeNode::::load_value_from_info(file, &info) { - Ok(value) => value, - Err(err) => { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(err))); - } - } - } else { - self.exhausted = true; - return Some(Err(BPlusTreeError::InvalidStructure("No data source available".into()))); - }; - - return Some(Ok((key, value))); - } - } -} - -/// Range scan iterator for `BPlusTreeUpdate` (mutable update handle). -struct RangeLeafIteratorUpdate<'a, K, V> { - tree: &'a mut BPlusTreeUpdate, - start_bound: Bound, - end_bound: Bound, - stack: TraversalStack, - current_leaf: Option>, - leaf_idx: usize, - initialized: bool, - exhausted: bool, -} - -impl<'a, K, V> RangeLeafIteratorUpdate<'a, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - impl_range_leaf_common!(BPlusTreeUpdate, 'a); - - fn descend_to_leaf(&mut self, mut offset: u64, mut start_key: Option<&K>) -> io::Result<()> { - loop { - let (node, pointers) = BPlusTreeNode::::deserialize_from_block( - &mut self.tree.file, - &mut self.tree.read_buffer, - offset, - false, - )?; - - if node.is_leaf { - let start_idx = if let Some(key) = start_key { - match self.start_bound { - Bound::Included(_) => node.keys.partition_point(|candidate| candidate < key), - Bound::Excluded(_) => node.keys.partition_point(|candidate| candidate <= key), - Bound::Unbounded => 0, - } - } else { - 0 - }; - self.load_leaf_from_node(node, start_idx); - return Ok(()); - } - - let child_idx = if let Some(key) = start_key { - get_entry_index_upper_bound(&node.keys, key) - } else { - 0 - }; - - let Some(ptrs) = pointers else { - self.exhausted = true; - return Ok(()); - }; - let Some(&next_offset) = ptrs.get(child_idx) else { - self.exhausted = true; - return Ok(()); - }; - self.stack.push((offset, child_idx.saturating_add(1))); - offset = next_offset; - start_key = None.or(start_key); - } - } - - fn initialize(&mut self) -> io::Result<()> { - if self.initialized { - return Ok(()); - } - self.initialized = true; - let start_key = match self.start_bound.clone() { - Bound::Included(key) | Bound::Excluded(key) => Some(key), - Bound::Unbounded => None, - }; - self.descend_to_leaf(self.tree.root_offset, start_key.as_ref()) - } - - fn advance_leaf(&mut self) -> io::Result<()> { - while let Some((offset, child_idx)) = self.stack.pop() { - let (_node, pointers) = BPlusTreeNode::::deserialize_from_block( - &mut self.tree.file, - &mut self.tree.read_buffer, - offset, - false, - )?; - - let Some(ptrs) = pointers else { - continue; - }; - let Some(&next_offset) = ptrs.get(child_idx) else { - continue; - }; - if child_idx + 1 < ptrs.len() { - self.stack.push((offset, child_idx + 1)); - } - self.descend_to_leaf(next_offset, None)?; - return Ok(()); - } - - self.exhausted = true; - Ok(()) - } -} - -impl Iterator for RangeLeafIteratorUpdate<'_, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = Result<(K, V), BPlusTreeError>; - - fn next(&mut self) -> Option { - if self.exhausted { - return None; - } - - if !self.initialized { - if let Err(e) = self.initialize() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - } - - loop { - let Some(node) = self.current_leaf.as_ref() else { - if self.exhausted { - return None; - } - if let Err(e) = self.advance_leaf() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - if self.exhausted { - return None; - } - continue; - }; - - if self.leaf_idx >= node.keys.len() { - self.current_leaf = None; - if self.exhausted { - return None; - } - if let Err(e) = self.advance_leaf() { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(e))); - } - if self.exhausted { - return None; - } - continue; - } - - let idx = self.leaf_idx; - self.leaf_idx += 1; - let key = node.keys[idx].clone(); - - if self.key_past_end(&key) { - self.current_leaf = None; - self.exhausted = true; - return None; - } - - let info = node.value_info[idx].clone(); - if self.tree.has_tombstones && info.is_tombstone() { - continue; - } - - let value = match BPlusTreeNode::::load_value_from_info(&mut self.tree.file, &info) { - Ok(value) => value, - Err(err) => { - self.exhausted = true; - return Some(Err(BPlusTreeError::Io(err))); - } - }; - - return Some(Ok((key, value))); - } - } -} - -pub struct BPlusTreeDiskIteratorOwned { - query: BPlusTreeQuery, - stack: TraversalStack, - leaf_keys: Vec, - leaf_values: Vec, - leaf_idx: usize, -} - -impl BPlusTreeDiskIteratorOwned -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - fn new(query: BPlusTreeQuery) -> Self { - let root_offset = query.root_offset; - Self { query, stack: smallvec![(root_offset, 0)], leaf_keys: Vec::new(), leaf_values: Vec::new(), leaf_idx: 0 } - } - - pub fn try_is_empty(&mut self) -> io::Result { - loop { - if self.leaf_idx < self.leaf_keys.len() { - return Ok(false); - } - match self.next_leaf()? { - Some((keys, values)) => { - self.leaf_keys = keys; - self.leaf_values = values; - self.leaf_idx = 0; - if !self.leaf_keys.is_empty() { - return Ok(false); - } - } - None => return Ok(true), - } - } - } - - pub fn is_empty(&mut self) -> bool { self.try_is_empty().unwrap_or(true) } - - fn next_leaf(&mut self) -> io::Result, Vec)>> { - loop { - let Some((offset, child_idx)) = self.stack.pop() else { return Ok(None) }; - - let (node, pointers) = if let Some(mmap) = &self.query.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? - } else if let Some(file) = &mut self.query.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.query.buffer, offset, false)? - } else { - return Err(io::Error::other("No data source available")); - }; - - if node.is_leaf { - let mut keys = Vec::with_capacity(node.keys.len()); - let mut vals = Vec::with_capacity(node.value_info.len()); - let has_tombstones = self.query.has_tombstones; - if let Some(mmap) = &self.query.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - if has_tombstones { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - if value_info.is_tombstone() { - continue; - } - let v = BPlusTreeNode::::load_value_from_info(&mut cursor, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } else { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - let v = BPlusTreeNode::::load_value_from_info(&mut cursor, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } - } else if let Some(file) = &mut self.query.file { - if has_tombstones { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - if value_info.is_tombstone() { - continue; - } - let v = BPlusTreeNode::::load_value_from_info(file, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } else { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - let v = BPlusTreeNode::::load_value_from_info(file, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } - } - return Ok(Some((keys, vals))); - } else if let Some(pters) = pointers { - if child_idx < pters.len() { - let next_ptr = pters[child_idx]; - self.stack.push((offset, child_idx + 1)); - self.stack.push((next_ptr, 0)); - } - } - } - } -} - -impl Iterator for BPlusTreeDiskIteratorOwned -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = (K, V); - - fn next(&mut self) -> Option { - loop { - if self.leaf_idx < self.leaf_keys.len() { - let key = self.leaf_keys[self.leaf_idx].clone(); - let value = self.leaf_values[self.leaf_idx].clone(); - self.leaf_idx += 1; - return Some((key, value)); - } - - match self.next_leaf() { - Ok(Some((keys, values))) => { - self.leaf_keys = keys; - self.leaf_values = values; - self.leaf_idx = 0; - } - _ => return None, - } - } - } -} - -/// Generic reader that can be either a sorted index iterator or a regular disk iterator. -/// Used for fallback logic (Sorted -> Unsorted). -pub enum PlaylistIteratorReader { - Sorted(super::sorted_index::BPlusTreeSortedIteratorOwned), - Unsorted(BPlusTreeDiskIteratorOwned), -} - -impl Iterator for PlaylistIteratorReader -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - SortKey: for<'de> Deserialize<'de>, -{ - type Item = io::Result<(K, V)>; - - fn next(&mut self) -> Option { - match self { - PlaylistIteratorReader::Sorted(iter) => iter.next(), - PlaylistIteratorReader::Unsorted(iter) => iter.next().map(Ok), - } - } -} - -pub struct BPlusTreeDiskIterator<'a, K, V> { - query: &'a mut BPlusTreeQuery, - stack: TraversalStack, // (node_offset, next_child_index) - leaf_keys: Vec, - leaf_values: Vec, - leaf_idx: usize, -} - -impl<'a, K, V> BPlusTreeDiskIterator<'a, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - fn new(query: &'a mut BPlusTreeQuery) -> Self { - let root_offset = query.root_offset; - Self { query, stack: smallvec![(root_offset, 0)], leaf_keys: Vec::new(), leaf_values: Vec::new(), leaf_idx: 0 } - } - - pub fn try_is_empty(&mut self) -> io::Result { - loop { - if self.leaf_idx < self.leaf_keys.len() { - return Ok(false); - } - match self.next_leaf()? { - Some((keys, values)) => { - self.leaf_keys = keys; - self.leaf_values = values; - self.leaf_idx = 0; - if !self.leaf_keys.is_empty() { - return Ok(false); - } - } - None => return Ok(true), - } - } - } - - pub fn is_empty(&mut self) -> bool { self.try_is_empty().unwrap_or(true) } - - /// Internal method to load the next leaf and return its content. - fn next_leaf(&mut self) -> io::Result, Vec)>> { - loop { - let Some((offset, child_idx)) = self.stack.pop() else { return Ok(None) }; - - let (node, pointers) = if let Some(mmap) = &self.query.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? - } else if let Some(file) = &mut self.query.file { - BPlusTreeNode::::deserialize_from_block(file, &mut self.query.buffer, offset, false)? - } else { - return Err(io::Error::other("No data source available")); - }; - - if node.is_leaf { - let mut keys = Vec::with_capacity(node.keys.len()); - let mut vals = Vec::with_capacity(node.value_info.len()); - let has_tombstones = self.query.has_tombstones; - if let Some(mmap) = &self.query.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - if has_tombstones { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - if value_info.is_tombstone() { - continue; - } - let v = BPlusTreeNode::::load_value_from_info(&mut cursor, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } else { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - let v = BPlusTreeNode::::load_value_from_info(&mut cursor, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } - } else if let Some(file) = &mut self.query.file { - if has_tombstones { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - if value_info.is_tombstone() { - continue; - } - let v = BPlusTreeNode::::load_value_from_info(file, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } else { - for (key, value_info) in node.keys.iter().zip(node.value_info.iter()) { - let v = BPlusTreeNode::::load_value_from_info(file, value_info)?; - keys.push(key.clone()); - vals.push(v); - } - } - } - return Ok(Some((keys, vals))); - } else if let Some(pters) = pointers { - if child_idx < pters.len() { - let next_ptr = pters[child_idx]; - self.stack.push((offset, child_idx + 1)); - self.stack.push((next_ptr, 0)); - } - } - } - } -} - -impl Iterator for BPlusTreeDiskIterator<'_, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = (K, V); - - fn next(&mut self) -> Option { - loop { - if self.leaf_idx < self.leaf_keys.len() { - let key = self.leaf_keys[self.leaf_idx].clone(); - let value = self.leaf_values[self.leaf_idx].clone(); - self.leaf_idx += 1; - return Some((key, value)); - } - - match self.next_leaf() { - Ok(Some((keys, values))) => { - self.leaf_keys = keys; - self.leaf_values = values; - self.leaf_idx = 0; - } - Err(_err) => { - // It is possible the tree is empty or the file is being written to, so we just return None - return None; - } - _ => return None, - } - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum FlushPolicy { - /// Flush and `sync_all` after each write operation (safest, slowest). - Immediate, - /// Flush after writes; call `commit()` (or a background commit loop) to `sync_all`. - Batch, - /// Flush only. Never call `sync_all` automatically. - None, -} - -#[derive(Debug, Clone, Copy)] -struct BatchRollbackState { - file_len: u64, - root_offset: u64, -} - -pub struct BPlusTreeUpdate { - file: BufReader, - read_buffer: Vec, - write_buffer: Vec, - serial_buffer: Vec, // Reusable buffer for serialization - cache: SplitNodeCache, - root_offset: u64, - has_tombstones: bool, - inner_order: usize, - leaf_order: usize, - flush_policy: FlushPolicy, - #[allow(dead_code)] - lock: FileLock, - _marker_k: PhantomData, - _marker_v: PhantomData, -} - -fn lock_path(filepath: &Path) -> PathBuf { - if let Some(stem) = filepath.file_stem() { - // filename with dot to hide - let mut name = OsString::from("."); - name.push(stem); - name.push(".lock"); - filepath.with_file_name(name) - } else { - // Fallback: without dot - filepath.with_extension("lock") - } -} - -struct FileLock { - // We hold the file handle to keep the advisory lock active. - // When this struct is dropped, the file handle closes and OS releases the lock. - _file: File, -} - -impl FileLock { - fn try_lock(filepath: &Path) -> io::Result { - // Sidecar Lock Pattern: Lock a separate .lock file, not the data file itself. - // This ensures implementation works on Windows where locked files cannot be renamed/deleted. - let lock_path_filename = lock_path(filepath); - - let file = OpenOptions::new() - .read(true) - .write(true) - .create(true) // Create if missing - .truncate(false) // Do not truncate, just open - .open(&lock_path_filename)?; - - // Try to acquire exclusive advisory lock. - // If another process holds it, this returns immediately with Error (WouldBlock). - file.try_lock_exclusive()?; - - Ok(Self { _file: file }) - } -} -// Drop implementation is implicit: closing the _file releases the lock. -// The .lock file remains on filesystem. - -impl BPlusTreeUpdate -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - pub fn try_new(filepath: &Path) -> io::Result { - if !filepath.exists() { - return Err(io::Error::new( - io::ErrorKind::NotFound, - format!("File not found {}", filepath.to_str().unwrap_or("?")), - )); - } - // Acquire lock first - let lock = FileLock::try_lock(filepath)?; - - let f = utils::open_read_write_file(filepath)?; - - // Verify Header - let mut header = [0u8; METADATA_DATA_START_POS]; - read_exact_at_offset(&f, &mut header, 0)?; - - if &header[0..4] != MAGIC { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid magic number")); - } - let version = u32::from_le_bytes( - header[4..8].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid version slice"))?, - ); - if version != STORAGE_VERSION { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Unsupported storage version: {version}"))); - } - let root_offset = u64::from_le_bytes( - header[8..16] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset slice"))?, - ); - let metadata_len_raw = u32::from_le_bytes( - header[16..20] - .try_into() - .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length slice"))?, - ); - let (metadata_len, has_tombstones) = decode_metadata_len_and_flags(metadata_len_raw); - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - let (inner_order, leaf_order) = default_orders(); - - let file = utils::file_reader(f); - - Ok(Self { - file, - read_buffer: vec![0u8; PAGE_SIZE_USIZE], - write_buffer: vec![0u8; PAGE_SIZE_USIZE], - serial_buffer: Vec::with_capacity(PAGE_SIZE_USIZE), - cache: SplitNodeCache::new(), - root_offset, - has_tombstones, - inner_order, - leaf_order, - flush_policy: FlushPolicy::Immediate, - lock, - _marker_k: PhantomData, - _marker_v: PhantomData, - }) - } - - /// 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_millis(1); - 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; } - - #[inline] - const fn should_sync_on_write(&self) -> bool { matches!(self.flush_policy, FlushPolicy::Immediate) } - - #[inline] - const fn should_sync_on_commit(&self) -> bool { !matches!(self.flush_policy, FlushPolicy::None) } - - /// 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 should_sync = self.should_sync_on_commit(); - { - 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 should_sync { - 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(()) - } - - fn read_metadata_len_from_header(&mut self) -> io::Result { - self.file.seek(SeekFrom::Start(METADATA_OFFSET_POS))?; - let mut len_buf = [0u8; 4]; - self.file.read_exact(&mut len_buf)?; - let (metadata_len, _) = decode_metadata_len_and_flags(u32::from_le_bytes(len_buf)); - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - Ok(metadata_len) - } - - fn write_metadata_len_with_current_flags(&mut self, metadata_len: u32) -> io::Result<()> { - if metadata_len as usize > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidInput, format!("Metadata too large: {metadata_len}"))); - } - - let metadata_len_with_flags = encode_metadata_len_with_flags(metadata_len, self.has_tombstones); - self.file.seek(SeekFrom::Start(METADATA_OFFSET_POS))?; - self.file.get_mut().write_all(&metadata_len_with_flags.to_le_bytes())?; - self.file.seek_relative(0)?; - Ok(()) - } - - /// Helper to access metadata - pub fn get_metadata(&mut self) -> io::Result { - let Ok(len) = self.read_metadata_len_from_header() else { - return Ok(BPlusTreeMetadata::Empty); - }; - if len == 0 { - return Ok(BPlusTreeMetadata::Empty); - } - let mut buf = vec![0u8; len as usize]; - self.file.read_exact(&mut buf)?; - - Ok(BPlusTreeMetadata::from_bytes(&buf)) - } - - /// Helper to set metadata - pub fn set_metadata(&mut self, data: &BPlusTreeMetadata) -> io::Result<()> { - let bytes = data.to_bytes(); - if bytes.len() > METADATA_MAX_SIZE { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - format!("Metadata too large: {} > {}", bytes.len(), METADATA_MAX_SIZE), - )); - } - - self.file.seek(SeekFrom::Start(METADATA_OFFSET_POS))?; - { - let file = self.file.get_mut(); - let metadata_len = - u32::try_from(bytes.len()).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?; - let metadata_len_with_flags = encode_metadata_len_with_flags(metadata_len, self.has_tombstones); - file.write_all(&metadata_len_with_flags.to_le_bytes())?; - file.write_all(&bytes)?; - file.flush()?; // Ensure it hits disk - } - // Resync/clear reader buffer after direct writes. - self.file.seek_relative(0)?; // same as self.file.seek(SeekFrom::Current(0))?; - - Ok(()) - } - - pub fn query(&mut self, key: &K) -> Result, BPlusTreeError> { - query_tree(&mut self.file, &mut self.read_buffer, &mut self.cache, key, self.root_offset) - } - - pub fn is_empty(&mut self) -> Result { - let (node, _) = - BPlusTreeNode::::deserialize_from_block(&mut self.file, &mut self.read_buffer, self.root_offset, false)?; - Ok(node.is_leaf && node.keys.is_empty()) - } - - pub fn len(&mut self) -> Result { - count_items::(&mut self.file, &mut self.read_buffer, &mut self.cache, self.root_offset) - } - - pub fn query_le(&mut self, key: &K) -> Result, BPlusTreeError> { - query_tree_le(&mut self.file, &mut self.read_buffer, &mut self.cache, key, self.root_offset) - } - - /// Iterates over key-value pairs within a given range using `right_sibling` pointers. - /// - /// This is more efficient than iterating the full tree and filtering when you only - /// need a subset of keys. - /// - /// Tombstones are skipped automatically. - pub fn range_iter( - &mut self, - start: Bound<&K>, - end: Bound<&K>, - ) -> impl Iterator> + '_ { - // We need to clone the bounds because the returned iterator has a different lifetime - let start_cloned = match start { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let end_cloned = match end { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - // Create a temporary BPlusTreeQuery-like struct to use RangeLeafIterator - RangeLeafIteratorUpdate::new(self, start_cloned, end_cloned) - } - - /// Returns a page of key-value pairs within a given range using offset/limit. - /// - /// More efficient than `range_iter` when you only need a subset, as it stops - /// scanning once `limit` items have been collected. - /// - /// Returns `(items, has_more)` where `has_more` indicates whether additional - /// items exist beyond the returned page. - pub fn range_page( - &mut self, - start: Bound<&K>, - end: Bound<&K>, - offset: usize, - limit: usize, - ) -> Result<(Vec<(K, V)>, bool), BPlusTreeError> { - let start_cloned = match start { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let end_cloned = match end { - Bound::Included(k) => Bound::Included(k.clone()), - Bound::Excluded(k) => Bound::Excluded(k.clone()), - Bound::Unbounded => Bound::Unbounded, - }; - let mut iter = RangeLeafIteratorUpdate::new(self, start_cloned, end_cloned); - iter.skip_items(offset)?; - - let mut collected: Vec<(K, V)> = Vec::with_capacity(limit); - let mut has_more = false; - - while collected.len() < limit { - match iter.next() { - Some(Ok(item)) => collected.push(item), - Some(Err(err)) => return Err(err), - None => return Ok((collected, false)), - } - } - - match iter.next() { - Some(Ok(_)) => has_more = true, - Some(Err(err)) => return Err(err), - None => {} - } - - Ok((collected, has_more)) - } - - pub fn update(&mut self, key: &K, value: V) -> Result { - let refs = [(key, &value)]; - self.update_batch(&refs) - } - - /// Update multiple items under the v2 write contract. - /// - /// Public callers use copy-on-write mode so rollback can restore the root - /// offset and truncate appended data without leaving overwritten blocks. - /// The in-place branch is retained as an internal fast path only and must - /// not be exposed by rollback-safe public APIs. - fn update_batch_recursive( - &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)?; - - if node.is_leaf { - // Track which items need the traditional COW approach (value grew beyond allocated space) - let mut needs_cow: Vec<(usize, &K, &V)> = Vec::new(); - // Track promoted packed values (need node rewrite but value already written) - let mut promoted: Vec<(usize, ValueInfo)> = Vec::new(); - - for (key, value) in items { - match node.keys.binary_search(key) { - Ok(idx) => { - 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)); - } - } - } else { - // Batch mode uses pure COW so rollback can safely truncate/revert root. - needs_cow.push((idx, *key, *value)); - } - } - Err(_) => return Err(BPlusTreeError::KeyNotFound), - } - } - - // If an internal in-place caller changed every value without - // promotions or COW, no node rewrite is needed. - if needs_cow.is_empty() && promoted.is_empty() { - // Flush to ensure in-place writes hit disk - self.file.get_mut().flush().map_err(BPlusTreeError::Io)?; - return Ok(offset); // Return original offset - no node changes - } - - // Apply promoted value_info updates - for (idx, new_info) in promoted { - node.value_info[idx] = new_info; - } - - // Handle COW items - need to write new value blocks - for (idx, _key, value) in needs_cow { - let (val_off, val_len) = self.insert_value_to_disk(value).map_err(BPlusTreeError::Io)?; - node.value_info[idx] = - ValueInfo { mode: ValueStorageMode::Single(val_off), length: val_len, cache: Mutex::new(None) }; - } - - let new_offset = self.write_node(&node).map_err(BPlusTreeError::Io)?; - Ok(new_offset) - } else { - let mut pointers = pointers_opt - .ok_or_else(|| BPlusTreeError::InvalidStructure("Internal node missing pointers".into()))?; - let mut any_child_changed = false; - - let mut current_idx = 0; - while current_idx < items.len() { - let first_key_in_group = items[current_idx].0; - let child_idx = get_entry_index_upper_bound::(&node.keys, first_key_in_group); - - let mut group_end = current_idx + 1; - while group_end < items.len() - && get_entry_index_upper_bound::(&node.keys, items[group_end].0) == child_idx - { - group_end += 1; - } - - 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, allow_in_place)?; - - if new_child_offset != original_child_offset { - pointers[child_idx] = new_child_offset; - any_child_changed = true; - } - - current_idx = group_end; - } - - // If no child offsets changed, no need to rewrite this internal node - if !any_child_changed { - return Ok(offset); - } - - let new_offset = self.write_internal_node(&node, &pointers).map_err(BPlusTreeError::Io)?; - Ok(new_offset) - } - } - - /// 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. - /// returns The final root offset after all updates, or an error if any update fails - pub fn update_batch(&mut self, items: &[(&K, &V)]) -> Result { - if items.is_empty() { - return Ok(self.root_offset); - } - - 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)); - - // Disable in-place updates for public 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 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.should_sync_on_write() { - 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) - } - } - } - - fn delete_batch_recursive(&mut self, offset: u64, keys: &[&K]) -> io::Result<(u64, usize)> { - let (mut node, pointers_opt) = - BPlusTreeNode::::deserialize_from_block(&mut self.file, &mut self.read_buffer, offset, false)?; - - if node.is_leaf { - let mut deleted = 0usize; - for key in keys { - if let Ok(idx) = node.keys.binary_search(key) { - let already_tombstoned = node.value_info.get(idx).is_some_and(ValueInfo::is_tombstone); - if !already_tombstoned { - node.value_info[idx] = ValueInfo::tombstone(); - deleted += 1; - } - } - } - - if deleted == 0 { - return Ok((offset, 0)); - } - - let new_offset = self.write_node(&node)?; - return Ok((new_offset, deleted)); - } - - let mut pointers = pointers_opt.ok_or_else(|| io::Error::other("Internal node missing pointers"))?; - let mut total_deleted = 0usize; - let mut any_child_changed = false; - let mut current_idx = 0usize; - - while current_idx < keys.len() { - let first_key_in_group = keys[current_idx]; - let child_idx = get_entry_index_upper_bound::(&node.keys, first_key_in_group); - - let mut group_end = current_idx + 1; - while group_end < keys.len() && get_entry_index_upper_bound::(&node.keys, keys[group_end]) == child_idx { - group_end += 1; - } - - let sub_keys = &keys[current_idx..group_end]; - let original_child_offset = pointers[child_idx]; - let (new_child_offset, deleted) = self.delete_batch_recursive(original_child_offset, sub_keys)?; - total_deleted += deleted; - - if new_child_offset != original_child_offset { - pointers[child_idx] = new_child_offset; - any_child_changed = true; - } - - current_idx = group_end; - } - - if !any_child_changed { - return Ok((offset, total_deleted)); - } - - let new_offset = self.write_internal_node(&node, &pointers)?; - Ok((new_offset, total_deleted)) - } - - pub fn delete(&mut self, key: &K) -> io::Result { - let deleted = self.delete_batch(&[key])?; - Ok(deleted > 0) - } - - pub fn delete_batch(&mut self, keys: &[&K]) -> io::Result { - if keys.is_empty() { - return Ok(0); - } - - let rollback_state = self.capture_batch_rollback_state()?; - let result = (|| -> io::Result { - let mut sorted_keys = keys.to_vec(); - sorted_keys.sort(); - sorted_keys.dedup_by(|a, b| *a == *b); - - let (new_root_offset, deleted) = self.delete_batch_recursive(self.root_offset, &sorted_keys)?; - - if new_root_offset != self.root_offset { - self.file.get_mut().seek(SeekFrom::Start(ROOT_OFFSET_POS))?; - self.file.get_mut().write_all(&new_root_offset.to_le_bytes())?; - self.root_offset = new_root_offset; - } - - if deleted > 0 { - self.has_tombstones = true; - let metadata_len = self.read_metadata_len_from_header()?; - self.write_metadata_len_with_current_flags(metadata_len)?; - } - - self.file.get_mut().flush()?; - if self.should_sync_on_write() { - self.file.get_mut().sync_all()?; - } - - Ok(deleted) - })(); - - 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!( - "delete_batch failed: {err}; rollback failed: {rollback_err}" - ))); - } - Err(err) - } - } - } - - /// Try to update a value in-place if possible. - /// Returns: - /// - `Ok(Success)` if in-place update succeeded (no node rewrite needed) - /// - `Ok(PromotedToSingle(info))` if packed value was promoted to Single (node rewrite needed) - /// - `Ok(NeedsCow)` if new value doesn't fit (caller should use full COW) - /// - `Err` on I/O error - /// - /// For Single storage: updates in-place if new value fits in existing space. - /// For Packed storage: - /// - If new serialized size equals old size exactly: updates in-place within the packed block - /// - Otherwise: promotes to Single storage mode (writes at EOF, returns new `ValueInfo`) - fn try_update_value_in_place(&mut self, value: &V, existing_info: &ValueInfo) -> io::Result { - // Serialize the new value first (needed for all paths) - let raw_bytes = binary_serialize(value)?; - - match existing_info.mode { - ValueStorageMode::Single(existing_offset) => { - // Single storage: compress and check if it fits in existing space - let (flag, payload) = compress_if_beneficial(&raw_bytes); - let new_stored_len = 1 + payload.len(); // flag + payload - let existing_len = existing_info.length as usize; - - if new_stored_len > existing_len { - return Ok(InPlaceUpdateResult::NeedsCow); // Doesn't fit - } - - // Write in-place: [flag:1][payload][zero-padding to existing_len] - self.file.seek(SeekFrom::Start(existing_offset))?; - self.file.get_mut().write_all(&[flag])?; - self.file.get_mut().write_all(&payload)?; - - // Zero-pad remaining space - let mut padding_len = existing_len - new_stored_len; - while padding_len > 0 { - let chunk_len = padding_len.min(ZERO_PADDING_CHUNK.len()); - self.file.get_mut().write_all(&ZERO_PADDING_CHUNK[..chunk_len])?; - padding_len -= chunk_len; - } - - // Force BufReader to discard its internal buffer by seeking to - // the underlying file position. stream_position() / seek(Current(0)) - // does NOT reliably invalidate stale buffered data. - let pos = self.file.get_mut().stream_position()?; - self.file.seek(SeekFrom::Start(pos))?; - - Ok(InPlaceUpdateResult::Success) - } - ValueStorageMode::Packed(block_offset, value_index) => { - // Packed storage: raw_bytes is the MessagePack-serialized value (no compression for packed) - let new_len = raw_bytes.len(); - let existing_len = existing_info.length as usize; - - if new_len == existing_len { - // Same size: can update in-place within the packed block - self.update_packed_value_in_place(block_offset, value_index, &raw_bytes, &existing_info.cache)?; - Ok(InPlaceUpdateResult::Success) - } else { - // Different size: promote to Single storage mode - // Write value at EOF as Single (with compression) - let (val_offset, val_len) = self.insert_value_to_disk(value)?; - let new_info = ValueInfo { - mode: ValueStorageMode::Single(val_offset), - length: val_len, - cache: Mutex::new(None), - }; - Ok(InPlaceUpdateResult::PromotedToSingle(new_info)) - } - } - ValueStorageMode::Tombstone => Ok(InPlaceUpdateResult::NeedsCow), - } - } - - /// Update a value in-place within a packed block. - /// The new value must have the exact same serialized size as the existing value. - fn update_packed_value_in_place( - &mut self, - block_offset: u64, - value_index: u16, - new_value_bytes: &[u8], - cache: &Mutex>, - ) -> io::Result<()> { - // Optimization: Use cached offset if available to skip read and scan - let cached_pos = { - let guard = cache.lock(); - if let Some(CacheData::PackedOffset(pos)) = guard.as_ref() { - Some(*pos) - } else { - None - } - }; - - if let Some(pos) = cached_pos { - // Jump directly to the offset (avoid reading entire block and linear scan) - self.file.seek(SeekFrom::Start(block_offset + u64::from(pos)))?; - self.file.get_mut().write_all(new_value_bytes)?; - // Force BufReader to discard its internal buffer. - let pos = self.file.get_mut().stream_position()?; - self.file.seek(SeekFrom::Start(pos))?; - return Ok(()); - } - - // Read the entire packed block - self.file.seek(SeekFrom::Start(block_offset))?; - let mut block_buffer = vec![0u8; PAGE_SIZE_USIZE]; - self.file.read_exact(&mut block_buffer)?; - - // Navigate to the target value's position - // Format: [COUNT:4B][LEN:4B][data...][LEN:4B][data...]... - let mut pos = 4; // Skip count - - for i in 0..=value_index { - if pos + 4 > PAGE_SIZE_USIZE { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Packed block corrupted: position {pos} exceeds block size"), - )); - } - - let len = u32::from_le_bytes(block_buffer[pos..pos + 4].try_into().map_err(to_io_error)?) as usize; - pos += 4; - - if i == value_index { - // Found target value - verify size matches - if len != new_value_bytes.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - format!("Size mismatch in packed update: expected {len}, got {}", new_value_bytes.len()), - )); - } - - // Update the value data in the buffer - block_buffer[pos..pos + len].copy_from_slice(new_value_bytes); - - // Write the entire block back (optimization: we could write just the slice, - // but we already have the whole block in-memory here and write_all(4KB) is fast) - self.file.seek(SeekFrom::Start(block_offset))?; - self.file.get_mut().write_all(&block_buffer)?; - - // Cache the offset for future updates - *cache.lock() = Some(CacheData::PackedOffset(u16::try_from(pos).map_err(to_io_error)?)); - - // Force BufReader to discard its internal buffer. - let pos = self.file.get_mut().stream_position()?; - self.file.seek(SeekFrom::Start(pos))?; - - return Ok(()); - } - - pos += len; - } - - Err(io::Error::new(io::ErrorKind::InvalidData, format!("Value index {value_index} not found in packed block"))) - } - - /// Insert or update multiple items in batch (upsert). If a key exists, it will be updated; - /// if it doesn't exist, it will be inserted. This is more efficient than calling `update()` - /// or `insert()` multiple times as it loads the tree once, performs all operations, and saves once. - /// returns The final root offset after all upserts, or an error if any operation fails - fn insert_value_to_disk(&mut self, value: &V) -> io::Result<(u64, u32)> { - let raw_bytes = binary_serialize(value)?; - - // Decide whether to compress based on size and effectiveness - let (flag, payload) = compress_if_beneficial(&raw_bytes); - - self.file.get_mut().seek(SeekFrom::End(0))?; - let offset = self.file.get_mut().stream_position()?; - - // Write: [flag:1][payload] - self.file.get_mut().write_all(&[flag])?; - // NOTE: For LZ4, payload already includes original length (prepended) - self.file.get_mut().write_all(&payload)?; - - // stored_len includes flag + payload - let stored_len = 1 + payload.len(); - Ok((offset, u32::try_from(stored_len).map_err(to_io_error)?)) - } - - fn write_node(&mut self, node: &BPlusTreeNode) -> io::Result { - self.file.get_mut().seek(SeekFrom::End(0))?; - let offset = self.file.get_mut().stream_position()?; - self.serial_buffer.clear(); - node.serialize_to_block(self.file.get_mut(), &mut self.write_buffer, &mut self.serial_buffer, offset)?; - Ok(offset) - } - - fn write_internal_node(&mut self, node: &BPlusTreeNode, pointers: &[u64]) -> io::Result { - self.file.get_mut().seek(SeekFrom::End(0))?; - let offset = self.file.get_mut().stream_position()?; - self.serial_buffer.clear(); - node.serialize_internal_with_offsets( - self.file.get_mut(), - &mut self.write_buffer, - &mut self.serial_buffer, - offset, - pointers, - )?; - Ok(offset) - } - - fn upsert_batch_recursive(&mut self, offset: u64, items: &[(&K, &V)]) -> io::Result<(u64, Vec<(K, u64)>)> { - let (mut node, pointers_opt) = BPlusTreeNode::::deserialize_from_block( - &mut self.file, - &mut self.read_buffer, - offset, - false, // shallow - )?; - - if node.is_leaf { - for (key, value) in items { - let (val_off, val_len) = self.insert_value_to_disk(value)?; - let new_info = - ValueInfo { mode: ValueStorageMode::Single(val_off), length: val_len, cache: Mutex::new(None) }; - - match node.keys.binary_search(key) { - Ok(idx) => { - node.value_info[idx] = new_info; - } - Err(idx) => { - node.keys.insert(idx, (*key).clone()); - node.value_info.insert(idx, new_info); - } - } - } - - let mut leaf_promotions = Vec::new(); - while node.keys.len() > self.leaf_order { - let median_idx = node.keys.len() / 2; - let mut right_node = BPlusTreeNode::new(true); - right_node.keys = node.keys.split_off(median_idx); - right_node.value_info = node.value_info.split_off(median_idx); - - let promoted_key = right_node.keys[0].clone(); - let right_offset = self.write_node(&right_node)?; - leaf_promotions.push((promoted_key, right_offset)); - } - - let new_leaf_offset = self.write_node(&node)?; - Ok((new_leaf_offset, leaf_promotions)) - } else { - let mut pointers = pointers_opt.ok_or_else(|| io::Error::other("Internal node missing pointers"))?; - let mut current_idx = 0; - - while current_idx < items.len() { - let first_key_in_group = items[current_idx].0; - let child_idx = get_entry_index_upper_bound::(&node.keys, first_key_in_group); - - let mut group_end = current_idx + 1; - while group_end < items.len() - && get_entry_index_upper_bound::(&node.keys, items[group_end].0) == child_idx - { - group_end += 1; - } - - let sub_items = &items[current_idx..group_end]; - let (new_child_offset, child_promotions) = - self.upsert_batch_recursive(pointers[child_idx], sub_items)?; - pointers[child_idx] = new_child_offset; - - for (median_key, right_child_offset) in child_promotions { - let insert_idx = get_entry_index_upper_bound::(&node.keys, &median_key); - node.keys.insert(insert_idx, median_key); - pointers.insert(insert_idx + 1, right_child_offset); - } - - current_idx = group_end; - } - - let mut node_promotions = Vec::new(); - while node.keys.len() > self.inner_order { - let median_idx = node.keys.len() / 2; - let mut right_node = BPlusTreeNode::new(false); - - let promoted_key = node.keys.remove(median_idx); - right_node.keys = node.keys.split_off(median_idx); - let right_pointers = pointers.split_off(median_idx + 1); - - let right_offset = self.write_internal_node(&right_node, &right_pointers)?; - node_promotions.push((promoted_key, right_offset)); - } - - let new_offset = self.write_internal_node(&node, &pointers)?; - Ok((new_offset, node_promotions)) - } - } - - fn build_higher_levels(&mut self, base_offset: u64, mut promotions: Vec<(K, u64)>) -> io::Result { - if promotions.is_empty() { - return Ok(base_offset); - } - promotions.sort_by(|a, b| a.0.cmp(&b.0)); - - let mut node = BPlusTreeNode::::new(false); - let mut pointers = vec![base_offset]; - for (key, ptr) in promotions { - node.keys.push(key); - pointers.push(ptr); - } - - if node.keys.len() <= self.inner_order { - return self.write_internal_node(&node, &pointers); - } - - let mut next_level_promotions = Vec::new(); - while node.keys.len() > self.inner_order { - let median_idx = node.keys.len() / 2; - let mut right_node = BPlusTreeNode::new(false); - let promoted_key = node.keys.remove(median_idx); - right_node.keys = node.keys.split_off(median_idx); - let right_pointers = pointers.split_off(median_idx + 1); - let right_offset = self.write_internal_node(&right_node, &right_pointers)?; - next_level_promotions.push((promoted_key, right_offset)); - } - let left_offset = self.write_internal_node(&node, &pointers)?; - self.build_higher_levels(left_offset, next_level_promotions) - } - - /// Insert or update multiple items in batch (upsert). - /// Uses disk-based recursive traversal for efficiency, processing each node only once. - pub fn upsert_batch(&mut self, items: &[(&K, &V)]) -> io::Result { - if items.is_empty() { - return Ok(self.root_offset); - } - - 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)?; - - // 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.should_sync_on_write() { - 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) - } - } - } - - 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. - /// - /// This method is designed for use with `spawn_blocking` where you want to: - /// 1. Serialize values in the async context (before `spawn_blocking`) - /// 2. Pass only `Vec` bytes into the blocking context (avoiding clones) - /// 3. Perform all I/O in a single blocking call - /// - /// The key type K still follows the tree bounds (Ord + Serialize + Deserialize), - /// but values are written as raw bytes without re-serialization. - /// - /// # Arguments - /// * `items` - (`key`, `value_bytes`) pairs where `value_bytes` is MessagePack-encoded. - /// Keys are already typed and used directly for tree traversal. - /// - /// # Returns - /// The final root offset after all upserts, or an error if any operation fails - pub fn upsert_batch_preserialized(&mut self, items: Vec<(K, Vec)>) -> io::Result { - if items.is_empty() { - return Ok(self.root_offset); - } - - 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.should_sync_on_write() { - 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) - } - } - } - - pub fn commit(&mut self) -> io::Result<()> { - self.file.get_mut().flush()?; - if self.should_sync_on_commit() { - self.file.get_mut().sync_all()?; - } - Ok(()) - } - - /// Recursive helper for `upsert_batch_preserialized`. - /// Items is a slice of (key, pre-serialized value bytes). - fn upsert_batch_preserialized_recursive( - &mut self, - offset: u64, - items: &[(K, Vec)], - ) -> io::Result<(u64, Vec<(K, u64)>)> { - let (mut node, pointers_opt) = BPlusTreeNode::::deserialize_from_block( - &mut self.file, - &mut self.read_buffer, - offset, - false, // shallow - )?; - - if node.is_leaf { - for (key, value_bytes) in items { - // Write pre-serialized value to disk with compression - let (val_off, val_len) = self.insert_preserialized_value_to_disk(value_bytes)?; - let new_info = - ValueInfo { mode: ValueStorageMode::Single(val_off), length: val_len, cache: Mutex::new(None) }; - - match node.keys.binary_search(key) { - Ok(idx) => { - node.value_info[idx] = new_info; - } - Err(idx) => { - node.keys.insert(idx, key.clone()); - node.value_info.insert(idx, new_info); - } - } - } - - let mut leaf_promotions = Vec::new(); - while node.keys.len() > self.leaf_order { - let median_idx = node.keys.len() / 2; - let mut right_node = BPlusTreeNode::new(true); - right_node.keys = node.keys.split_off(median_idx); - right_node.value_info = node.value_info.split_off(median_idx); - - let promoted_key = right_node.keys[0].clone(); - let right_offset = self.write_node(&right_node)?; - leaf_promotions.push((promoted_key, right_offset)); - } - - let new_leaf_offset = self.write_node(&node)?; - Ok((new_leaf_offset, leaf_promotions)) - } else { - let mut pointers = pointers_opt.ok_or_else(|| io::Error::other("Internal node missing pointers"))?; - let mut current_idx = 0; - - while current_idx < items.len() { - let first_key_in_group = &items[current_idx].0; - let child_idx = get_entry_index_upper_bound::(&node.keys, first_key_in_group); - - let mut group_end = current_idx + 1; - while group_end < items.len() - && get_entry_index_upper_bound::(&node.keys, &items[group_end].0) == child_idx - { - group_end += 1; - } - - let sub_items = &items[current_idx..group_end]; - let (new_child_offset, child_promotions) = - self.upsert_batch_preserialized_recursive(pointers[child_idx], sub_items)?; - pointers[child_idx] = new_child_offset; - - for (median_key, right_child_offset) in child_promotions { - let insert_idx = get_entry_index_upper_bound::(&node.keys, &median_key); - node.keys.insert(insert_idx, median_key); - pointers.insert(insert_idx + 1, right_child_offset); - } - - current_idx = group_end; - } - - let mut node_promotions = Vec::new(); - while node.keys.len() > self.inner_order { - let median_idx = node.keys.len() / 2; - let mut right_node = BPlusTreeNode::new(false); - - let promoted_key = node.keys.remove(median_idx); - right_node.keys = node.keys.split_off(median_idx); - let right_pointers = pointers.split_off(median_idx + 1); - - let right_offset = self.write_internal_node(&right_node, &right_pointers)?; - node_promotions.push((promoted_key, right_offset)); - } - - let new_offset = self.write_internal_node(&node, &pointers)?; - Ok((new_offset, node_promotions)) - } - } - - /// Insert a pre-serialized value (already `MessagePack` encoded) to disk. - /// Applies compression if beneficial. - fn insert_preserialized_value_to_disk(&mut self, value_bytes: &[u8]) -> io::Result<(u64, u32)> { - // Apply compression if beneficial - let (flag, payload) = compress_if_beneficial(value_bytes); - - self.file.get_mut().seek(SeekFrom::End(0))?; - let offset = self.file.get_mut().stream_position()?; - - // Write: [flag:1][payload] - self.file.get_mut().write_all(&[flag])?; - self.file.get_mut().write_all(&payload)?; - - let stored_len = 1 + payload.len(); - Ok((offset, u32::try_from(stored_len).map_err(to_io_error)?)) - } - - /// Garbage Collection: Compacts the file by rewriting only live blocks sequentially. - pub fn compact(&mut self, filepath: &Path) -> io::Result<()> { - let mut temp_file = NamedTempFile::new_in(filepath.parent().unwrap_or(Path::new(".")))?; - - // 1. Read existing metadata from source (manually to avoid full load) - let mut metadata = Vec::new(); - { - if let Ok(mut src) = File::open(filepath) { - if src.seek(SeekFrom::Start(METADATA_OFFSET_POS)).is_ok() { - let mut lbuf = [0u8; 4]; - if src.read_exact(&mut lbuf).is_ok() { - let (l, _) = decode_metadata_len_and_flags(u32::from_le_bytes(lbuf)); - if l > 0 && l as usize <= METADATA_MAX_SIZE { - let mut b = vec![0u8; l as usize]; - if src.read_exact(&mut b).is_ok() { - metadata = b; - } else { - error!("Failed to read metadata bytes during compaction"); - } - } else if l > u32::try_from(METADATA_MAX_SIZE).unwrap_or(4000) { - error!("Metadata too large during compaction: {l}"); - } - } - } - } - } - - // 2. Write Header placeholder - temp_file.seek(SeekFrom::Start(0))?; - // Construct full header block - let mut header = [0u8; PAGE_SIZE_USIZE]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&STORAGE_VERSION.to_le_bytes()); - // Root offset placeholder (will be filled later) - // Root offset placeholder (will be filled later) - let metadata_len = u32::try_from(metadata.len()).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - let metadata_len_with_flags = encode_metadata_len_with_flags(metadata_len, false); - header[16..20].copy_from_slice(&metadata_len_with_flags.to_le_bytes()); - if !metadata.is_empty() { - header[METADATA_DATA_START_POS..METADATA_DATA_START_POS + metadata.len()].copy_from_slice(&metadata); - } - temp_file.write_all(&header)?; - - // temp_file.seek(SeekFrom::Start(HEADER_SIZE))?; // Skip to start of content (implied by write_all(4096)) - - let mut current_offset = HEADER_SIZE; - let mut leaf_pointers: Vec<(K, u64)> = Vec::new(); - let mut current_leaf = BPlusTreeNode::::new(true); - - // 2. Iterate source and write values + leaf nodes immediately (Streaming) - { - let mut query = BPlusTreeQuery::::try_new(filepath).map_err(to_io_error)?; - let mut write_buffer = std::io::BufWriter::new(&mut temp_file); - let mut node_buffer = vec![0u8; PAGE_SIZE_USIZE]; - let mut serial_buf = Vec::with_capacity(PAGE_SIZE_USIZE); - - for (k, v) in query.iter() { - let value_bytes = binary_serialize(&v)?; - let val_offset = current_offset; - - // Write value header (flag + payload) - let (flag, payload) = compress_if_beneficial(&value_bytes); - write_buffer.write_all(&[flag])?; - write_buffer.write_all(&payload)?; - - let stored_len = u32::try_from(1 + payload.len()).map_err(to_io_error)?; - current_offset += u64::from(stored_len); - - current_leaf.keys.push(k); - current_leaf.value_info.push(ValueInfo { - mode: ValueStorageMode::Single(val_offset), - length: stored_len, - cache: Mutex::new(None), - }); - - if current_leaf.keys.len() >= self.leaf_order { - let first_key = current_leaf.keys[0].clone(); - let node_offset = current_offset; - current_offset = current_leaf.serialize_to_block( - &mut write_buffer, - &mut node_buffer, - &mut serial_buf, - node_offset, - )?; - leaf_pointers.push((first_key, node_offset)); - current_leaf = BPlusTreeNode::new(true); - } - } - - // Handle trailing leaf - if !current_leaf.keys.is_empty() { - let first_key = current_leaf.keys[0].clone(); - let node_offset = current_offset; - current_offset = current_leaf.serialize_to_block( - &mut write_buffer, - &mut node_buffer, - &mut serial_buf, - node_offset, - )?; - leaf_pointers.push((first_key, node_offset)); - } - write_buffer.flush()?; - } - - // 3. Build Internal Levels - let tree = BPlusTree::::new(); - let mut node_buffer = vec![0u8; PAGE_SIZE_USIZE]; - let root_offset = - tree.build_levels_from_pointers(&mut temp_file, leaf_pointers, current_offset, &mut node_buffer)?; - - // 4. Update Header - temp_file.seek(SeekFrom::Start(ROOT_OFFSET_POS))?; - temp_file.write_all(&root_offset.to_le_bytes())?; - - temp_file.flush()?; - // 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)?; - - // 6. Refresh state - self.root_offset = root_offset; - self.has_tombstones = false; - let file = utils::open_read_write_file(filepath)?; - self.file = utils::file_reader(file); - self.cache.clear(); - - Ok(()) - } -} - -/// 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: Arc>>, - flush_policy: FlushPolicy, - dirty: Arc, - background_commit_shutdown: Arc, - background_commit_handle: 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: Arc::new(Mutex::new(updater)), - flush_policy, - dirty: Arc::new(AtomicBool::new(false)), - background_commit_shutdown: Arc::new(AtomicBool::new(false)), - background_commit_handle: Mutex::new(None), - }) - } - - pub fn upsert_prepared(&self, items: Vec<(K, Vec)>) -> io::Result { - let result = self.updater.lock().upsert_batch_encoded(items); - if result.is_ok() { - self.mark_dirty_after_write(); - } - result - } - - pub fn upsert(&self, items: &[(&K, &V)]) -> io::Result { - let prepared = BPlusTreeUpdate::::prepare_upsert_batch(items)?; - self.upsert_prepared(prepared) - } - - #[inline] - fn mark_dirty_after_write(&self) { - match self.flush_policy { - FlushPolicy::Batch => self.dirty.store(true, Ordering::Release), - FlushPolicy::Immediate | FlushPolicy::None => self.dirty.store(false, Ordering::Release), - } - } - - pub fn start_background_commit(&self, interval: Duration) -> io::Result<()> { - if self.flush_policy != FlushPolicy::Batch { - return Err(io::Error::new(io::ErrorKind::InvalidInput, "background commit requires FlushPolicy::Batch")); - } - if interval.is_zero() { - return Err(io::Error::new(io::ErrorKind::InvalidInput, "background commit interval must be > 0")); - } - - let mut handle_slot = self.background_commit_handle.lock(); - if handle_slot.is_some() { - return Ok(()); - } - - self.background_commit_shutdown.store(false, Ordering::Release); - let updater = Arc::clone(&self.updater); - let dirty = Arc::clone(&self.dirty); - let shutdown = Arc::clone(&self.background_commit_shutdown); - - let handle = std::thread::Builder::new() - .name("bplustree-commit".to_string()) - .spawn(move || { - while !shutdown.load(Ordering::Acquire) { - std::thread::park_timeout(interval); - if shutdown.load(Ordering::Acquire) { - break; - } - if !dirty.swap(false, Ordering::AcqRel) { - continue; - } - if let Err(err) = updater.lock().commit() { - error!("Background B+Tree commit failed: {err}"); - dirty.store(true, Ordering::Release); - } - } - - if dirty.swap(false, Ordering::AcqRel) { - if let Err(err) = updater.lock().commit() { - error!("Final background B+Tree commit failed during shutdown: {err}"); - } - } - }) - .map_err(io::Error::other)?; - *handle_slot = Some(handle); - Ok(()) - } - - pub fn stop_background_commit(&self) -> io::Result<()> { - self.background_commit_shutdown.store(true, Ordering::Release); - if let Some(handle) = self.background_commit_handle.lock().take() { - handle.thread().unpark(); - handle.join().map_err(|_| io::Error::other("background B+Tree commit thread panicked"))?; - } - Ok(()) - } - - /// Explicit durability barrier. - pub fn flush_now(&self) -> io::Result<()> { self.commit() } - - pub fn commit(&self) -> io::Result<()> { - let result = self.updater.lock().commit(); - if result.is_ok() { - self.dirty.store(false, Ordering::Release); - } - result - } - - /// Alias for `commit()`. - pub fn shutdown(&self) -> io::Result<()> { - self.stop_background_commit()?; - self.commit() - } -} - -impl Drop for BPlusTreeSerialWriter { - fn drop(&mut self) { - self.background_commit_shutdown.store(true, Ordering::Release); - if self.dirty.load(Ordering::Acquire) { - warn!("Dropping dirty B+Tree serial writer without explicit shutdown; pending batch writes may be uncommitted"); - } - if let Some(handle) = self.background_commit_handle.lock().take() { - handle.thread().unpark(); - // Do not block indefinitely from Drop. Call shutdown() explicitly - // when the caller needs a synchronous final commit barrier. - drop(handle); - } - } -} - -#[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]>, - current_values: Option<&'a [V]>, - index: usize, -} - -impl<'a, K, V> BPlusTreeIterator<'a, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - pub fn new(tree: &'a BPlusTree) -> Self { - let stack = vec![&tree.root]; - Self { stack, current_keys: None, current_values: None, index: 0 } - } -} - -impl<'a, K, V> Iterator for BPlusTreeIterator<'a, K, V> -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = (&'a K, &'a V); - - fn next(&mut self) -> Option { - loop { - // Try to return next item from current leaf - if let Some(keys) = self.current_keys { - if let Some(values) = self.current_values { - if self.index < keys.len() { - let key = &keys[self.index]; - let value = &values[self.index]; - self.index += 1; - return Some((key, value)); - } - } - } - - // Current leaf exhausted, find next leaf - loop { - let node = self.stack.pop()?; - - if node.is_leaf { - // Found a leaf node - self.current_keys = Some(&node.keys); - self.current_values = Some(&node.values); - self.index = 0; - break; // Exit inner loop to process this leaf - } - // Push children in reverse order to maintain left-to-right traversal - for child in node.children.iter().rev() { - self.stack.push(child); - } - } - } - } -} - -impl BPlusTree -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - pub fn iter(&self) -> BPlusTreeIterator<'_, K, V> { BPlusTreeIterator::new(self) } -} - -impl<'a, K, V> IntoIterator for &'a BPlusTree -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, -{ - type Item = (&'a K, &'a V); - type IntoIter = BPlusTreeIterator<'a, K, V>; - - fn into_iter(self) -> Self::IntoIter { self.iter() } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::repository::bplustree::{BPlusTree, BPlusTreeQuery, BPlusTreeUpdate}; - use parking_lot::Mutex; - use serde::{de::Deserializer, ser::Error as SerError, Deserialize, Serialize, Serializer}; - use shared::{model::UUIDType, utils::generate_random_string}; - use std::{collections::HashSet, io, sync::Arc}; - use tempfile::tempdir; - - #[cfg(unix)] - 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 { - id: u32, - 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<()> { - let test_size = 500; - let content = generate_random_string(1024); - let mut tree = BPlusTree::::new(); - for i in 0u32..=test_size { - tree.insert(i, Record { id: i, data: format!("{content} {i}") }); - } - - // // Traverse the tree - // tree.traverse(|node| { - // println!("Node: {:?}", node); - // }); - - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_insert_test.bin"); - // Serialize the tree to a file - tree.store(&filepath)?; - // Deserialize the tree from the file - tree = BPlusTree::::load(&filepath)?; - - // Query the tree - for i in 0u32..=test_size { - let found = tree.query(&i); - assert!(found.is_some(), "{content} {i} not found"); - assert!(found.unwrap().eq(&Record { id: i, data: format!("{content} {i}") }), "{content} {i} not found"); - } - - let mut tree_query: BPlusTreeQuery = BPlusTreeQuery::try_new(&filepath)?; - for i in 0u32..=test_size { - let found = tree_query.query(&i).expect("Query failed"); - assert!(found.is_some(), "{content} {i} not found"); - let entry = found.unwrap(); - assert!(entry.eq(&Record { id: i, data: format!("{content} {i}") }), "{content} {i} not found"); - } - - let mut tree_update: BPlusTreeUpdate = BPlusTreeUpdate::try_new(&filepath)?; - - for i in 0u32..=test_size { - if let Ok(Some(record)) = tree_update.query(&i) { - let new_record = Record { id: record.id, data: format!("{content} {}", record.id + 9000) }; - tree_update.update(&i, new_record).map_err(BPlusTreeError::to_io)?; - } else { - panic!("{content} {i} not found"); - } - } - - // Verify with Query - let mut tree_query: BPlusTreeQuery = BPlusTreeQuery::try_new(&filepath)?; - - for i in 0u32..=test_size { - let found = tree_query.query(&i).expect("Query failed"); - assert!(found.is_some(), "{content} {i} not found"); - let entry = found.unwrap(); - let expected = Record { id: i, data: format!("{content} {}", i + 9000) }; - assert!(entry.eq(&expected), "Entry not equal {entry:?} != {expected:?}"); - } - - Ok(()) - } - - #[test] - fn insert_duplicate_test() { - let content = "Entry"; - let mut tree = BPlusTree::::new(); - for i in 0u32..=500 { - tree.insert(i, Record { id: i, data: format!("{content} {i}") }); - } - for i in 0u32..=500 { - tree.insert(i, Record { id: i, data: format!("{content} {}", i + 1) }); - } - - tree.traverse(|keys, values| { - keys.iter().zip(values.iter()).for_each(|(k, v)| { - assert!(format!("{content} {}", k + 1).eq(&v.data), "Wrong entry"); - }); - }); - } - - #[test] - fn test_upsert_batch() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("upsert_batch_test.bin"); - - // 1. Create initial tree - let mut tree = BPlusTree::::new(); - tree.insert(1, Record { id: 1, data: "original 1".to_string() }); - tree.insert(2, Record { id: 2, data: "original 2".to_string() }); - tree.store(&filepath)?; - - // 2. Open for update and upsert batch - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - let r1_new = Record { id: 1, data: "updated 1".to_string() }; - let r3_new = Record { id: 3, data: "new 3".to_string() }; - - update.upsert_batch(&[(&1, &r1_new), (&3, &r3_new)])?; - - // 3. Verify with query - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.query(&1).unwrap(), Some(r1_new)); - assert_eq!(query.query(&2).unwrap(), Some(Record { id: 2, data: "original 2".to_string() })); - assert_eq!(query.query(&3).unwrap(), Some(r3_new)); - - Ok(()) - } - - #[test] - fn len_test() -> io::Result<()> { - let test_size = 100; - let mut tree = BPlusTree::::new(); - - // Initial state - assert_eq!(tree.len(), 0); - assert!(tree.is_empty()); - - for i in 1..=test_size { - tree.insert(i, Record { id: i, data: format!("data {i}") }); - assert_eq!(tree.len(), i as usize); - assert!(!tree.is_empty()); - } - - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("len_test.bin"); - tree.store(&filepath)?; - - // Test BPlusTreeQuery len - let mut query: BPlusTreeQuery = BPlusTreeQuery::try_new(&filepath)?; - assert_eq!(query.len().expect("Query len failed"), test_size as usize); - assert!(!query.is_empty().expect("Query is_empty failed")); - - // Test BPlusTreeUpdate len and modifications - let mut update: BPlusTreeUpdate = BPlusTreeUpdate::try_new(&filepath)?; - assert_eq!(update.len().expect("Update len failed"), test_size as usize); - - // Update existing key - length should stay same - update.update(&1, Record { id: 1, data: "updated".to_string() }).map_err(BPlusTreeError::to_io)?; - assert_eq!(update.len().expect("Update len failed after update"), test_size as usize); - - // Insert new key - length should increase - update.upsert_batch(&[(&(test_size + 1), &Record { id: test_size + 1, data: "new".to_string() })])?; - assert_eq!(update.len().expect("Update len failed after insert"), (test_size + 1) as usize); - - Ok(()) - } - - #[test] - fn iterator_test() -> io::Result<()> { - let mut tree = BPlusTree::::new(); - let mut entry_set = HashSet::new(); - for i in 0u32..=500 { - tree.insert(i, Record { id: i, data: format!("Entry {i}") }); - entry_set.insert(i); - } - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_iterator_test.bin"); - // Serialize the tree to a file - tree.store(&filepath)?; - let tree: BPlusTree = BPlusTree::load(&filepath)?; - - // Traverse the tree - for (key, value) in &tree { - assert!(format!("Entry {key}").eq(&value.data), "Wrong entry"); - entry_set.remove(key); - } - assert!(entry_set.is_empty()); - Ok(()) - } - - #[test] - fn persistence_update_and_iterate_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_update_iter.bin"); - let content = "InitialContent"; - let mut tree = BPlusTree::::new(); - - // Initial store - for i in 0u32..100 { - tree.insert(i, Record { id: i, data: format!("{content} {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - // Update via BPlusTreeUpdate - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - for i in 0u32..100 { - if i % 2 == 0 { - tree_update - .update(&i, Record { id: i, data: format!("UpdatedContent {i}") }) - .map_err(BPlusTreeError::to_io)?; - } - } - - // Reload and Verify via Query - let mut tree_query: BPlusTreeQuery = BPlusTreeQuery::try_new(&filepath)?; - for i in 0u32..100 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find key"); - if i % 2 == 0 { - assert_eq!(val.data, format!("UpdatedContent {i}")); - } else { - assert_eq!(val.data, format!("{content} {i}")); - } - } - - // Reload and Verify via Iterator - let reloaded_tree = BPlusTree::::load(&filepath)?; - let mut count = 0; - for (key, value) in &reloaded_tree { - if *key % 2 == 0 { - assert_eq!( - value.data, - format!("UpdatedContent {key}"), - "Iterator returned wrong value for updated key {key}" - ); - } else { - assert_eq!( - value.data, - format!("{content} {key}"), - "Iterator returned wrong value for original key {key}" - ); - } - count += 1; - } - assert_eq!(count, 100, "Iterator did not yield all entries"); - - Ok(()) - } - - #[test] - fn update_cow_size_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_size_test.bin"); - let mut tree = BPlusTree::::new(); - - // Use incompressible data > SMALL_VALUE_THRESHOLD (256 bytes) to ensure Single storage. - // Repetitive strings like "x".repeat(300) compress too well with LZ4 and end up - // smaller than 256 bytes, causing them to be stored as Packed (not Single). - // Random-looking strings don't compress well and stay above the threshold. - let padding: String = generate_random_string(400); - for i in 0u32..10 { - // Each record has unique data to prevent compression - tree.insert(i, Record { id: i, data: format!("{padding}{i}") }); - } - tree.store(&filepath)?; - - let initial_size = std::fs::metadata(&filepath)?.len(); - - // Public v2 update() uses COW semantics even when an equal-size in-place - // rewrite would be possible. This keeps rollback behavior honest. - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - let same_size_padding: String = generate_random_string(400); - for i in 0u32..10 { - tree_update - .update(&i, Record { id: i, data: format!("{same_size_padding}{i}") }) - .map_err(BPlusTreeError::to_io)?; - } - - let size_after_same_update = std::fs::metadata(&filepath)?.len(); - assert!(size_after_same_update > initial_size, "Same-size public updates should use COW and grow the file"); - drop(tree_update); - - // Reload and verify the COW updates worked - let mut tree_query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0u32..10 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find key"); - assert!(val.data.starts_with(&same_size_padding), "Updated value should contain new padding"); - } - drop(tree_query); - - // Smaller updates also stay on the public COW path. - // Using 200-char random string (smaller than 400 but > threshold for incompressibility) - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - let smaller_padding: String = generate_random_string(200); - for i in 0u32..10 { - tree_update - .update(&i, Record { id: i, data: format!("{smaller_padding}{i}") }) - .map_err(BPlusTreeError::to_io)?; - } - - let size_after_smaller_update = std::fs::metadata(&filepath)?.len(); - assert!( - size_after_smaller_update > size_after_same_update, - "Smaller public updates should use COW and continue appending" - ); - - // Update with larger size data - should trigger COW and file growth - // 5000 chars is much larger than the original ~400 byte allocation - let larger_padding: String = generate_random_string(5000); - for i in 0u32..1 { - tree_update - .update(&i, Record { id: i, data: format!("{larger_padding}{i}") }) - .map_err(BPlusTreeError::to_io)?; - } - - let size_after_larger_update = std::fs::metadata(&filepath)?.len(); - assert!(size_after_larger_update > size_after_smaller_update, "Larger updates should use COW and grow the file"); - - // Final verification: Compact should shrink the file - tree_update.compact(&filepath)?; - let size_after_compact = std::fs::metadata(&filepath)?.len(); - assert!(size_after_compact < size_after_larger_update, "Compaction should reduce file size"); - drop(tree_update); - - // Final data check after compact - let mut final_query = BPlusTreeQuery::::try_new(&filepath)?; - // Key 0 was updated with larger padding - assert!(final_query.query(&0).unwrap().unwrap().data.starts_with(&larger_padding)); - // Keys 1-9 were updated with smaller padding - for i in 1u32..10 { - assert!(final_query.query(&i).unwrap().unwrap().data.starts_with(&smaller_padding)); - } - - Ok(()) - } - - #[test] - fn cow_deep_tree_compaction_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("deep_tree.idx"); - - let test_size = 500u32; // Enough to force multiple levels - let mut tree = BPlusTree::new(); - for i in 0..test_size { - tree.insert(i, Record { id: i, data: format!("Content {i}") }); - } - tree.store(&filepath)?; - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // 1. Initial Queries - for i in (0..test_size).step_by(50) { - let val = tree_update.query(&i).expect("Query failed").expect("Should find initial key"); - assert_eq!(val.data, format!("Content {i}")); - } - - // 2. Multiple Updates (COW) - for i in (0..test_size).step_by(10) { - tree_update - .update(&i, Record { id: i, data: format!("UpdatedContent {i}") }) - .map_err(BPlusTreeError::to_io)?; - } - - // 3. Verify Query Integrity (Must return NEW values) - for i in (0..test_size).step_by(10) { - let val = tree_update.query(&i).expect("Query failed").expect("Should find updated key"); - assert_eq!(val.data, format!("UpdatedContent {i}")); - } - - // 4. Verify Query Integrity for non-updated keys (Must return OLD values) - for i in (1..test_size).step_by(11) { - if i % 10 == 0 { - continue; - } // skip updated ones - let val = tree_update.query(&i).expect("Query failed").expect("Should find original key"); - assert_eq!(val.data, format!("Content {i}")); - } - - let size_before_compact = std::fs::metadata(&filepath)?.len(); - - // 5. GC / Compaction - tree_update.compact(&filepath)?; - - let size_after_compact = std::fs::metadata(&filepath)?.len(); - assert!(size_after_compact < size_before_compact, "Compaction should reclaimed space from COW path copies"); - - // 6. Final verification after GC - let mut final_query = BPlusTreeQuery::::try_new(&filepath)?; - for i in (0..test_size).step_by(10) { - let val = final_query.query(&i).expect("Query failed").expect("Should find updated key after GC"); - assert_eq!(val.data, format!("UpdatedContent {i}")); - } - - Ok(()) - } - - #[test] - fn query_le_cow_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("le_cow.idx"); - - // 1. Build initial tree with gaps - let mut tree = BPlusTree::new(); - for i in (0..100u32).step_by(10) { - tree.insert(i, Record { id: i, data: format!("Content {i}") }); - } - tree.store(&filepath)?; - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Initial LE check - assert_eq!(tree_update.query_le(&15).unwrap().unwrap().id, 10); - assert_eq!(tree_update.query_le(&5).unwrap().unwrap().id, 0); - - // 2. COW Update - tree_update.update(&10, Record { id: 10, data: "NewVal".to_string() }).map_err(BPlusTreeError::to_io)?; - - // 3. Verify LE returns the LATEST value - let val = tree_update.query_le(&15).expect("Query failed").expect("Should find LE key after COW update"); - assert_eq!(val.id, 10); - assert_eq!(val.data, "NewVal"); - - Ok(()) - } - - #[test] - fn disk_iterator_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("disk_it.idx"); - - let mut tree = BPlusTree::new(); - let test_size = 500u32; - for i in 0..test_size { - tree.insert(i, Record { id: i, data: format!("Value {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - - // 1. Test Iterator - let mut count = 0; - for (k, v) in query.iter() { - assert_eq!(k, count); - assert_eq!(v.data, format!("Value {count}")); - count += 1; - } - assert_eq!(count, test_size); - - // 2. Test Traverse helper - let mut traverse_count = 0; - query.traverse(|keys, values| { - for (k, v) in keys.iter().zip(values.iter()) { - assert_eq!(*k, traverse_count); - assert_eq!(v.data, format!("Value {traverse_count}")); - traverse_count += 1; - } - })?; - assert_eq!(traverse_count, test_size); - - Ok(()) - } - - #[test] - fn update_batch_basic_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_update_batch.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree - for i in 0u32..50 { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - // Test batch update - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Prepare batch updates - let updates: Vec<(u32, Record)> = (0u32..50) - .filter(|i| i % 5 == 0) - .map(|i| (i, Record { id: i, data: format!("BatchUpdated {i}") })) - .collect(); - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - tree_update.update_batch(&update_refs).map_err(BPlusTreeError::to_io)?; - drop(tree_update); - - // Verify all updates - let mut tree_query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0u32..50 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find key"); - if i % 5 == 0 { - assert_eq!(val.data, format!("BatchUpdated {i}"), "Batch updated key {i} should have new value"); - } else { - assert_eq!(val.data, format!("Initial {i}"), "Non-updated key {i} should have original value"); - } - } - - Ok(()) - } - - #[test] - fn update_batch_empty_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_update_batch_empty.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree - for i in 0u32..10 { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - let initial_root = tree_update.root_offset; - - // Test empty batch - should be no-op - let empty_batch: Vec<(&u32, &Record)> = vec![]; - let result = tree_update.update_batch(&empty_batch).map_err(BPlusTreeError::to_io)?; - - assert_eq!(result, initial_root, "Empty batch should not change root offset"); - - // Verify data unchanged - for i in 0u32..10 { - let val = tree_update.query(&i).expect("Query failed").expect("Should find key"); - assert_eq!(val.data, format!("Initial {i}")); - } - - 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 batch_updates = [(0u32, "X".repeat(6_000)), (u32::MAX, "missing".to_string())]; - let refs: Vec<(&u32, &String)> = batch_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 batch_updater = BPlusTreeUpdate::::try_new(&filepath)?; - let update_items = [(1u32, updated.clone()), (u32::MAX, "missing".to_string())]; - let refs: Vec<(&u32, &String)> = update_items.iter().map(|(k, v)| (k, v)).collect(); - - let result = batch_updater.update_batch(&refs); - assert!(matches!(result, Err(BPlusTreeError::KeyNotFound))); - drop(batch_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()?; - let filepath = tempdir.path().join("tree_update_batch_large.bin"); - let mut tree = BPlusTree::::new(); - - let test_size = 200u32; - - // Create initial tree - for i in 0..test_size { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Prepare large batch update (every other item) - let updates: Vec<(u32, Record)> = (0..test_size) - .filter(|i| i % 2 == 0) - .map(|i| (i, Record { id: i, data: format!("BatchUpdated {i}") })) - .collect(); - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - // Perform batch update - tree_update.update_batch(&update_refs).map_err(BPlusTreeError::to_io)?; - drop(tree_update); - - // Verify all updates via iterator - let reloaded_tree = BPlusTree::::load(&filepath)?; - let mut count = 0; - for (key, value) in &reloaded_tree { - if *key % 2 == 0 { - assert_eq!(value.data, format!("BatchUpdated {key}"), "Even keys should be batch updated"); - } else { - assert_eq!(value.data, format!("Initial {key}"), "Odd keys should remain unchanged"); - } - count += 1; - } - assert_eq!(count, test_size, "Should have all entries"); - - Ok(()) - } - - #[test] - fn update_batch_with_compaction_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_update_batch_compact.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree with larger data - let large_data = "x".repeat(1000); - for i in 0u32..100 { - tree.insert(i, Record { id: i, data: large_data.clone() }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Batch update with smaller data - let small_data = "y".repeat(50); - let updates: Vec<(u32, Record)> = - (0u32..100).map(|i| (i, Record { id: i, data: small_data.clone() })).collect(); - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - tree_update.update_batch(&update_refs).map_err(BPlusTreeError::to_io)?; - - let size_before_compact = std::fs::metadata(&filepath)?.len(); - - // Compact to reclaim space - tree_update.compact(&filepath)?; - - let size_after_compact = std::fs::metadata(&filepath)?.len(); - assert!(size_after_compact < size_before_compact, "Compaction should reduce file size after batch update"); - - // Verify all data is correct after compaction - drop(tree_update); - let mut tree_query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0u32..100 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find key after compaction"); - assert_eq!(val.data, small_data, "Data should be updated after compaction"); - } - - Ok(()) - } - - #[test] - fn compact_reopen_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_compact_reopen.bin"); - let mut tree = BPlusTree::::new(); - - // Initial write - tree.insert(1, Record { id: 1, data: "Initial".to_string() }); - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // 1. Write something - let r2 = Record { id: 2, data: "BeforeCompact".to_string() }; - tree_update.upsert_batch(&[(&2, &r2)])?; - - // 2. Compact (this replaces the file) - tree_update.compact(&filepath)?; - - // 3. Write something else - let r3 = Record { id: 3, data: "AfterCompact".to_string() }; - tree_update.upsert_batch(&[(&3, &r3)])?; - - drop(tree_update); - - // Verify all data is present in the NEW file - let mut tree_check = BPlusTreeQuery::::try_new(&filepath)?; - - assert!(tree_check.query(&1).map_err(BPlusTreeError::to_io)?.is_some(), "Should have key 1"); - assert!(tree_check.query(&2).map_err(BPlusTreeError::to_io)?.is_some(), "Should have key 2"); - assert!( - tree_check.query(&3).map_err(BPlusTreeError::to_io)?.is_some(), - "Should have key 3 - if missing, file handle wasn't updated" - ); - - 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)?; - - query.refresh()?; - let refreshed = query.query(&1).map_err(BPlusTreeError::to_io)?.expect("missing replaced key"); - assert_eq!(refreshed.data, "cccc"); - - Ok(()) - } - - #[cfg(unix)] - #[test] - fn query_reopens_reader_after_atomic_replace_without_mmap() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_query_reader_replace_target.bin"); - let replacement_path = tempdir.path().join("tree_query_reader_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)?; - query.mmap = None; - query.file = Some(utils::file_reader(File::open(&filepath)?)); - query.file_identity = Some(FileIdentity::from_metadata(&std::fs::metadata(&filepath)?)); - - 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)?; - - query.refresh()?; - let refreshed = query.query(&1).map_err(BPlusTreeError::to_io)?.expect("missing replaced key"); - assert_eq!(refreshed.data, "cccc"); - - Ok(()) - } - - #[test] - fn query_refresh_leaves_retry_timestamp_unchanged_on_io_error() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_query_refresh_error.bin"); - - let mut tree = BPlusTree::::new(); - tree.insert(1, Record { id: 1, data: "aaaa".to_string() }); - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - let last_refresh_at = query.last_refresh_at; - query.filepath = tempdir.path().join("missing_tree.bin"); - - let err = query.refresh().expect_err("refresh should fail when the database path is missing"); - assert_eq!(err.kind(), io::ErrorKind::NotFound); - assert_eq!(query.last_refresh_at, last_refresh_at); - - Ok(()) - } - - #[test] - fn upsert_batch_mixed_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_batch_mixed.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree with keys 0-49 - for i in 0u32..50 { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Prepare upsert batch: update existing keys 0-24, insert new keys 50-74 - let mut updates: Vec<(u32, Record)> = Vec::new(); - - // Updates to existing keys - for i in 0u32..25 { - updates.push((i, Record { id: i, data: format!("Updated {i}") })); - } - - // Inserts for new keys - for i in 50u32..75 { - updates.push((i, Record { id: i, data: format!("Inserted {i}") })); - } - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - tree_update.upsert_batch(&update_refs)?; - drop(tree_update); - - // Verify all 75 entries exist with correct values - let mut tree_query = BPlusTreeQuery::::try_new(&filepath)?; - - // Check updated keys (0-24) - for i in 0u32..25 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find updated key"); - assert_eq!(val.data, format!("Updated {i}"), "Key {i} should be updated"); - } - - // Check unchanged keys (25-49) - for i in 25u32..50 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find unchanged key"); - assert_eq!(val.data, format!("Initial {i}"), "Key {i} should remain unchanged"); - } - - // Check inserted keys (50-74) - for i in 50u32..75 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find inserted key"); - assert_eq!(val.data, format!("Inserted {i}"), "Key {i} should be inserted"); - } - - Ok(()) - } - - #[test] - fn upsert_batch_all_new_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_batch_new.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree with unrelated keys - for i in 0u32..10 { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Upsert all new keys (100-149) - let updates: Vec<(u32, Record)> = - (100u32..150).map(|i| (i, Record { id: i, data: format!("New {i}") })).collect(); - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - tree_update.upsert_batch(&update_refs)?; - drop(tree_update); - - // Verify all keys exist - let mut tree_query = BPlusTreeQuery::::try_new(&filepath)?; - - // Original keys should still exist - for i in 0u32..10 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find original key"); - assert_eq!(val.data, format!("Initial {i}")); - } - - // New keys should be inserted - for i in 100u32..150 { - let val = tree_query.query(&i).expect("Query failed").expect("Should find new key"); - assert_eq!(val.data, format!("New {i}")); - } - - Ok(()) - } - - #[test] - fn upsert_batch_all_existing_test() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_batch_existing.bin"); - let mut tree = BPlusTree::::new(); - - // Create initial tree - for i in 0u32..100 { - tree.insert(i, Record { id: i, data: format!("Initial {i}") }); - } - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Upsert all existing keys (should behave like update) - let updates: Vec<(u32, Record)> = - (0u32..100).map(|i| (i, Record { id: i, data: format!("Updated {i}") })).collect(); - - let update_refs: Vec<(&u32, &Record)> = updates.iter().map(|(k, v)| (k, v)).collect(); - - tree_update.upsert_batch(&update_refs)?; - drop(tree_update); - - // Verify all values were updated - let reloaded_tree = BPlusTree::::load(&filepath)?; - let mut count = 0; - for (key, value) in &reloaded_tree { - assert_eq!(value.data, format!("Updated {key}"), "All keys should be updated"); - count += 1; - } - assert_eq!(count, 100, "Should have exactly 100 entries"); - - Ok(()) - } - - #[test] - fn test_value_packing_efficiency() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("packing_test.bin"); - let mut tree = BPlusTree::::new(); - - // Insert 1000 small values (approx 50 bytes each) - let small_value = "x".repeat(50); - let count = 1000; - for i in 0..count { - tree.insert(i, small_value.clone()); - } - - tree.store(&filepath)?; - - let file_size = std::fs::metadata(&filepath)?.len(); - - // Expected size without packing: - // 1000 items * 4096 bytes/block = 4,096,000 bytes (~4MB) - // Plus internal nodes - let unpacked_size_estimate = u64::from(count) * u64::try_from(super::PAGE_SIZE_USIZE).unwrap(); - - println!("File size with packing: {file_size} bytes"); - println!("Estimated unpacked size: {unpacked_size_estimate} bytes"); - - // We expect significant savings. - // 1000 items * ~60 bytes / 4096 bytes/block ~= 15 blocks - // Plus tree structure overhead. - // Let's be conservative and say it should be less than 10% of unpacked size. - assert!(file_size < unpacked_size_estimate / 10, "Packing should reduce size by at least 90%"); - - Ok(()) - } - - #[test] - fn test_mixed_value_packing() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("mixed_packing.bin"); - let mut tree = BPlusTree::::new(); - - // Insert mixed values: - // 0-99: Small (50 bytes) -> Packed - // 100-109: Large (5000 bytes) -> Single (2 blocks) - // 110-209: Small (50 bytes) -> Packed - - // Insert in order - for i in 0..100 { - tree.insert(i, "s".repeat(50)); - } - for i in 100..110 { - tree.insert(i, "L".repeat(5000)); - } - for i in 110..210 { - tree.insert(i, "s".repeat(50)); - } - - tree.store(&filepath)?; - - // Verify we can read them back correctly - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - - for i in 0..100 { - let val = query.query(&i).expect("Query failed").expect("Should find small value"); - assert_eq!(val.len(), 50); - } - for i in 100..110 { - let val = query.query(&i).expect("Query failed").expect("Should find large value"); - assert_eq!(val.len(), 5000); - } - for i in 110..210 { - let val = query.query(&i).expect("Query failed").expect("Should find small value 2"); - assert_eq!(val.len(), 50); - } - - Ok(()) - } - #[test] - fn test_upsert_huge_values_chunking() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_huge.bin"); - - // Initialize tree - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Insert values > PAGE_SIZE_USIZE (4096). - // 10K value -> 3 chunks. - let val1 = "A".repeat(10000); - let val2 = "B".repeat(10000); - - let updates = [(1, val1.clone()), (2, val2.clone())]; - - let update_refs: Vec<(&u32, &String)> = updates.iter().map(|(k, v)| (k, v)).collect(); - tree_update.upsert_batch(&update_refs)?; - drop(tree_update); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.query(&1).unwrap(), Some(val1)); - assert_eq!(query.query(&2).unwrap(), Some(val2)); - - Ok(()) - } - - #[test] - fn test_upsert_deep_split() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_split.bin"); - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Insert 5000 items. - // 5000 items ensures at least Root -> Internal -> Leaf split (Height 2 or 3). - - let count = 5000; - let mut updates = Vec::with_capacity(count); - for i in 0..count { - let val = u32::try_from(i).unwrap(); - updates.push((val, val)); // value matches key - } - - // Split into batches to test multiple batch ops - for chunk in updates.chunks(1000) { - let chunk_refs: Vec<(&u32, &u32)> = chunk.iter().map(|(k, v)| (k, v)).collect(); - tree_update.upsert_batch(&chunk_refs)?; - } - drop(tree_update); - - // Validation - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..count { - let k = u32::try_from(i).unwrap(); - let val = query.query(&k).unwrap(); - assert_eq!(val, Some(k)); - } - Ok(()) - } - - #[test] - fn test_upsert_batch_overwrites() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_upsert_overwrite.bin"); - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Batch contains same key multiple times - let updates = - [(1, "First".to_string()), (1, "Second".to_string()), (2, "Two".to_string()), (1, "Third".to_string())]; - - let update_refs: Vec<(&u32, &String)> = updates.iter().map(|(k, v)| (k, v)).collect(); - tree_update.upsert_batch(&update_refs)?; - drop(tree_update); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.query(&1).unwrap(), Some("Third".to_string())); - assert_eq!(query.query(&2).unwrap(), Some("Two".to_string())); - Ok(()) - } - - #[test] - fn test_compaction_packing_limits() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_compact_pack.bin"); - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - drop(tree); - - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Insert 200 items of 100 bytes. - // Upsert creates Single blocks blocks block-aligned. - // File size > 200 * 4096 = 800KB. - - let val = "x".repeat(100); - let count = 200; - let mut updates = Vec::new(); - for i in 0..count { - updates.push((i, val.clone())); - } - // Use individual upsert calls to create fragmentation (one path write per update) - for (k, v) in updates { - let refs = [(&k, &v)]; - tree_update.upsert_batch(&refs)?; - } - - let size_before = std::fs::metadata(&filepath)?.len(); - // Each individual update writes a full path, increasing file size significantly. - assert!(size_before > u64::from(count) * 4000); - - // Now Compact - tree_update.compact(&filepath)?; - drop(tree_update); - - let size_after = std::fs::metadata(&filepath)?.len(); - // 200 items * 100 bytes = 20KB payload. - // Should pack into ~5-6 blocks (4KB each). - - println!("Size before: {size_before}, Size after: {size_after}"); - assert!(size_after < size_before / 10, "Compaction should pack values"); - assert!(size_after < 100 * 1024, "File should be small"); // < 100KB - - // Verify data - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..count { - assert_eq!(query.query(&i).unwrap(), Some(val.clone())); - } - - Ok(()) - } - - #[test] - fn test_large_keys_multiblock_node() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_multiblock.bin"); - - let mut tree = BPlusTree::::new(); - - // 5 keys of 2000 bytes each. Total ~10KB keys. - // Should span ~3 blocks (4KB each). - for i in 0..5 { - let key = format!("{:04}{}", i, "a".repeat(2000)); - tree.insert(key, i); - } - - tree.store(&filepath)?; - drop(tree); - - let loaded = BPlusTree::::load(&filepath)?; - for i in 0..5 { - let key = format!("{:04}{}", i, "a".repeat(2000)); - assert_eq!(loaded.query(&key), Some(i).as_ref()); - } - Ok(()) - } - - #[test] - fn test_upsert_multiblock_node() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_multiblock_upsert.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - drop(tree); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - - // Upsert large keys - let mut batch = Vec::new(); - let keys: Vec = (0..5).map(|i| format!("{:04}{}", i, "b".repeat(2000))).collect(); - let vals: Vec = (0..5).collect(); - - for i in 0..5 { - batch.push((&keys[i], &vals[i])); - } - - updater.upsert_batch(&batch)?; - drop(updater); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..5 { - let val = query.query(&keys[i]).unwrap(); - assert_eq!(val, Some(vals[i])); - } - Ok(()) - } - - #[test] - fn test_store_and_query_many_variable_string_keys() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_many_variable_string_keys.bin"); - - let count = 40_000usize; - let mut tree = BPlusTree::::new(); - for i in 0..count { - let key_suffix_len = 10 + (i % 97); - let key = format!("k{i:06}_{}", "x".repeat(key_suffix_len)); - tree.insert(key, format!("v{i:06}")); - } - - tree.store(&filepath)?; - drop(tree); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in (0..count).step_by(17) { - let key_suffix_len = 10 + (i % 97); - let key = format!("k{i:06}_{}", "x".repeat(key_suffix_len)); - assert_eq!(query.query(&key).map_err(super::BPlusTreeError::to_io)?, Some(format!("v{i:06}"))); - } - - Ok(()) - } - - #[test] - fn test_node_serialization_overhead() -> io::Result<()> { - use crate::{ - repository::bplustree::{ValueInfo, ValueStorageMode, PAGE_SIZE_USIZE}, - utils::binary_serialize, - }; - - // Simulate a leaf node with u32 keys and ValueInfo - let key_counts = [10, 30, 50, 80, 100]; - - for count in key_counts { - let keys: Vec = (0..count).collect(); - let value_info: Vec = (0..count) - .map(|i| ValueInfo { - mode: ValueStorageMode::Packed(u64::from(i) * 4096, (i % 16) as u16), - length: 100, - cache: Mutex::new(None), - }) - .collect(); - - let keys_serialized = binary_serialize(&keys)?; - let info_serialized = binary_serialize(&value_info)?; - - // Total content: flag(1) + keys_len(4) + keys + info_len(4) + info - let total = 1 + 4 + keys_serialized.len() + 4 + info_serialized.len(); - let fits_in_block = total <= PAGE_SIZE_USIZE; - - println!( - "Keys={}: keys_bytes={}, info_bytes={}, total={}, fits_in_block={}", - count, - keys_serialized.len(), - info_serialized.len(), - total, - fits_in_block - ); - } - Ok(()) - } - - #[test] - fn test_internal_node_size_estimate_never_underestimates_msgpack_pointers() -> io::Result<()> { - // Many children force pointer encoding overhead to dominate. - // This guards against size underestimation in internal-node layout. - let key_count = 599usize; - let mut node = BPlusTreeNode::::new(false); - node.keys = (0..key_count).map(|i| format!("k{i:04}")).collect(); - node.children = (0..=key_count).map(|_| BPlusTreeNode::::new(true)).collect(); - - let mut serial_buf = Vec::new(); - let estimated_size = node.calculate_serialized_size(&mut serial_buf)?; - - let mut file = tempfile::tempfile()?; - let mut buffer = Vec::new(); - let base = u64::from(u32::MAX) + 1_000_000; - let child_offsets: Vec = (0..node.children.len()).map(|i| base + (i as u64) * 10_000).collect(); - let actual_size = - node.serialize_internal_with_offsets(&mut file, &mut buffer, &mut serial_buf, 0, &child_offsets)?; - - assert!( - estimated_size >= actual_size, - "internal node size estimate underflowed: estimated={estimated_size}, actual={actual_size}" - ); - Ok(()) - } - - /// Test that packed value updates work correctly through the public COW path. - #[test] - fn test_packed_value_update() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_packed_update.bin"); - - // Create and store tree with small values that will be packed - let mut tree = BPlusTree::::new(); - - // Insert 50 small values (< 256 bytes) that will use packed storage - let small_val = "x".repeat(50); // 50 bytes, well under SMALL_VALUE_THRESHOLD - for i in 0..50 { - tree.insert(i, small_val.clone()); - } - - tree.store(&filepath)?; - drop(tree); - - // Get initial file size - let size_initial = std::fs::metadata(&filepath)?.len(); - - // Open for update - let mut tree_update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Test 1: Same-size update - let same_size_val = "y".repeat(50); // Same size as original - let refs1 = [(&5u32, &same_size_val)]; - tree_update.update_batch(&refs1).map_err(super::BPlusTreeError::to_io)?; - - // Test 2: Different-size update - let larger_val = "z".repeat(100); // Larger than original - let refs2 = [(&10u32, &larger_val)]; - tree_update.update_batch(&refs2).map_err(super::BPlusTreeError::to_io)?; - - // Test 3: Smaller-size update - let smaller_val = "w".repeat(30); // Smaller than original - let refs3 = [(&15u32, &smaller_val)]; - tree_update.update_batch(&refs3).map_err(super::BPlusTreeError::to_io)?; - - drop(tree_update); - - // Get file size after updates - let size_after = std::fs::metadata(&filepath)?.len(); - - // File size should have grown slightly (promoted values written at EOF) - // but not dramatically since most values are still packed - println!("Size initial: {size_initial}, Size after: {size_after}"); - assert!(size_after >= size_initial, "File should not shrink"); - - // Verify all data is correct - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - - // Updated values - assert_eq!(query.query(&5).unwrap(), Some(same_size_val.clone()), "Same-size update failed"); - assert_eq!(query.query(&10).unwrap(), Some(larger_val.clone()), "Larger-size update failed"); - assert_eq!(query.query(&15).unwrap(), Some(smaller_val.clone()), "Smaller-size update failed"); - - // Unchanged values - assert_eq!(query.query(&0).unwrap(), Some(small_val.clone()), "Unchanged value 0 incorrect"); - assert_eq!(query.query(&20).unwrap(), Some(small_val.clone()), "Unchanged value 20 incorrect"); - assert_eq!(query.query(&49).unwrap(), Some(small_val.clone()), "Unchanged value 49 incorrect"); - - Ok(()) - } - - #[test] - fn test_flush_policy() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_flush.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..10 { - tree.insert(i, format!("value_{i}")); - } - tree.store(&filepath)?; - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Test None policy - should not error - update.flush_policy = super::FlushPolicy::None; - for i in 0..5 { - update.update(&i, format!("new_{i}")).map_err(super::BPlusTreeError::to_io)?; - } - - // Verify values within same session - for i in 0..5 { - assert_eq!(update.query(&i).map_err(super::BPlusTreeError::to_io)?.unwrap(), format!("new_{i}")); - } - - // Test Batch policy - update.flush_policy = super::FlushPolicy::Batch; - let batch = [(&5u32, &"batch_5".to_string()), (&6u32, &"batch_6".to_string())]; - update.update_batch(&batch).map_err(super::BPlusTreeError::to_io)?; - - assert_eq!(update.query(&5).map_err(super::BPlusTreeError::to_io)?.unwrap(), "batch_5"); - - Ok(()) - } - - #[test] - fn serial_writer_batch_marks_dirty_and_commit_clears_dirty() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("serial_writer_dirty.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - - let writer = BPlusTreeSerialWriter::::new(&filepath, FlushPolicy::Batch)?; - let key = 7u32; - let value = "value_7".to_string(); - writer.upsert(&[(&key, &value)])?; - assert!(writer.dirty.load(Ordering::Acquire)); - - writer.commit()?; - assert!(!writer.dirty.load(Ordering::Acquire)); - - writer.shutdown()?; - Ok(()) - } - - #[test] - fn serial_writer_background_commit_requires_batch_policy() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("serial_writer_bg_policy.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - - let writer = BPlusTreeSerialWriter::::new(&filepath, FlushPolicy::Immediate)?; - let result = writer.start_background_commit(Duration::from_millis(10)); - assert!(result.is_err()); - - Ok(()) - } - - #[test] - fn serial_writer_background_commit_flushes_batch_writes() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("serial_writer_bg_flush.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - - let writer = BPlusTreeSerialWriter::::new(&filepath, FlushPolicy::Batch)?; - writer.start_background_commit(Duration::from_millis(10))?; - - let key = 3u32; - let value = "batch_value".to_string(); - writer.upsert(&[(&key, &value)])?; - assert!(writer.dirty.load(Ordering::Acquire)); - - for _ in 0..40 { - if !writer.dirty.load(Ordering::Acquire) { - break; - } - std::thread::sleep(Duration::from_millis(5)); - } - assert!(!writer.dirty.load(Ordering::Acquire)); - - writer.shutdown()?; - Ok(()) - } - - #[test] - fn test_cache_population_and_reuse() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_cache.bin"); - - // Create tree with mixed storage - let mut tree = BPlusTree::::new(); - // 1. Single storage (large-ish value) - let large_val = "A".repeat(500); - tree.insert(1, large_val.clone()); - // 2. Packed storage (small value) - let small_val = "B".repeat(20); - tree.insert(2, small_val.clone()); - - tree.store(&filepath)?; - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - - // --- Single Storage Cache --- - // First read populates cache - let val1 = update.query(&1).map_err(super::BPlusTreeError::to_io)?.unwrap(); - assert_eq!(val1, large_val); - - // Subsequent update of DIFFERENT key should not affect key 1's cache - let _ = update.update(&10, "unrelated".into()); // Might error if 10 not found, but we check cache reuse - - // Update key 1 with SAME SIZE (should use/populate cache) - let large_val_2 = "C".repeat(500); - update.update(&1, large_val_2.clone()).map_err(super::BPlusTreeError::to_io)?; - assert_eq!(update.query(&1).map_err(super::BPlusTreeError::to_io)?.unwrap(), large_val_2); - - // --- Packed Storage Cache --- - // First read of key 2 populates PackedOffset cache - let _ = update.query(&2).map_err(super::BPlusTreeError::to_io)?; - - // Update key 2 with SAME SIZE (should use Cached Offset) - let small_val_2 = "D".repeat(20); - update.update(&2, small_val_2.clone()).map_err(super::BPlusTreeError::to_io)?; - - // Verify update - assert_eq!(update.query(&2).map_err(super::BPlusTreeError::to_io)?.unwrap(), small_val_2); - - Ok(()) - } - - #[test] - fn test_concurrent_cache_access() -> io::Result<()> { - use std::{sync::Arc, thread}; - - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("tree_concurrent.bin"); - - let mut tree = BPlusTree::::new(); - let val = "concurrent_test_value".to_string(); - for i in 0..100 { - tree.insert(i, val.clone()); - } - tree.store(&filepath)?; - - // We test concurrent READS on BPlusTreeQuery. - // Note: BPlusTreeQuery::query requires &mut self for buffer management, - // so we use a Mutex to protect it. While this serializes the query() calls, - // it still tests thread-safety of the shared ValueInfo Mutexes internally - // if they were somehow shared (though they currently aren't). - let query = Arc::new(parking_lot::Mutex::new(super::BPlusTreeQuery::::try_new(&filepath)?)); - - let mut handles = Vec::new(); - for _t in 0..10 { - let q = Arc::clone(&query); - let v = val.clone(); - let handle = thread::spawn(move || { - for i in 0..100 { - let mut guard = q.lock(); - let res = guard.query(&i).expect("Query failed"); - assert_eq!(res, Some(v.clone())); - } - }); - handles.push(handle); - } - - for h in handles { - h.join().expect("Thread panicked"); - } - - Ok(()) - } - - #[test] - fn test_page_initialization() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - assert_eq!(page.header.page_type, PageType::Leaf); - assert_eq!(page.header.cell_count, 0); - assert_eq!(page.header.free_start, PAGE_HEADER_SIZE); - assert_eq!(page.header.free_end, PAGE_SIZE); - assert_eq!(page.free_space(), PAGE_SIZE_USIZE - PAGE_HEADER_SIZE_USIZE); - } - - #[test] - fn test_insert_get() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let mut page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - - let val1 = b"hello"; - let val2 = b"world"; - - // Insert length-prefixed for test realism - let mut cell1 = Vec::new(); - cell1.extend_from_slice(&u32::try_from(val1.len()).unwrap().to_le_bytes()); - cell1.extend_from_slice(val1); - - let mut cell2 = Vec::new(); - cell2.extend_from_slice(&u32::try_from(val2.len()).unwrap().to_le_bytes()); - cell2.extend_from_slice(val2); - - page.insert_at_index(0, &cell1).unwrap(); - page.insert_at_index(1, &cell2).unwrap(); - - assert_eq!(page.header.cell_count, 2); - - let read1 = page.get_cell(0).expect("Get cell 0"); - assert_eq!(&read1[4..], val1); - - let read2 = page.get_cell(1).expect("Get cell 1"); - assert_eq!(&read2[4..], val2); - } - - #[test] - fn test_split_off() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let mut page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - - let payload = vec![0xAAu8; 500]; - let mut cell = Vec::new(); - cell.extend_from_slice(&u32::try_from(payload.len()).unwrap().to_le_bytes()); - cell.extend_from_slice(&payload); - - for i in 0..6 { - page.insert_at_index(i, &cell).unwrap(); - } - - assert_eq!(page.header.cell_count, 6); - - let new_page_bytes = page.split_off().expect("Split failed").expect("Should have split"); - - // Check original page - assert_eq!(page.header.cell_count, 3); - - // Check new page - let header = PageHeader::deserialize(&new_page_bytes[..PAGE_HEADER_SIZE_USIZE]).expect("Deserialize failed"); - assert_eq!(header.cell_count, 3); - } - - #[test] - fn test_split_off_edge_cases() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let mut page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - - // Case 0: Split empty page -> Should Error - let res = page.split_off(); - assert!(matches!(res, Err(PageError::InvalidIndex))); - - // Case 1: Split single item page -> Should return None (no-op) - let val = b"item"; - let mut cell = Vec::new(); - cell.extend_from_slice(&u32::try_from(val.len()).unwrap().to_le_bytes()); - cell.extend_from_slice(val); - page.insert_at_index(0, &cell).unwrap(); - - let res = page.split_off(); - match res { - Ok(None) => { - assert_eq!(page.header.cell_count, 1); // Original page untouched - } - Ok(Some(_)) => panic!("Split of single item should result in None"), - Err(e) => panic!("Split of single item should result in no-op, not error: {e:?}"), - } - } - - #[test] - fn internal_split_preserves_child_key_invariants() { - let mut node = BPlusTreeNode::::new(false); - node.keys = (1..=65).collect(); - node.children = (0..=65) - .map(|key| { - let mut child = BPlusTreeNode::new(true); - child.keys.push(key); - child.values.push(key); - child - }) - .collect(); - - let right = node.split(64); - - assert_eq!(node.children.len(), node.keys.len() + 1); - assert_eq!(right.children.len(), right.keys.len() + 1); - assert_eq!(BPlusTreeNode::::find_leaf_entry(&right), Some(&33)); - } - - #[test] - fn file_deserializer_rejects_unknown_node_flag() { - let mut bytes = Vec::new(); - bytes.push(3u8); - bytes.extend_from_slice(&1u32.to_le_bytes()); - bytes.push(0x90); - bytes.extend_from_slice(&1u32.to_le_bytes()); - bytes.push(0x90); - - let mut cursor = io::Cursor::new(bytes); - let mut buffer = Vec::new(); - let result = BPlusTreeNode::::deserialize_from_block(&mut cursor, &mut buffer, 0, false); - - assert!(matches!(result, Err(err) if err.kind() == io::ErrorKind::InvalidData)); - } - - #[test] - fn slice_deserializer_rejects_empty_slice() { - let mut cursor = io::Cursor::new(Vec::::new()); - let result = BPlusTreeNode::::deserialize_from_block_slice(&[], None, &mut cursor, false); - - assert!(matches!(result, Err(err) if err.kind() == io::ErrorKind::UnexpectedEof)); - } - - #[test] - fn slotted_page_rejects_cell_offsets_inside_header_or_slot_area() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let mut page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - page.header.cell_count = 1; - page.header.free_start = PAGE_HEADER_SIZE + u16::try_from(SLOT_SIZE).expect("slot size fits"); - page.commit(); - - let invalid_offset = PAGE_HEADER_SIZE + u16::try_from(SLOT_SIZE).expect("slot size fits"); - page.data[PAGE_HEADER_SIZE_USIZE..PAGE_HEADER_SIZE_USIZE + SLOT_SIZE] - .copy_from_slice(&invalid_offset.to_le_bytes()); - - assert!(page.get_cell(0).is_none()); - } - - #[test] - fn msgpack_u32_scan_handles_non_negative_signed_encodings() { - assert_eq!( - ::compare_at_position(&7, &[0xd2, 0, 0, 0, 5]), - Some((std::cmp::Ordering::Greater, 5)) - ); - assert_eq!( - ::compare_at_position(&7, &[0xd1, 0, 7]), - Some((std::cmp::Ordering::Equal, 3)) - ); - assert_eq!( - ::compare_at_position(&7, &[0xd0, 9]), - Some((std::cmp::Ordering::Less, 2)) - ); - } - - #[test] - fn small_tree_test() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("small_tree.bin"); - - let mut tree = BPlusTree::::new(); - tree.insert("key1".to_string(), "val1".to_string()); - tree.insert("key2".to_string(), "val2".to_string()); - - assert_eq!(tree.len(), 2); - - tree.store(&filepath)?; - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - assert_eq!(update.len().unwrap(), 2); - - let res1 = update.query(&"key1".to_string()).unwrap(); - assert_eq!(res1, Some("val1".to_string())); - - let res2 = update.query(&"key2".to_string()).unwrap(); - assert_eq!(res2, Some("val2".to_string())); - - Ok(()) - } - - #[test] - fn test_metadata() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("metadata_test.bin"); - - // 1. Test in-memory BPlusTree metadata - let mut tree = BPlusTree::::new(); - assert!(matches!(tree.get_metadata(), BPlusTreeMetadata::Empty)); - - let meta = BPlusTreeMetadata::TargetIdMapping(12345); - tree.set_metadata(meta.clone()); - assert_eq!(tree.get_metadata(), &meta); - - // 2. Persist and check reloaded metadata - tree.store(&filepath)?; - let loaded = BPlusTree::::load(&filepath)?; - assert_eq!(loaded.get_metadata(), &meta); - drop(loaded); - - // 3. Test BPlusTreeUpdate metadata - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - assert_eq!(update.get_metadata()?, meta); - - let new_meta = BPlusTreeMetadata::TargetIdMapping(67890); - update.set_metadata(&new_meta)?; - assert_eq!(update.get_metadata()?, new_meta); - drop(update); - - // Reload and verify again - let loaded2 = BPlusTree::::load(&filepath)?; - assert_eq!(loaded2.get_metadata(), &new_meta); - - Ok(()) - } - - #[test] - fn test_query_zero_copy() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("zero_copy_test.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..100 { - tree.insert(i, format!("value_{i}")); - } - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..100 { - let res = query.query_zero_copy(&i).expect("Zero copy query failed"); - assert_eq!(res, Some(format!("value_{i}"))); - } - - // Test key not found - assert_eq!(query.query_zero_copy(&101).unwrap(), None); - - 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. - query.refresh()?; - 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 query_refresh_throttle_keeps_hot_queries_on_cached_root_until_forced_refresh() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("query_refresh_throttle.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)?; - query.set_refresh_interval(Duration::from_secs(60)); - assert_eq!(query.query(&999).map_err(BPlusTreeError::to_io)?, None); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let inserted = "throttled_value".to_string(); - updater.upsert_batch(&[(&999u32, &inserted)])?; - drop(updater); - - assert_eq!(query.query(&999).map_err(BPlusTreeError::to_io)?, None); - - query.refresh()?; - assert_eq!(query.query(&999).map_err(BPlusTreeError::to_io)?, Some(inserted)); - - Ok(()) - } - - #[test] - fn query_cache_keeps_internal_nodes_separate_from_leaf_lru() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("query_split_cache.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..5000 { - tree.insert(i, format!("value_{i}")); - } - tree.store(&filepath)?; - drop(tree); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..5000 { - assert_eq!(query.query(&i).map_err(BPlusTreeError::to_io)?, Some(format!("value_{i}"))); - } - - assert!(!query.node_cache.internal.is_empty(), "internal nodes should use the sticky internal cache"); - assert!(!query.node_cache.leaf.is_empty(), "leaf nodes should use the leaf LRU cache"); - - Ok(()) - } - - #[test] - fn test_query_le_query_struct() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("query_le_test.bin"); - - let mut tree = BPlusTree::::new(); - for i in (0..100).step_by(10) { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - - // Exact match - assert_eq!(query.query_le(&20).unwrap().unwrap(), "val_20"); - - // In gap - assert_eq!(query.query_le(&25).unwrap().unwrap(), "val_20"); - - // Before all - assert_eq!(query.query_le(&0).unwrap().unwrap(), "val_0"); - - // After all - assert_eq!(query.query_le(&1000).unwrap().unwrap(), "val_90"); - - Ok(()) - } - - #[test] - fn test_range_page_reports_has_more_and_returns_exact_slice() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("range_page_test.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..10u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - let (page, has_more) = query - .range_page(Bound::Included(&0), Bound::Included(&9), 3, 3) - .map_err(BPlusTreeError::to_io)?; - - assert_eq!( - page, - vec![ - (3, String::from("val_3")), - (4, String::from("val_4")), - (5, String::from("val_5")), - ] - ); - assert!(has_more); - - let (tail, tail_has_more) = query - .range_page(Bound::Included(&0), Bound::Included(&9), 9, 3) - .map_err(BPlusTreeError::to_io)?; - assert_eq!(tail, vec![(9, String::from("val_9"))]); - assert!(!tail_has_more); - - Ok(()) - } - - #[test] - fn test_range_iter_seeks_into_following_leaf_when_start_key_is_past_first_leaf() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("range_iter_seek_test.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..512u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - let rows: Vec<(u32, String)> = query - .range_iter(Bound::Included(&255), Bound::Included(&260)) - .collect::, _>>() - .map_err(BPlusTreeError::to_io)?; - - assert_eq!( - rows, - vec![ - (255, String::from("val_255")), - (256, String::from("val_256")), - (257, String::from("val_257")), - (258, String::from("val_258")), - (259, String::from("val_259")), - (260, String::from("val_260")), - ] - ); - - Ok(()) - } - - #[test] - fn test_delete_batch_hides_entries_from_query_iter_and_len() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("delete_tombstone_visibility.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..240u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - drop(tree); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let deleted_keys: Vec = (0..240u32).filter(|k| k % 3 == 0).collect(); - let delete_refs: Vec<&u32> = deleted_keys.iter().collect(); - let deleted = updater.delete_batch(&delete_refs)?; - assert_eq!(deleted, deleted_keys.len()); - assert_eq!(updater.len().map_err(BPlusTreeError::to_io)?, 240usize - deleted_keys.len()); - assert_eq!(updater.query(&0).map_err(BPlusTreeError::to_io)?, None); - assert_eq!(updater.query(&1).map_err(BPlusTreeError::to_io)?, Some("val_1".to_string())); - drop(updater); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.len().map_err(BPlusTreeError::to_io)?, 240usize - deleted_keys.len()); - - let iter_keys: Vec = query.iter().map(|(k, _)| k).collect(); - assert_eq!(iter_keys.len(), 240usize - deleted_keys.len()); - for key in iter_keys { - assert_ne!(key % 3, 0); - } - - Ok(()) - } - - #[test] - fn test_sorted_index_iterator_skips_tombstoned_entries_without_rebuild() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("sorted_index_tombstone_skip.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..120u32 { - tree.insert(i, Record { id: i, data: format!("val_{i}") }); - } - tree.store_with_index(&filepath, |record| record.id)?; - drop(tree); - - let deleted_keys: Vec = (0..120u32).filter(|key| key % 4 == 0).collect(); - let delete_refs: Vec<&u32> = deleted_keys.iter().collect(); - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let deleted = updater.delete_batch(&delete_refs)?; - assert_eq!(deleted, deleted_keys.len()); - drop(updater); - - let query = BPlusTreeQuery::::try_new(&filepath)?; - let sorted_iter = query.disk_iter_sorted::()?; - let sorted_keys: Vec = - sorted_iter.map(|entry| entry.map(|(key, _)| key)).collect::>>()?; - - assert_eq!(sorted_keys.len(), 120usize - deleted_keys.len()); - for key in &sorted_keys { - assert_ne!(key % 4, 0, "tombstoned key must be skipped in sorted iterator"); - } - - Ok(()) - } - - #[test] - fn test_tombstone_header_flag_transitions_on_delete_and_compact() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("tombstone_header_flag.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..64u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - drop(tree); - - let query_before = BPlusTreeQuery::::try_new(&filepath)?; - assert!(!query_before.has_tombstones()); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let delete_key = 3u32; - let deleted = updater.delete(&delete_key)?; - assert!(deleted); - drop(updater); - - let query_after_delete = BPlusTreeQuery::::try_new(&filepath)?; - assert!(query_after_delete.has_tombstones()); - - let mut updater_for_compact = BPlusTreeUpdate::::try_new(&filepath)?; - updater_for_compact.compact(&filepath)?; - drop(updater_for_compact); - - let query_after_compact = BPlusTreeQuery::::try_new(&filepath)?; - assert!(!query_after_compact.has_tombstones()); - - Ok(()) - } - - #[test] - fn test_query_le_skips_tombstones_across_leaves() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("query_le_tombstone.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..640u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - drop(tree); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let deleted_keys: Vec = (0..640u32).filter(|k| k % 5 == 0).collect(); - let delete_refs: Vec<&u32> = deleted_keys.iter().collect(); - let deleted = updater.delete_batch(&delete_refs)?; - assert_eq!(deleted, deleted_keys.len()); - drop(updater); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for q in 0..640u32 { - let mut expected_key: Option = None; - for candidate in (0..=q).rev() { - if candidate % 5 != 0 { - expected_key = Some(candidate); - break; - } - } - let expected_value = expected_key.map(|candidate| format!("val_{candidate}")); - let actual = query.query_le(&q).map_err(BPlusTreeError::to_io)?; - assert_eq!(actual, expected_value, "query_le mismatch for key {q}"); - } - - assert_eq!(query.query_le(&639).map_err(BPlusTreeError::to_io)?, Some("val_639".to_string())); - assert_eq!(query.query_le(&640).map_err(BPlusTreeError::to_io)?, Some("val_639".to_string())); - assert_eq!(query.query_le(&0).map_err(BPlusTreeError::to_io)?, None); - - Ok(()) - } - - #[test] - fn test_load_skips_tombstones_without_error() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("load_tombstones.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..512u32 { - tree.insert(i, format!("val_{i}")); - } - tree.store(&filepath)?; - drop(tree); - - let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; - let deleted_keys: Vec = (0..512u32).filter(|k| k % 4 == 0).collect(); - let delete_refs: Vec<&u32> = deleted_keys.iter().collect(); - let deleted = updater.delete_batch(&delete_refs)?; - assert_eq!(deleted, deleted_keys.len()); - drop(updater); - - let loaded = BPlusTree::::load(&filepath)?; - let expected_live = 512usize - deleted_keys.len(); - assert_eq!(loaded.len(), expected_live); - - assert!(loaded.query(&0).is_none()); - assert_eq!(loaded.query(&1).cloned(), Some("val_1".to_string())); - - let le_key = loaded.find_le(&4).map(|(k, _)| *k); - assert_eq!(le_key, Some(3)); - - let mut seen = 0usize; - for (key, value) in &loaded { - assert_ne!(*key % 4, 0); - assert_eq!(value, &format!("val_{key}")); - seen += 1; - } - assert_eq!(seen, expected_live); - - Ok(()) - } - - #[test] - fn test_load_non_existent_errors() { - let tempdir = tempdir().unwrap(); - let filepath = tempdir.path().join("non_existent_file.bin"); - - // Load non-existent - let res_load = BPlusTree::::load(&filepath); - assert!(res_load.is_err()); - - // Update try_new non-existent - let res_update = BPlusTreeUpdate::::try_new(&filepath); - assert!(res_update.is_err()); - - // Query try_new non-existent - let res_query = BPlusTreeQuery::::try_new(&filepath); - assert!(res_query.is_err()); - } - - #[test] - fn test_update_key_not_found_error() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("not_found_err.bin"); - - let mut tree = BPlusTree::::new(); - tree.insert(1, "one".into()); - tree.store(&filepath)?; - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - let res = update.update(&2, "two".into()); - assert!(matches!(res, Err(BPlusTreeError::KeyNotFound))); - - Ok(()) - } - - #[test] - fn test_empty_tree_operations() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("empty_tree_ops.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert!(query.is_empty().unwrap()); - assert_eq!(query.len().unwrap(), 0); - assert_eq!(query.query(&1).unwrap(), None); - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - assert!(update.is_empty().unwrap()); - assert_eq!(update.len().unwrap(), 0); - - Ok(()) - } - - #[test] - fn test_query_zero_copy_string() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("zero_copy_string.bin"); - - let mut tree = BPlusTree::::new(); - for i in 0..50 { - tree.insert(format!("key_{i:03}"), i); - } - tree.store(&filepath)?; - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - for i in 0..50 { - let k = format!("key_{i:03}"); - let res = query.query_zero_copy(&k).expect("Zero copy query failed"); - assert_eq!(res, Some(i)); - } - - assert_eq!(query.query_zero_copy(&"key_051".to_string()).unwrap(), None); - - Ok(()) - } - - #[test] - fn test_slotted_page_compact_manual() { - let mut data = [0u8; PAGE_SIZE_USIZE]; - let mut page = SlottedPage::new(&mut data, PageType::Leaf).expect("Init failed"); - - let cell1 = vec![0x01; 100]; - let cell2 = vec![0x02; 100]; - let cell3 = vec![0x03; 100]; - - let mut c1 = Vec::new(); - c1.extend_from_slice(&100u32.to_le_bytes()); - c1.extend_from_slice(&cell1); - - let mut c2 = Vec::new(); - c2.extend_from_slice(&100u32.to_le_bytes()); - c2.extend_from_slice(&cell2); - - let mut c3 = Vec::new(); - c3.extend_from_slice(&100u32.to_le_bytes()); - c3.extend_from_slice(&cell3); - - page.insert_at_index(0, &c1).unwrap(); - page.insert_at_index(1, &c2).unwrap(); - page.insert_at_index(2, &c3).unwrap(); - - // Compacting a non-fragmented page should be fine - page.compact().expect("Compact failed"); - assert_eq!(page.header.cell_count, 3); - assert_eq!(&page.get_cell(0).unwrap()[4..], &cell1); - assert_eq!(&page.get_cell(1).unwrap()[4..], &cell2); - assert_eq!(&page.get_cell(2).unwrap()[4..], &cell3); - } - - #[test] - fn test_upsert_batch_preserialized() -> io::Result<()> { - let tempdir = tempdir()?; - let filepath = tempdir.path().join("preserialized_test.bin"); - - let mut tree = BPlusTree::::new(); - tree.store(&filepath)?; - - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - - // Manually serialize records - let r1 = Record { id: 1, data: "preserialized_1".to_string() }; - let r2 = Record { id: 2, data: "preserialized_2".to_string() }; - - let r1_bytes = binary_serialize(&r1)?; - let r2_bytes = binary_serialize(&r2)?; - - update.upsert_batch_preserialized(vec![(1, r1_bytes), (2, r2_bytes)])?; - - // Verify with query - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.query(&1).unwrap(), Some(r1)); - assert_eq!(query.query(&2).unwrap(), Some(r2)); - - // Test mixed: existing and new - let r1_new = Record { id: 1, data: "updated_preserialized_1".to_string() }; - let r3 = Record { id: 3, data: "new_preserialized_3".to_string() }; - - let r1_new_bytes = binary_serialize(&r1_new)?; - let r3_bytes = binary_serialize(&r3)?; - - update.upsert_batch_preserialized(vec![(1, r1_new_bytes), (3, r3_bytes)])?; - - let mut query2 = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query2.query(&1).unwrap(), Some(r1_new)); - assert_eq!(query2.query(&2).unwrap(), Some(Record { id: 2, data: "preserialized_2".to_string() })); - assert_eq!(query2.query(&3).unwrap(), Some(r3)); - - Ok(()) - } - - #[test] - fn update_single_key_uses_cow_for_equal_size_values() -> io::Result<()> { - let tempdir = tempfile::tempdir()?; - let filepath = tempdir.path().join("update_single_key_cow.bin"); - - let mut tree = BPlusTree::::new(); - tree.insert(1, "aaaa".to_string()); - tree.store(&filepath)?; - - let original_len = std::fs::metadata(&filepath)?.len(); - let mut update = BPlusTreeUpdate::::try_new(&filepath)?; - let original_root = update.root_offset; - - let new_root = update.update(&1, "bbbb".to_string()).map_err(BPlusTreeError::to_io)?; - drop(update); - - let updated_len = std::fs::metadata(&filepath)?.len(); - assert_ne!(new_root, original_root, "single-key update must use v2 COW semantics"); - assert!(updated_len > original_len, "COW update must append replacement nodes/value data"); - - let mut query = BPlusTreeQuery::::try_new(&filepath)?; - assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some("bbbb".to_string())); - Ok(()) - } - - #[test] - fn default_orders_are_conservative_and_key_type_independent() { - let u32_tree = BPlusTree::::new(); - let string_tree = BPlusTree::::new(); - let arc_str_tree = BPlusTree::, u32>::new(); - let uuid_tree = BPlusTree::::new(); - - assert_eq!(u32_tree.inner_order, string_tree.inner_order); - assert_eq!(u32_tree.leaf_order, string_tree.leaf_order); - assert_eq!(u32_tree.inner_order, arc_str_tree.inner_order); - assert_eq!(u32_tree.leaf_order, arc_str_tree.leaf_order); - assert_eq!(u32_tree.inner_order, uuid_tree.inner_order); - assert_eq!(u32_tree.leaf_order, uuid_tree.leaf_order); - assert!(u32_tree.inner_order >= 2); - assert!(u32_tree.leaf_order >= 2); - } - - #[test] - fn explicit_orders_override_v2_defaults() { - let tree = BPlusTree::::new_with_orders(5, 7); - - assert_eq!(tree.inner_order, 5); - assert_eq!(tree.leaf_order, 7); - } -} diff --git a/backend/src/repository/bplustree/common.rs b/backend/src/repository/bplustree/common.rs new file mode 100644 index 000000000..c30f4945f --- /dev/null +++ b/backend/src/repository/bplustree/common.rs @@ -0,0 +1,103 @@ +use log::warn; +use memmap2::{Advice, Mmap}; +use std::{ + ffi::OsString, + fs::File, + io, + path::{Path, PathBuf}, +}; + +#[cfg(not(unix))] +use std::io::{Read, Seek, SeekFrom}; +#[cfg(unix)] +use std::os::unix::fs::FileExt; + +fn advise_mmap(mmap: &Mmap, advice: Advice, context: &str) { + if let Err(err) = mmap.advise(advice) { + warn!("Failed to apply mmap advice {advice:?} for {context}: {err}"); + } +} + +pub(crate) fn mmap_with_advice(file: &File, advice: Advice, context: &str) -> Option { + // SAFETY: Every v3 persisted query holds its shared sidecar lock for the mapping lifetime, and v3 writers require + // the exclusive lock. The v3 temporary verifier maps a private synchronized file that is not mutated while mapped. + // Legacy v2 callers retain their existing invariant that a mapped file is never truncated in place. + let mmap = unsafe { + match Mmap::map(file) { + Ok(mmap) => mmap, + Err(err) => { + warn!("Failed to mmap B+Tree for {context}; falling back to buffered file I/O: {err}"); + return None; + } + } + }; + advise_mmap(&mmap, advice, context); + Some(mmap) +} + +pub(crate) fn read_exact_at_offset(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { + #[cfg(unix)] + file.read_exact_at(buf, offset)?; + #[cfg(not(unix))] + { + let mut file = file; + let current_pos = file.stream_position()?; + file.seek(SeekFrom::Start(offset))?; + file.read_exact(buf)?; + file.seek(SeekFrom::Start(current_pos))?; + } + Ok(()) +} + +pub(crate) fn sidecar_lock_path(filepath: &Path) -> PathBuf { + if let Some(stem) = filepath.file_stem() { + let mut name = OsString::from("."); + name.push(stem); + name.push(".lock"); + filepath.with_file_name(name) + } else { + filepath.with_extension("lock") + } +} + +#[derive(Debug)] +pub enum BPlusTreeError { + Io(io::Error), + Corrupted(String), + InvalidStructure(String), + KeyNotFound, +} + +impl std::fmt::Display for BPlusTreeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Io(err) => write!(f, "I/O error: {err}"), + Self::Corrupted(msg) => write!(f, "Data corrupted: {msg}"), + Self::InvalidStructure(msg) => write!(f, "Invalid structure: {msg}"), + Self::KeyNotFound => write!(f, "Key not found"), + } + } +} + +impl std::error::Error for BPlusTreeError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Io(err) => Some(err), + Self::Corrupted(_) | Self::InvalidStructure(_) | Self::KeyNotFound => None, + } + } +} + +impl From for BPlusTreeError { + fn from(err: io::Error) -> Self { Self::Io(err) } +} + +impl BPlusTreeError { + pub fn to_io(self) -> io::Error { + match self { + Self::Io(error) => error, + Self::KeyNotFound => io::Error::new(io::ErrorKind::NotFound, "Key not found"), + error => io::Error::new(io::ErrorKind::InvalidData, error), + } + } +} diff --git a/backend/src/repository/bplustree_migration.rs b/backend/src/repository/bplustree/migration.rs similarity index 72% rename from backend/src/repository/bplustree_migration.rs rename to backend/src/repository/bplustree/migration.rs index 83a52dd0c..95b011c28 100644 --- a/backend/src/repository/bplustree_migration.rs +++ b/backend/src/repository/bplustree/migration.rs @@ -1,32 +1,33 @@ //! Startup migration helpers for legacy repository files. //! -//! The current B+Tree implementation writes storage format v2. This migrator is -//! intentionally limited to normalizing legacy v1 headers into v2-compatible -//! files at startup. It is format-version aware, but it is not a generic typed -//! v2 -> v3 rewrite tool; that migration must be added at repository call sites -//! that know the concrete key and value types. +//! The current B+Tree implementation writes storage format v3. This migrator +//! recognizes the repository's closed set of concrete key/value schemas and +//! rewrites legacy v1/v2 databases into verified v3 files before publishing +//! them atomically. -use super::{ - bplustree::{BPlusTree, MAGIC, STORAGE_VERSION}, +use super::v3::{BPlusTree, MAGIC, STORAGE_VERSION}; +use crate::repository::{ + qos_snapshot_repository::{QosAggregationCheckpoint, QosSnapshotRecord}, storage_const, + target_id_mapping::VirtualIdRecord, }; +use crate::api::model::{MetadataRetryDbKey, MetadataRetryDbValue}; use fs2::FileExt as _; use log::{info, trace, warn}; -use shared::model::{ClusterFlags, ConfigPaths, NetworkAccessDto, ProxyType, ProxyUserStatus}; +use shared::model::{ + ClusterFlags, ConfigPaths, EpgChannel, M3uPlaylistItem, NetworkAccessDto, ProxyType, ProxyUserStatus, UUIDType, + XtreamPlaylistItem, +}; use std::{ collections::{HashSet, VecDeque}, ffi::OsStr, fs::OpenOptions, io::{self, Read, Seek, SeekFrom, Write}, path::{Path, PathBuf}, + sync::Arc, }; const LEGACY_STORAGE_VERSION: u32 = 1; -const METADATA_LEN_OFFSET: u64 = 16; -const METADATA_MAX_SIZE: u32 = 4000; -const HEADER_FLAG_HAS_METADATA_FLAGS: u32 = 1 << 31; -const HEADER_FLAG_HAS_TOMBSTONES: u32 = 1 << 30; -const HEADER_METADATA_LEN_MASK: u32 = !(HEADER_FLAG_HAS_METADATA_FLAGS | HEADER_FLAG_HAS_TOMBSTONES); const MARKER_FILE_GUARD_PREFIX: &str = ".db_mergeto_v"; const MARKER_FILE_GUARD_PREFIX_LEGACY_ALT: &str = ".db_mergedto"; const MARKER_FILE_API_USER_GUARD: &str = ".userdb_mergeto_v6"; @@ -38,7 +39,7 @@ pub struct BPlusTreeMigrationStats { pub scanned_files: usize, pub bplustree_files: usize, pub migrated_files: usize, - pub skipped_by_marker: bool, + pub marker_was_current: bool, } #[derive(Debug)] @@ -60,17 +61,13 @@ impl BPlusTreeStartupMigrator { let resolved_roots = Self::resolve_scan_roots(&self.roots); let roots_fingerprint = Self::roots_fingerprint(&resolved_roots); if let Some(marker_path) = &self.migration_marker_path { - if Self::marker_matches(marker_path, &roots_fingerprint)? { - stats.skipped_by_marker = true; - return Ok(stats); - } + stats.marker_was_current = Self::marker_matches(marker_path, &roots_fingerprint)?; } - for root in &resolved_roots { let files = Self::collect_db_files_for_root(root)?; for file in files { stats.scanned_files = stats.scanned_files.saturating_add(1); - match Self::migrate_file_if_needed(&file)? { + match Self::migrate_file_if_needed(&file, &resolved_roots)? { FileMigrationOutcome::NotBPlusTree => {} FileMigrationOutcome::AlreadyCurrent | FileMigrationOutcome::Locked => { stats.bplustree_files = stats.bplustree_files.saturating_add(1); @@ -220,6 +217,10 @@ impl BPlusTreeStartupMigrator { if !file_type.is_file() { continue; } + if Self::is_abandoned_v3_temporary(&path) { + std::fs::remove_file(&path)?; + continue; + } if path.extension().and_then(OsStr::to_str).is_some_and(|ext| ext.eq_ignore_ascii_case("db")) { files.push(path); } @@ -229,6 +230,22 @@ impl BPlusTreeStartupMigrator { Ok(files) } + fn is_abandoned_v3_temporary(path: &Path) -> bool { + let Some(name) = path.file_name().and_then(OsStr::to_str) else { + return false; + }; + let Some(prefix) = name.strip_suffix(".v3.tmp") else { + return false; + }; + let Some((database_name, transaction)) = prefix.rsplit_once('.') else { + return false; + }; + Path::new(database_name) + .extension() + .is_some_and(|extension| extension.eq_ignore_ascii_case("db") || extension.eq_ignore_ascii_case("idx")) + && uuid::Uuid::parse_str(transaction).is_ok() + } + fn marker_matches(marker_path: &Path, expected_fingerprint: &str) -> io::Result { let Some(stored_fingerprint) = Self::read_migration_marker_fingerprint(marker_path)? else { return Ok(false); @@ -281,7 +298,7 @@ impl BPlusTreeStartupMigrator { Ok(()) } - fn migrate_file_if_needed(path: &Path) -> io::Result { + fn migrate_file_if_needed(path: &Path, roots: &[PathBuf]) -> io::Result { let mut read_file = OpenOptions::new().read(true).open(path)?; let file_len = read_file.metadata()?.len(); if file_len < 8 { @@ -297,9 +314,10 @@ impl BPlusTreeStartupMigrator { let version = u32::from_le_bytes(header[4..8].try_into().map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?); if version == STORAGE_VERSION { + let _ = super::v3::BPlusTreeQuery::::try_new(path)?; return Ok(FileMigrationOutcome::AlreadyCurrent); } - if version != LEGACY_STORAGE_VERSION { + if version != LEGACY_STORAGE_VERSION && version != 2 { return Err(io::Error::new( io::ErrorKind::InvalidData, format!( @@ -332,7 +350,7 @@ impl BPlusTreeStartupMigrator { if locked_version == STORAGE_VERSION { return Ok(FileMigrationOutcome::AlreadyCurrent); } - if locked_version != LEGACY_STORAGE_VERSION { + if locked_version != LEGACY_STORAGE_VERSION && locked_version != 2 { return Err(io::Error::new( io::ErrorKind::InvalidData, format!( @@ -342,38 +360,93 @@ impl BPlusTreeStartupMigrator { )); } - let _flags_written = Self::normalize_metadata_flags(&mut file)?; - file.sync_data()?; - - file.seek(SeekFrom::Start(4))?; - file.write_all(&STORAGE_VERSION.to_le_bytes())?; - file.flush()?; - file.sync_data()?; + Self::migrate_typed_path(path, roots, locked_version).map_err(|error| { + io::Error::new(error.kind(), format!("Failed to migrate B+Tree {}: {error}", path.display())) + })?; Ok(FileMigrationOutcome::Migrated) } - fn normalize_metadata_flags(file: &mut std::fs::File) -> io::Result { - file.seek(SeekFrom::Start(METADATA_LEN_OFFSET))?; - let mut metadata_len_raw = [0u8; 4]; - file.read_exact(&mut metadata_len_raw)?; - let raw = u32::from_le_bytes(metadata_len_raw); + fn migrate_typed_path(path: &Path, roots: &[PathBuf], version: u32) -> io::Result<()> { + use super::v3::migration::{migrate_v2_typed, migrate_v2_typed_with_index}; - let metadata_len = raw & HEADER_METADATA_LEN_MASK; - if metadata_len > METADATA_MAX_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); - } - - let mut normalized = metadata_len | HEADER_FLAG_HAS_METADATA_FLAGS; - normalized &= !HEADER_FLAG_HAS_TOMBSTONES; - - if normalized == raw { - return Ok(false); - } - - file.seek(SeekFrom::Start(METADATA_LEN_OFFSET))?; - file.write_all(&normalized.to_le_bytes())?; - Ok(true) + let kind = Self::migration_kind(path, roots).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported typed B+Tree migration for {} (storage version {version})", path.display()), + ) + })?; + match kind { + TypedMigration::MetadataRetry => migrate_v2_typed::(path)?, + TypedMigration::QosSnapshot => migrate_v2_typed::(path)?, + TypedMigration::QosCheckpoint => migrate_v2_typed::(path)?, + TypedMigration::GeoIp => migrate_v2_typed::(path)?, + TypedMigration::IdMapping => migrate_v2_typed::(path)?, + TypedMigration::UuidMapping => migrate_v2_typed::(path)?, + TypedMigration::TargetM3u => { + migrate_v2_typed_with_index::(path, |item| item.source_ordinal)? + } + TypedMigration::InputM3u => migrate_v2_typed::, M3uPlaylistItem>(path)?, + TypedMigration::Library => migrate_v2_typed::(path)?, + TypedMigration::TargetXtream => { + migrate_v2_typed_with_index::(path, |item| item.source_ordinal)? + } + TypedMigration::InputXtream => migrate_v2_typed::(path)?, + TypedMigration::Epg => migrate_v2_typed::, EpgChannel>(path)?, + }; + Ok(()) } + + fn migration_kind(path: &Path, roots: &[PathBuf]) -> Option { + let relative = roots.iter().find_map(|root| path.strip_prefix(root).ok())?; + let components = relative.iter().filter_map(OsStr::to_str).collect::>(); + match components.as_slice() { + ["metadata_retry_state.db"] => Some(TypedMigration::MetadataRetry), + ["qos_snapshot.db"] => Some(TypedMigration::QosSnapshot), + ["qos_snapshot_meta.db"] => Some(TypedMigration::QosCheckpoint), + ["geoip.db"] => Some(TypedMigration::GeoIp), + [directory, "metadata_retry_state.db"] if directory.starts_with("input_") => { + Some(TypedMigration::MetadataRetry) + } + [directory, name] if directory.starts_with("input_") && name.starts_with("m3u_") => { + Some(TypedMigration::InputM3u) + } + [directory, name] + if directory.starts_with("input_") + && (name.starts_with("lib_") || name.starts_with("media_server_")) => + { + Some(TypedMigration::Library) + } + [directory, name] + if directory.starts_with("input_") + && matches!(*name, "live.db" | "video.db" | "series.db") => + { + Some(TypedMigration::InputXtream) + } + [_, "id_mapping.db"] => Some(TypedMigration::IdMapping), + [_, "id_mapping.uuid.db"] => Some(TypedMigration::UuidMapping), + [_, "m3u", "m3u.db"] => Some(TypedMigration::TargetM3u), + [_, "m3u" | "xtream", "epg.db"] => Some(TypedMigration::Epg), + [_, "xtream", "live.db" | "video.db" | "series.db"] => Some(TypedMigration::TargetXtream), + _ => None, + } + } + +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TypedMigration { + MetadataRetry, + QosSnapshot, + QosCheckpoint, + GeoIp, + IdMapping, + UuidMapping, + TargetM3u, + InputM3u, + Library, + TargetXtream, + InputXtream, + Epg, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -693,6 +766,44 @@ fn create_user_db_merge_guard(merge_guard_path: &Path) -> io::Result<()> { pub(crate) fn user_db_merge_guard_path(config_dir: &Path) -> PathBuf { config_dir.join(MARKER_FILE_API_USER_GUARD) } +fn migrate_legacy_user_schema( + db_path: &Path, + merge_guard_path: &Path, + map: Map, +) -> io::Result +where + SourceV: serde::Serialize + for<'de> serde::Deserialize<'de> + Clone, + Map: FnMut(SourceV) -> StoredApiUserV6, +{ + if super::v3::migration::migrate_v2_typed_map::(db_path, map).is_err() + { + return Ok(false); + } + create_user_db_merge_guard(merge_guard_path)?; + Ok(true) +} + +fn migrate_current_user_schema( + db_path: &Path, + merge_guard_path: &Path, + map: Map, +) -> io::Result +where + SourceV: for<'de> serde::Deserialize<'de>, + Map: Fn(&SourceV) -> StoredApiUserV6, +{ + let Ok(tree) = BPlusTree::::load(db_path) else { + return Ok(false); + }; + let mut v6_tree = BPlusTree::new(); + for (key, user) in &tree { + v6_tree.insert(key.clone(), map(user)); + } + v6_tree.store(db_path)?; + create_user_db_merge_guard(merge_guard_path)?; + Ok(true) +} + /// Migrates the user database file from V1-V5 schema to V6 (current) in /// place and creates a merge-guard file so config-driven merges are skipped /// until the operator explicitly removes it. @@ -704,57 +815,65 @@ fn migrate_user_db_schema(db_path: &Path, merge_guard_path: &Path) -> io::Result return Ok(false); } - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v6_tree: BPlusTree = BPlusTree::new(); - for (key, v5) in &tree { - v6_tree.insert(key.clone(), StoredApiUserV6::from_v5(v5)); + let storage_version = super::v3::migration::storage_version(db_path)? + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "user database is not a B+Tree"))?; + if storage_version <= 2 { + if super::v3::migration::migrate_v2_typed::(db_path).is_ok() { + return Ok(false); } - create_user_db_merge_guard(merge_guard_path)?; - v6_tree.store(db_path)?; - return Ok(true); + if migrate_legacy_user_schema::(db_path, merge_guard_path, |user| { + StoredApiUserV6::from_v5(&user) + })? { + return Ok(true); + } + if migrate_legacy_user_schema::(db_path, merge_guard_path, |user| { + StoredApiUserV6::from_v4(&user) + })? { + return Ok(true); + } + if migrate_legacy_user_schema::(db_path, merge_guard_path, |user| { + StoredApiUserV6::from_v3(&user) + })? { + return Ok(true); + } + if migrate_legacy_user_schema::(db_path, merge_guard_path, |user| { + StoredApiUserV6::from_v2(&user) + })? { + return Ok(true); + } + if migrate_legacy_user_schema::(db_path, merge_guard_path, |user| { + StoredApiUserV6::from_v1(&user) + })? { + return Ok(true); + } + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Legacy user DB at '{}' could not be read as V1, V2, V3, V4, V5, or V6", db_path.display()), + )); + } + if storage_version != STORAGE_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported user B+Tree storage version {storage_version} in {}", db_path.display()), + )); } if BPlusTree::::load(db_path).is_ok() { return Ok(false); } - - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v6_tree: BPlusTree = BPlusTree::new(); - for (key, v4) in &tree { - v6_tree.insert(key.clone(), StoredApiUserV6::from_v4(v4)); - } - create_user_db_merge_guard(merge_guard_path)?; - v6_tree.store(db_path)?; + if migrate_current_user_schema::(db_path, merge_guard_path, StoredApiUserV6::from_v5)? { return Ok(true); } - - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v6_tree: BPlusTree = BPlusTree::new(); - for (key, v3) in &tree { - v6_tree.insert(key.clone(), StoredApiUserV6::from_v3(v3)); - } - create_user_db_merge_guard(merge_guard_path)?; - v6_tree.store(db_path)?; + if migrate_current_user_schema::(db_path, merge_guard_path, StoredApiUserV6::from_v4)? { return Ok(true); } - - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v6_tree: BPlusTree = BPlusTree::new(); - for (key, v2) in &tree { - v6_tree.insert(key.clone(), StoredApiUserV6::from_v2(v2)); - } - create_user_db_merge_guard(merge_guard_path)?; - v6_tree.store(db_path)?; + if migrate_current_user_schema::(db_path, merge_guard_path, StoredApiUserV6::from_v3)? { return Ok(true); } - - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v6_tree: BPlusTree = BPlusTree::new(); - for (key, v1) in &tree { - v6_tree.insert(key.clone(), StoredApiUserV6::from_v1(v1)); - } - create_user_db_merge_guard(merge_guard_path)?; - v6_tree.store(db_path)?; + if migrate_current_user_schema::(db_path, merge_guard_path, StoredApiUserV6::from_v2)? { + return Ok(true); + } + if migrate_current_user_schema::(db_path, merge_guard_path, StoredApiUserV6::from_v1)? { return Ok(true); } @@ -781,12 +900,11 @@ fn run_all_startup_migrations( storage_dir: &Path, config_dir: &Path, ) -> io::Result { - let marker_path = bplustree_migration_marker_path(storage_dir); - let bplustree = BPlusTreeStartupMigrator::new_with_marker(roots.to_vec(), marker_path).run()?; - let user_db_path = config_dir.join(storage_const::API_USER_DB_FILE); let merge_guard_path = user_db_merge_guard_path(config_dir); let user_db_migrated = migrate_user_db_schema(&user_db_path, &merge_guard_path)?; + let marker_path = bplustree_migration_marker_path(storage_dir); + let bplustree = BPlusTreeStartupMigrator::new_with_marker(roots.to_vec(), marker_path).run()?; Ok(AllStartupMigrationStats { bplustree, user_db_migrated }) } @@ -810,9 +928,10 @@ pub fn run_startup_migrations(config_paths: &ConfigPaths) { match run_all_startup_migrations(&roots, &storage_dir, &config_dir) { Ok(stats) => { - if stats.bplustree.skipped_by_marker { - trace!("B+Tree startup migration skipped (marker already present)"); - } else if stats.bplustree.migrated_files > 0 { + if stats.bplustree.marker_was_current { + trace!("B+Tree startup migration marker was current; database headers were still scanned"); + } + if stats.bplustree.migrated_files > 0 { info!( "B+Tree startup migration completed: migrated {} file(s) ({} B+Tree files checked, {} .db files scanned)", stats.bplustree.migrated_files, @@ -920,18 +1039,24 @@ mod tests { BPlusTreeStartupMigrator::roots_fingerprint(&resolved) } + fn write_legacy_geoip(root: &Path, version: u32) -> io::Result { + let path = root.join("geoip.db"); + let mut tree = super::super::v2::BPlusTree::new(); + tree.insert(1u32, (2u32, String::from("DE"))); + tree.store(&path)?; + if version == LEGACY_STORAGE_VERSION { + let mut file = OpenOptions::new().write(true).open(&path)?; + file.seek(SeekFrom::Start(4))?; + file.write_all(&version.to_le_bytes())?; + file.sync_all()?; + } + Ok(path) + } + #[test] fn startup_migrator_upgrades_legacy_bplustree_files() -> io::Result<()> { let temp = tempdir()?; - let db_path = temp.path().join("legacy.db"); - - let mut file = OpenOptions::new().create(true).truncate(true).read(true).write(true).open(&db_path)?; - let mut header = [0u8; 4096]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&LEGACY_STORAGE_VERSION.to_le_bytes()); - file.write_all(&header)?; - file.flush()?; - drop(file); + let db_path = write_legacy_geoip(temp.path(), LEGACY_STORAGE_VERSION)?; let stats = migrate_bplustree_databases(&[temp.path().to_path_buf()])?; assert_eq!(stats.scanned_files, 1); @@ -945,11 +1070,8 @@ mod tests { version_bytes[4..8].try_into().map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?, ); assert_eq!(version, STORAGE_VERSION); - let metadata_len_raw = u32::from_le_bytes( - version_bytes[16..20].try_into().map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?, - ); - assert_ne!(metadata_len_raw & HEADER_FLAG_HAS_METADATA_FLAGS, 0); - assert_eq!(metadata_len_raw & HEADER_FLAG_HAS_TOMBSTONES, 0); + let query = BPlusTree::::load(&db_path)?; + assert_eq!(query.query(&1), Some(&(2, String::from("DE")))); Ok(()) } @@ -971,35 +1093,138 @@ mod tests { Ok(()) } + #[test] + fn startup_migrator_rejects_an_unknown_valid_legacy_tree_without_changing_it() -> io::Result<()> { + let temp = tempdir()?; + let path = temp.path().join("backup.db"); + let mut tree = super::super::v2::BPlusTree::new(); + tree.insert(1u32, String::from("preserve")); + tree.store(&path)?; + let before = std::fs::read(&path)?; + + assert!(migrate_bplustree_databases(&[temp.path().to_path_buf()]).is_err()); + assert_eq!(std::fs::read(&path)?, before); + Ok(()) + } + + #[test] + fn startup_migrator_reports_a_locked_legacy_tree_without_modifying_it() -> io::Result<()> { + let temp = tempdir()?; + let path = write_legacy_geoip(temp.path(), 2)?; + let before = std::fs::read(&path)?; + let file = OpenOptions::new().read(true).write(true).open(&path)?; + file.lock_exclusive()?; + + let stats = migrate_bplustree_databases(&[temp.path().to_path_buf()])?; + assert_eq!(stats.bplustree_files, 1); + assert_eq!(stats.migrated_files, 0); + assert_eq!(std::fs::read(&path)?, before); + fs2::FileExt::unlock(&file)?; + Ok(()) + } + + #[test] + fn startup_migrator_only_removes_exact_abandoned_v3_temporary_names() -> io::Result<()> { + let temp = tempdir()?; + let abandoned = temp.path().join(format!("geoip.db.{}.v3.tmp", uuid::Uuid::new_v4())); + let unrelated = temp.path().join("geoip.db.not-a-uuid.v3.tmp"); + std::fs::write(&abandoned, b"partial")?; + std::fs::write(&unrelated, b"keep")?; + + let stats = migrate_bplustree_databases(&[temp.path().to_path_buf()])?; + assert_eq!(stats.scanned_files, 0); + assert!(!abandoned.exists()); + assert!(unrelated.exists()); + Ok(()) + } + + #[test] + fn typed_registry_matches_only_the_persisted_path_families() -> io::Result<()> { + let temp = tempdir()?; + let roots = [temp.path().to_path_buf()]; + let cases = [ + ("input_news/metadata_retry_state.db", TypedMigration::MetadataRetry), + ("qos_snapshot.db", TypedMigration::QosSnapshot), + ("qos_snapshot_meta.db", TypedMigration::QosCheckpoint), + ("geoip.db", TypedMigration::GeoIp), + ("target/id_mapping.db", TypedMigration::IdMapping), + ("target/id_mapping.uuid.db", TypedMigration::UuidMapping), + ("target/m3u/m3u.db", TypedMigration::TargetM3u), + ("input_news/m3u_news.db", TypedMigration::InputM3u), + ("input_news/lib_news.db", TypedMigration::Library), + ("input_news/media_server_news.db", TypedMigration::Library), + ("target/xtream/live.db", TypedMigration::TargetXtream), + ("target/xtream/video.db", TypedMigration::TargetXtream), + ("input_news/live.db", TypedMigration::InputXtream), + ("input_news/video.db", TypedMigration::InputXtream), + ("target/m3u/epg.db", TypedMigration::Epg), + ("target/xtream/epg.db", TypedMigration::Epg), + ]; + for (relative, expected) in cases { + assert_eq!(BPlusTreeStartupMigrator::migration_kind(&temp.path().join(relative), &roots), Some(expected)); + } + assert_eq!(BPlusTreeStartupMigrator::migration_kind(&temp.path().join("backup/tree.db"), &roots), None); + assert_eq!(BPlusTreeStartupMigrator::migration_kind(&temp.path().join("input_news/vod.db"), &roots), None); + Ok(()) + } + + #[test] + fn startup_migrator_exercises_every_typed_registry_family() -> io::Result<()> { + let temp = tempdir()?; + let mut paths = Vec::new(); + macro_rules! empty_v2 { + ($relative:literal, $key:ty, $value:ty) => {{ + let path = temp.path().join($relative); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let mut tree = super::super::v2::BPlusTree::<$key, $value>::new(); + tree.store(&path)?; + paths.push(path); + }}; + } + empty_v2!("input_news/metadata_retry_state.db", MetadataRetryDbKey, MetadataRetryDbValue); + empty_v2!("qos_snapshot.db", String, QosSnapshotRecord); + empty_v2!("qos_snapshot_meta.db", u8, QosAggregationCheckpoint); + empty_v2!("geoip.db", u32, (u32, String)); + empty_v2!("target/id_mapping.db", u32, VirtualIdRecord); + empty_v2!("target/id_mapping.uuid.db", UUIDType, u32); + empty_v2!("target/m3u/m3u.db", u32, M3uPlaylistItem); + empty_v2!("input_news/m3u_news.db", Arc, M3uPlaylistItem); + empty_v2!("input_news/lib_news.db", UUIDType, XtreamPlaylistItem); + empty_v2!("input_news/media_server_news.db", UUIDType, XtreamPlaylistItem); + empty_v2!("target/xtream/live.db", u32, XtreamPlaylistItem); + empty_v2!("target/xtream/video.db", u32, XtreamPlaylistItem); + empty_v2!("target/xtream/series.db", u32, XtreamPlaylistItem); + empty_v2!("input_news/live.db", u32, XtreamPlaylistItem); + empty_v2!("input_news/video.db", u32, XtreamPlaylistItem); + empty_v2!("input_news/series.db", u32, XtreamPlaylistItem); + empty_v2!("target/m3u/epg.db", Arc, EpgChannel); + empty_v2!("target/xtream/epg.db", Arc, EpgChannel); + + let stats = migrate_bplustree_databases(&[temp.path().to_path_buf()])?; + assert_eq!(stats.migrated_files, paths.len()); + for path in paths { + assert_eq!(super::super::v3::migration::storage_version(&path)?, Some(STORAGE_VERSION)); + } + assert!(temp.path().join("target/m3u/m3u.idx").exists()); + assert!(temp.path().join("target/xtream/live.idx").exists()); + Ok(()) + } + #[test] fn startup_migrator_writes_marker_after_success() -> io::Result<()> { let temp = tempdir()?; let temp_other = tempdir()?; - let db_path = temp.path().join("legacy.db"); - let db_path_other = temp_other.path().join("legacy_other.db"); - let mut file = OpenOptions::new().create(true).truncate(true).read(true).write(true).open(&db_path)?; - let mut header = [0u8; 4096]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&LEGACY_STORAGE_VERSION.to_le_bytes()); - file.write_all(&header)?; - file.flush()?; - drop(file); - - let mut file_other = - OpenOptions::new().create(true).truncate(true).read(true).write(true).open(&db_path_other)?; - let mut header_other = [0u8; 4096]; - header_other[0..4].copy_from_slice(MAGIC); - header_other[4..8].copy_from_slice(&LEGACY_STORAGE_VERSION.to_le_bytes()); - file_other.write_all(&header_other)?; - file_other.flush()?; - drop(file_other); + let _ = write_legacy_geoip(temp.path(), LEGACY_STORAGE_VERSION)?; + let _ = write_legacy_geoip(temp_other.path(), LEGACY_STORAGE_VERSION)?; let stats = migrate_bplustree_databases_with_marker( &[temp.path().to_path_buf(), temp_other.path().to_path_buf()], temp.path(), )?; assert_eq!(stats.migrated_files, 2); - assert!(!stats.skipped_by_marker); + assert!(!stats.marker_was_current); let marker = bplustree_migration_marker_path(temp.path()); assert!(marker.exists()); @@ -1011,16 +1236,9 @@ mod tests { } #[test] - fn startup_migrator_skips_when_marker_exists() -> io::Result<()> { + fn startup_migrator_scans_when_marker_exists() -> io::Result<()> { let temp = tempdir()?; - let db_path = temp.path().join("legacy.db"); - let mut file = OpenOptions::new().create(true).truncate(true).read(true).write(true).open(&db_path)?; - let mut header = [0u8; 4096]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&LEGACY_STORAGE_VERSION.to_le_bytes()); - file.write_all(&header)?; - file.flush()?; - drop(file); + let db_path = write_legacy_geoip(temp.path(), LEGACY_STORAGE_VERSION)?; let marker = bplustree_migration_marker_path(temp.path()); let roots = [temp.path().to_path_buf()]; @@ -1028,10 +1246,10 @@ mod tests { BPlusTreeStartupMigrator::write_migration_marker(&marker, &fingerprint)?; let stats = migrate_bplustree_databases_with_marker(&roots, temp.path())?; - assert_eq!(stats.scanned_files, 0); - assert_eq!(stats.bplustree_files, 0); - assert_eq!(stats.migrated_files, 0); - assert!(stats.skipped_by_marker); + assert_eq!(stats.scanned_files, 1); + assert_eq!(stats.bplustree_files, 1); + assert_eq!(stats.migrated_files, 1); + assert!(stats.marker_was_current); let mut check = OpenOptions::new().read(true).open(&db_path)?; let mut version_bytes = [0u8; 8]; @@ -1039,7 +1257,7 @@ mod tests { let version = u32::from_le_bytes( version_bytes[4..8].try_into().map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?, ); - assert_eq!(version, LEGACY_STORAGE_VERSION); + assert_eq!(version, STORAGE_VERSION); Ok(()) } @@ -1059,7 +1277,7 @@ mod tests { assert!(legacy_per_root_marker.exists()); let stats = migrate_bplustree_databases_with_marker(&roots, marker_dir)?; - assert!(stats.skipped_by_marker); + assert!(stats.marker_was_current); assert!(!legacy_per_root_marker.exists()); assert!(global_marker.exists()); @@ -1070,15 +1288,7 @@ mod tests { fn startup_migrator_does_not_skip_when_marker_fingerprint_differs() -> io::Result<()> { let temp_a = tempdir()?; let temp_b = tempdir()?; - let db_path = temp_a.path().join("legacy.db"); - - let mut file = OpenOptions::new().create(true).truncate(true).read(true).write(true).open(&db_path)?; - let mut header = [0u8; 4096]; - header[0..4].copy_from_slice(MAGIC); - header[4..8].copy_from_slice(&LEGACY_STORAGE_VERSION.to_le_bytes()); - file.write_all(&header)?; - file.flush()?; - drop(file); + let _ = write_legacy_geoip(temp_a.path(), LEGACY_STORAGE_VERSION)?; let marker = bplustree_migration_marker_path(temp_a.path()); let old_roots = [temp_a.path().to_path_buf()]; @@ -1087,7 +1297,7 @@ mod tests { let current_roots = [temp_a.path().to_path_buf(), temp_b.path().to_path_buf()]; let stats = migrate_bplustree_databases_with_marker(¤t_roots, temp_a.path())?; - assert!(!stats.skipped_by_marker); + assert!(!stats.marker_was_current); assert_eq!(stats.migrated_files, 1); Ok(()) @@ -1099,7 +1309,8 @@ mod tests { let db_path = temp.path().join(storage_const::API_USER_DB_FILE); let merge_guard_path = user_db_merge_guard_path(temp.path()); - let mut v2_tree: BPlusTree = BPlusTree::new(); + let mut v2_tree: super::super::v2::BPlusTree = + super::super::v2::BPlusTree::new(); v2_tree.insert( "alice".to_string(), StoredApiUserV2 { @@ -1292,7 +1503,8 @@ mod tests { let db_path = temp.path().join(storage_const::API_USER_DB_FILE); let merge_guard_path = user_db_merge_guard_path(temp.path()); - let mut v6_tree: BPlusTree = BPlusTree::new(); + let mut v6_tree: super::super::v2::BPlusTree = + super::super::v2::BPlusTree::new(); v6_tree.insert( "erin".to_string(), StoredApiUserV6 { diff --git a/backend/src/repository/bplustree/mod.rs b/backend/src/repository/bplustree/mod.rs new file mode 100644 index 000000000..c37792e93 --- /dev/null +++ b/backend/src/repository/bplustree/mod.rs @@ -0,0 +1,11 @@ +pub(crate) mod common; +mod migration; +pub(super) mod sorted_index; +#[cfg(test)] +mod stress; +pub(crate) mod v2; +pub(crate) mod v3; + +pub use common::BPlusTreeError; +pub use self::migration::*; +pub use v3::*; diff --git a/backend/src/repository/bplustree/sorted_index.rs b/backend/src/repository/bplustree/sorted_index.rs new file mode 100644 index 000000000..e8d1f03dc --- /dev/null +++ b/backend/src/repository/bplustree/sorted_index.rs @@ -0,0 +1,477 @@ +pub(crate) mod v4 { + use crate::{ + repository::bplustree::v3::{BPlusTreeQuery, Locator}, + utils::{binary_deserialize, binary_serialize}, + }; + use serde::{Deserialize, Serialize}; + use std::{ + fs::{File, OpenOptions}, + io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write}, + marker::PhantomData, + num::TryFromIntError, + path::Path, + }; + + const MAGIC: &[u8; 4] = b"SIDX"; + const VERSION: u32 = 4; + const TREE_VERSION: u32 = 3; + const HEADER_LEN: usize = 64; + const HEADER_LEN_U32: u32 = 64; + const HEADER_CRC_OFFSET: usize = 56; + const ENTRY_PREFIX_LEN: usize = 8; + const ENTRY_PREFIX_LEN_U64: u64 = 8; + const ENTRY_FIXED_BODY_LEN: usize = 24; + + fn invalid_data(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, message.into()) + } + + fn invalid_integer(error: TryFromIntError) -> io::Error { invalid_data(error.to_string()) } + + fn header(database_id: [u8; 16], generation: u64, count: u64) -> [u8; HEADER_LEN] { + let mut bytes = [0; HEADER_LEN]; + bytes[0..4].copy_from_slice(MAGIC); + bytes[4..8].copy_from_slice(&VERSION.to_le_bytes()); + bytes[8..12].copy_from_slice(&TREE_VERSION.to_le_bytes()); + bytes[12..16].copy_from_slice(&HEADER_LEN_U32.to_le_bytes()); + bytes[16..32].copy_from_slice(&database_id); + bytes[32..40].copy_from_slice(&generation.to_le_bytes()); + bytes[40..48].copy_from_slice(&count.to_le_bytes()); + let checksum = crc32fast::hash(&bytes); + bytes[HEADER_CRC_OFFSET..HEADER_CRC_OFFSET + 4].copy_from_slice(&checksum.to_le_bytes()); + bytes + } + + fn read_u32(bytes: &[u8], offset: usize) -> io::Result { + bytes + .get(offset..offset + 4) + .and_then(|value| value.try_into().ok()) + .map(u32::from_le_bytes) + .ok_or_else(|| invalid_data("sorted-index u32 field is truncated")) + } + + fn read_u64(bytes: &[u8], offset: usize) -> io::Result { + bytes + .get(offset..offset + 8) + .and_then(|value| value.try_into().ok()) + .map(u64::from_le_bytes) + .ok_or_else(|| invalid_data("sorted-index u64 field is truncated")) + } + + fn decode_header( + mut bytes: [u8; HEADER_LEN], + expected_database_id: [u8; 16], + expected_generation: u64, + ) -> io::Result { + if bytes[0..4] != *MAGIC + || read_u32(&bytes, 4)? != VERSION + || read_u32(&bytes, 8)? != TREE_VERSION + || read_u32(&bytes, 12)? != HEADER_LEN_U32 + { + return Err(invalid_data("unsupported sorted-index header")); + } + if bytes[48..56].iter().chain(&bytes[60..64]).any(|byte| *byte != 0) { + return Err(invalid_data("sorted-index reserved header bytes must be zero")); + } + let expected_checksum = read_u32(&bytes, HEADER_CRC_OFFSET)?; + bytes[HEADER_CRC_OFFSET..HEADER_CRC_OFFSET + 4].fill(0); + if crc32fast::hash(&bytes) != expected_checksum { + return Err(invalid_data("sorted-index header checksum mismatch")); + } + if bytes[16..32] != expected_database_id || read_u64(&bytes, 32)? != expected_generation { + return Err(invalid_data("sorted-index tree identity or generation mismatch")); + } + read_u64(&bytes, 40) + } + + pub(crate) struct Writer { + output: BufWriter, + database_id: [u8; 16], + generation: u64, + count: u64, + body: Vec, + _marker: PhantomData<(SortKey, K)>, + } + + impl Writer + where + SortKey: Serialize, + K: Serialize, + { + pub(crate) fn new(path: &Path, database_id: [u8; 16], generation: u64) -> io::Result { + let file = OpenOptions::new().write(true).create_new(true).open(path)?; + let mut output = BufWriter::new(file); + output.write_all(&header(database_id, generation, 0))?; + Ok(Self { output, database_id, generation, count: 0, body: Vec::new(), _marker: PhantomData }) + } + + pub(crate) fn push(&mut self, sort_key: &SortKey, primary_key: &K, locator: Locator) -> io::Result<()> { + let sort_key = binary_serialize(sort_key)?; + let primary_key = binary_serialize(primary_key)?; + let sort_key_len = u32::try_from(sort_key.len()).map_err(invalid_integer)?; + let primary_key_len = u32::try_from(primary_key.len()).map_err(invalid_integer)?; + self.body.clear(); + self.body.extend_from_slice(&sort_key_len.to_le_bytes()); + self.body.extend_from_slice(&primary_key_len.to_le_bytes()); + self.body.extend_from_slice(&locator.encode()); + self.body.extend_from_slice(&sort_key); + self.body.extend_from_slice(&primary_key); + let body_len = u32::try_from(self.body.len()).map_err(invalid_integer)?; + self.output.write_all(&body_len.to_le_bytes())?; + self.output.write_all(&crc32fast::hash(&self.body).to_le_bytes())?; + self.output.write_all(&self.body)?; + self.count = self.count.checked_add(1).ok_or_else(|| invalid_data("sorted-index count overflow"))?; + Ok(()) + } + + pub(crate) fn finish(mut self) -> io::Result { + self.output.flush()?; + let mut file = self.output.into_inner().map_err(std::io::IntoInnerError::into_error)?; + file.seek(SeekFrom::Start(0))?; + file.write_all(&header(self.database_id, self.generation, self.count))?; + file.sync_all()?; + Ok(self.count) + } + } + + #[cfg(test)] + pub(crate) struct Entry { + pub(crate) sort_key: SortKey, + pub(crate) primary_key: K, + pub(crate) locator: Locator, + } + + struct BorrowedEntry<'a, SortKey, K> { + sort_key: SortKey, + primary_key: K, + serialized_primary_key: &'a [u8], + locator: Locator, + } + + pub(crate) struct Reader { + input: BufReader, + remaining_entries: u64, + remaining_bytes: u64, + failed: bool, + body: Vec, + _marker: PhantomData<(SortKey, K)>, + } + + impl Reader + where + SortKey: for<'de> Deserialize<'de>, + K: for<'de> Deserialize<'de>, + { + pub(crate) fn open( + path: &Path, + expected_database_id: [u8; 16], + expected_generation: u64, + ) -> io::Result { + let file = File::open(path)?; + let length = file.metadata()?.len(); + if length < HEADER_LEN as u64 { + return Err(invalid_data("sorted-index header is truncated")); + } + let mut input = BufReader::new(file); + let mut encoded = [0; HEADER_LEN]; + input.read_exact(&mut encoded)?; + let remaining_entries = decode_header(encoded, expected_database_id, expected_generation)?; + Ok(Self { + input, + remaining_entries, + remaining_bytes: length - HEADER_LEN as u64, + failed: false, + body: Vec::new(), + _marker: PhantomData, + }) + } + + #[cfg(test)] + pub(crate) fn remaining(&self) -> u64 { self.remaining_entries } + + #[cfg(test)] + pub(crate) fn read_next(&mut self) -> io::Result>> { + let result = match self.read_next_borrowed() { + Ok(entry) => Ok(entry.map(|entry| Entry { + sort_key: entry.sort_key, + primary_key: entry.primary_key, + locator: entry.locator, + })), + Err(error) => Err(error), + }; + if result.is_err() { + self.failed = true; + self.remaining_entries = 0; + } + result + } + + fn read_next_borrowed(&mut self) -> io::Result>> { + if self.failed { + return Ok(None); + } + if self.remaining_entries == 0 { + if self.remaining_bytes == 0 { + return Ok(None); + } + self.failed = true; + self.remaining_bytes = 0; + return Err(invalid_data("sorted-index contains trailing bytes")); + } + self.read_next_inner() + } + + fn read_next_inner(&mut self) -> io::Result>> { + if self.remaining_bytes < ENTRY_PREFIX_LEN_U64 { + self.failed = true; + return Err(invalid_data("sorted-index entry prefix is truncated")); + } + let mut prefix = [0; ENTRY_PREFIX_LEN]; + if let Err(error) = self.input.read_exact(&mut prefix) { + self.failed = true; + return Err(error); + } + self.remaining_bytes -= ENTRY_PREFIX_LEN_U64; + let body_len = u64::from(read_u32(&prefix, 0)?); + let expected_crc = read_u32(&prefix, 4)?; + if body_len < ENTRY_FIXED_BODY_LEN as u64 || body_len > self.remaining_bytes { + self.failed = true; + return Err(invalid_data("sorted-index entry length is invalid")); + } + let body_len = match usize::try_from(body_len) { + Ok(body_len) => body_len, + Err(error) => { + self.failed = true; + return Err(invalid_integer(error)); + } + }; + self.body.resize(body_len, 0); + if let Err(error) = self.input.read_exact(&mut self.body) { + self.failed = true; + return Err(error); + } + self.remaining_bytes -= u64::try_from(body_len).map_err(invalid_integer)?; + self.remaining_entries -= 1; + if crc32fast::hash(&self.body) != expected_crc { + return Err(invalid_data("sorted-index entry checksum mismatch")); + } + let sort_key_len = usize::try_from(read_u32(&self.body, 0)?).map_err(invalid_integer)?; + let primary_key_len = usize::try_from(read_u32(&self.body, 4)?).map_err(invalid_integer)?; + let expected_len = ENTRY_FIXED_BODY_LEN + .checked_add(sort_key_len) + .and_then(|length| length.checked_add(primary_key_len)) + .ok_or_else(|| invalid_data("sorted-index entry length overflow"))?; + if expected_len != body_len { + return Err(invalid_data("sorted-index body length mismatch")); + } + let locator = Locator::decode( + self.body.get(8..24).ok_or_else(|| invalid_data("sorted-index locator is truncated"))?, + )?; + let sort_key_end = 24 + sort_key_len; + let sort_key = binary_deserialize( + self.body.get(24..sort_key_end).ok_or_else(|| invalid_data("sort key is truncated"))?, + )?; + let serialized_primary_key = self + .body + .get(sort_key_end..expected_len) + .ok_or_else(|| invalid_data("primary key is truncated"))?; + let primary_key = binary_deserialize(serialized_primary_key)?; + Ok(Some(BorrowedEntry { sort_key, primary_key, serialized_primary_key, locator })) + } + } + + pub(crate) struct OwnedIterator { + reader: Reader, + query: BPlusTreeQuery, + previous_sort_key: Option, + finished: bool, + } + + impl OwnedIterator + where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, + SortKey: Ord + for<'de> Deserialize<'de>, + { + pub(crate) fn open(query: BPlusTreeQuery, index_path: &Path) -> io::Result { + let (database_id, generation) = query.snapshot_identity(); + let reader = Reader::open(index_path, database_id, generation)?; + Ok(Self { reader, query, previous_sort_key: None, finished: false }) + } + + #[cfg(test)] + pub(crate) fn remaining(&self) -> u64 { self.reader.remaining() } + } + + impl Iterator for OwnedIterator + where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, + SortKey: Ord + for<'de> Deserialize<'de>, + { + type Item = io::Result<(K, V)>; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + let result = match self.reader.read_next_borrowed() { + Ok(Some(entry)) => { + if self.previous_sort_key.as_ref().is_some_and(|previous| previous > &entry.sort_key) { + Err(invalid_data("sorted-index entries are out of order")) + } else { + let BorrowedEntry { sort_key, primary_key, serialized_primary_key, locator } = entry; + self.previous_sort_key = Some(sort_key); + self.query + .read_locator_value(locator, serialized_primary_key) + .map(|value| (primary_key, value)) + } + } + Ok(None) => { + self.finished = true; + return None; + } + Err(error) => Err(error), + }; + match result { + Ok(entry) => Some(Ok(entry)), + Err(error) => Some(Err(error)), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::repository::bplustree::v3::{BPlusTree as V3Tree, BPlusTreeQuery as V3Query, Locator}; + use crate::utils::binary_serialize; + use std::{fs, io, io::Write}; + use tempfile::tempdir; + + #[test] + fn v4_header_entries_and_identity_are_exact_and_checked() -> io::Result<()> { + let dir = tempdir()?; + let path = dir.path().join("v4.idx"); + let database_id = [0x5a; 16]; + let locator = Locator { leaf_page_id: 7, slot_index: 3, serialized_key_crc32: 0x1122_3344 }; + let mut writer = v4::Writer::::new(&path, database_id, 9)?; + writer.push(&String::from("sort"), &42, locator)?; + assert_eq!(writer.finish()?, 1); + + let bytes = fs::read(&path)?; + assert_eq!(&bytes[0..4], b"SIDX"); + assert_eq!(u32::from_le_bytes(bytes[4..8].try_into().map_err(io::Error::other)?), 4); + assert_eq!(u32::from_le_bytes(bytes[8..12].try_into().map_err(io::Error::other)?), 3); + assert_eq!(u32::from_le_bytes(bytes[12..16].try_into().map_err(io::Error::other)?), 64); + assert_eq!(&bytes[16..32], &database_id); + assert_eq!(u64::from_le_bytes(bytes[32..40].try_into().map_err(io::Error::other)?), 9); + assert_eq!(u64::from_le_bytes(bytes[40..48].try_into().map_err(io::Error::other)?), 1); + assert!(bytes[48..56].iter().chain(&bytes[60..64]).all(|byte| *byte == 0)); + + let mut reader = v4::Reader::::open(&path, database_id, 9)?; + let entry = reader.read_next()?.ok_or_else(|| io::Error::other("v4 entry missing"))?; + assert_eq!(entry.sort_key, "sort"); + assert_eq!(entry.primary_key, 42); + assert_eq!(entry.locator, locator); + assert!(reader.read_next()?.is_none()); + assert!(v4::Reader::::open(&path, [0; 16], 9).is_err()); + assert!(v4::Reader::::open(&path, database_id, 10).is_err()); + Ok(()) + } + + #[test] + fn v4_entry_corruption_is_reported_once_then_reader_fuses() -> io::Result<()> { + let dir = tempdir()?; + let path = dir.path().join("v4-corrupt.idx"); + let database_id = [7; 16]; + let mut writer = v4::Writer::::new(&path, database_id, 1)?; + writer.push(&1, &2, Locator::for_key(3, 0, &binary_serialize(&2u32)?)?)?; + writer.finish()?; + let mut bytes = fs::read(&path)?; + *bytes.get_mut(72).ok_or_else(|| io::Error::other("v4 body missing"))? ^= 1; + fs::write(&path, bytes)?; + + let mut reader = v4::Reader::::open(&path, database_id, 1)?; + assert!(reader.read_next().is_err()); + assert!(reader.read_next()?.is_none()); + Ok(()) + } + + #[test] + fn v4_trailing_bytes_are_reported_once_then_reader_fuses() -> io::Result<()> { + let dir = tempdir()?; + let path = dir.path().join("v4-trailing.idx"); + let database_id = [8; 16]; + let mut writer = v4::Writer::::new(&path, database_id, 1)?; + writer.push(&1, &2, Locator::for_key(3, 0, &binary_serialize(&2u32)?)?)?; + writer.finish()?; + let mut file = std::fs::OpenOptions::new().append(true).open(&path)?; + file.write_all(&[0xff])?; + file.sync_all()?; + + let mut reader = v4::Reader::::open(&path, database_id, 1)?; + assert!(reader.read_next()?.is_some()); + assert!(reader.read_next().is_err()); + assert!(reader.read_next()?.is_none()); + Ok(()) + } + + #[test] + fn v4_sorted_iterator_validates_locators_and_fuses_after_a_late_error() -> io::Result<()> { + let dir = tempdir()?; + let database = dir.path().join("tree.db"); + let index = dir.path().join("tree.idx"); + let mut tree = V3Tree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.store(&database)?; + let mut query = V3Query::::try_new(&database)?; + let entries = query.collect_with_locators()?; + let (database_id, generation) = query.snapshot_identity(); + drop(query); + + let mut writer = v4::Writer::::new(&index, database_id, generation)?; + writer.push(&1, &entries[0].0, entries[0].2)?; + writer.push( + &2, + &entries[1].0, + Locator { slot_index: u16::MAX, ..entries[1].2 }, + )?; + writer.finish()?; + + let query = V3Query::::try_new(&database)?; + let mut iterator = v4::OwnedIterator::::open(query, &index)?; + assert_eq!(iterator.next().transpose()?, Some((1, String::from("one")))); + assert!(iterator.next().is_some_and(|entry| entry.is_err())); + assert!(iterator.next().is_none()); + Ok(()) + } + + #[test] + fn v4_sorted_iterator_rejects_out_of_order_sort_keys() -> io::Result<()> { + let dir = tempdir()?; + let database = dir.path().join("tree.db"); + let index = dir.path().join("tree.idx"); + let mut tree = V3Tree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.store(&database)?; + let mut query = V3Query::::try_new(&database)?; + let entries = query.collect_with_locators()?; + let (database_id, generation) = query.snapshot_identity(); + drop(query); + + let mut writer = v4::Writer::::new(&index, database_id, generation)?; + writer.push(&2, &entries[1].0, entries[1].2)?; + writer.push(&1, &entries[0].0, entries[0].2)?; + writer.finish()?; + + let query = V3Query::::try_new(&database)?; + let mut iterator = v4::OwnedIterator::::open(query, &index)?; + assert_eq!(iterator.next().transpose()?, Some((2, String::from("two")))); + assert!(iterator.next().is_some_and(|entry| entry.is_err())); + assert!(iterator.next().is_none()); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree_stress.rs b/backend/src/repository/bplustree/stress.rs similarity index 73% rename from backend/src/repository/bplustree_stress.rs rename to backend/src/repository/bplustree/stress.rs index 95be08613..5567486f6 100644 --- a/backend/src/repository/bplustree_stress.rs +++ b/backend/src/repository/bplustree/stress.rs @@ -1,4 +1,4 @@ -use super::bplustree::{BPlusTree, BPlusTreeQuery, BPlusTreeSerialWriter, BPlusTreeUpdate, FlushPolicy}; +use super::{BPlusTree, BPlusTreeQuery, BPlusTreeSerialWriter, BPlusTreeUpdate, FlushPolicy}; use rand::{distr::Alphanumeric, prelude::*}; use std::{ io::Write, @@ -220,7 +220,7 @@ fn stress_test_bplustree() { writeln!(log_file, "\n[Phase 6] Compaction...").unwrap(); let mut tree_updater = BPlusTreeUpdate::::try_new(&filepath).unwrap(); let start = Instant::now(); - tree_updater.compact(&filepath).unwrap(); + tree_updater.compact().unwrap(); let duration = start.elapsed(); writeln!(log_file, "Time: {duration:.2?}").unwrap(); let size_phase6 = std::fs::metadata(&filepath).unwrap().len(); @@ -440,7 +440,8 @@ fn stress_test_bplustree() { let iter_start = Instant::now(); let mut iterated_count = 0usize; let mut iterator = iter_query.iter(); - for (_k, _v) in iterator.by_ref() { + for entry in iterator.by_ref() { + let (_k, _v) = entry.unwrap(); iterated_count += 1; if iterated_count >= iterator_target { break; @@ -510,7 +511,7 @@ fn stress_test_bplustree() { let final_size_before_compact = std::fs::metadata(&filepath).unwrap().len(); let final_compact_start = Instant::now(); - final_updater.compact(&filepath).unwrap(); + final_updater.compact().unwrap(); let final_compact_duration = final_compact_start.elapsed(); let final_size_after_compact = std::fs::metadata(&filepath).unwrap().len(); writeln!( @@ -522,3 +523,210 @@ fn stress_test_bplustree() { ) .unwrap(); } + +#[derive(Clone, serde::Deserialize, serde::Serialize)] +struct ComparisonValue { + sort_key: u32, + payload: String, +} + +fn comparison_dataset(count: usize) -> Vec<(String, ComparisonValue)> { + const SIZES: [usize; 5] = [96, 640, 48, 160, 320]; + (0..count) + .map(|index| { + let profile = index % SIZES.len(); + ( + format!("{profile}:{index:08}"), + ComparisonValue { + sort_key: u32::try_from(count - index).unwrap_or(u32::MAX), + payload: char::from(b'a' + u8::try_from(profile).unwrap_or(0)).to_string().repeat(SIZES[profile]), + }, + ) + }) + .collect() +} + +fn store_comparison_v2(path: &std::path::Path, entries: &[(String, ComparisonValue)]) -> std::io::Result<()> { + let mut tree = super::v2::BPlusTree::new(); + for (key, value) in entries { + tree.insert(key.clone(), value.clone()); + } + tree.store(path).map(|_| ()) +} + +fn store_comparison_v3( + path: &std::path::Path, + entries: &[(String, ComparisonValue)], +) -> std::io::Result<()> { + let mut tree = super::v3::BPlusTree::new(); + for (key, value) in entries { + tree.insert(key.clone(), value.clone()); + } + tree.store_with_index(path, |value| value.sort_key).map(|_| ()) +} + +fn comparison_runs( + name: &str, + mut operation: impl FnMut() -> std::io::Result<()>, +) -> std::io::Result> { + let mut runs = Vec::with_capacity(5); + for _ in 0..5 { + let started = Instant::now(); + operation()?; + runs.push(started.elapsed()); + } + let individual = runs.iter().map(|duration| format!("{duration:?}")).collect::>().join(", "); + println!("{name}: [{individual}], median={:?}", comparison_median(&runs)); + Ok(runs) +} + +fn comparison_median(runs: &[Duration]) -> Duration { + let mut sorted = runs.to_vec(); + sorted.sort_unstable(); + sorted[sorted.len() / 2] +} + +#[allow(clippy::too_many_lines)] +#[ignore = "explicit release-only v2/v3 persistence comparison"] +#[test] +fn bplustree_v2_v3_comparison() -> std::io::Result<()> { + use std::{hint::black_box, ops::Bound}; + + const ENTRY_COUNT: usize = 10_000; + const LOOKUP_COUNT: usize = 4_000; + const UPDATE_COUNT: usize = 1_000; + + let directory = tempfile::tempdir()?; + let v2_path = directory.path().join("comparison-v2.db"); + let v3_path = directory.path().join("comparison-v3.db"); + let entries = comparison_dataset(ENTRY_COUNT); + store_comparison_v2(&v2_path, &entries)?; + store_comparison_v3(&v3_path, &entries)?; + + println!("profiles: mapping=48B, M3U=96B, QoS=160B, user=320B, Xtream=640B; entries={ENTRY_COUNT}"); + println!("initial sizes: v2={} v3={}", std::fs::metadata(&v2_path)?.len(), std::fs::metadata(&v3_path)?.len()); + + let lookup_keys = entries + .iter() + .step_by((ENTRY_COUNT / LOOKUP_COUNT).max(1)) + .take(LOOKUP_COUNT) + .map(|(key, _)| key.clone()) + .collect::>(); + let v2_point = comparison_runs("v2 point lookup", || { + let mut query = super::v2::BPlusTreeQuery::::try_new(&v2_path)?; + for key in &lookup_keys { + let value = query + .range_iter(Bound::Included(key), Bound::Included(key)) + .next() + .transpose() + .map_err(super::v2::BPlusTreeError::to_io)?; + black_box(value); + } + Ok(()) + })?; + let v3_point = comparison_runs("v3 point lookup", || { + let mut query = super::v3::BPlusTreeQuery::::try_new(&v3_path)?; + for key in &lookup_keys { + black_box(query.query(key).map_err(super::BPlusTreeError::to_io)?); + } + Ok(()) + })?; + let held_query = super::v3::BPlusTreeQuery::::try_new(&v3_path)?; + comparison_runs("v3 shared snapshot clone + point lookup", || { + for key in &lookup_keys { + let mut query = held_query.try_clone()?; + black_box(query.query(key).map_err(super::BPlusTreeError::to_io)?); + } + Ok(()) + })?; + drop(held_query); + + let v2_scan = comparison_runs("v2 full scan", || { + let mut query = super::v2::BPlusTreeQuery::::try_new(&v2_path)?; + let count = query + .range_iter(Bound::Unbounded, Bound::Unbounded) + .collect::, _>>() + .map_err(super::v2::BPlusTreeError::to_io)? + .len(); + black_box(count); + Ok(()) + })?; + let v3_scan = comparison_runs("v3 mmap full scan", || { + let mut query = super::v3::BPlusTreeQuery::::try_new(&v3_path)?; + black_box(query.iter().collect::>>()?.len()); + Ok(()) + })?; + + comparison_runs("v2 sorted full scan", || { + let mut query = super::v2::BPlusTreeQuery::::try_new(&v2_path)?; + let mut values = query + .range_iter(Bound::Unbounded, Bound::Unbounded) + .collect::, _>>() + .map_err(super::v2::BPlusTreeError::to_io)?; + values.sort_unstable_by_key(|(_, value)| value.sort_key); + black_box(values.len()); + Ok(()) + })?; + let index_path = crate::repository::storage::get_file_path_for_db_index(&v3_path); + comparison_runs("v3 locator sorted full scan", || { + let query = super::v3::BPlusTreeQuery::::try_new(&v3_path)?; + let iterator = super::sorted_index::v4::OwnedIterator::::open(query, &index_path)?; + black_box(iterator.collect::>>()?.len()); + Ok(()) + })?; + + let smaller = entries + .iter() + .take(UPDATE_COUNT) + .map(|(key, value)| { + let mut value = value.clone(); + value.payload.truncate(32); + (key.clone(), value) + }) + .collect::>(); + let growing = entries + .iter() + .take(UPDATE_COUNT) + .map(|(key, value)| { + let mut value = value.clone(); + value.payload.push_str(&"z".repeat(1_024)); + (key.clone(), value) + }) + .collect::>(); + + for (label, updates) in [("equal/smaller", &smaller), ("growing", &growing)] { + let filename_label = label.replace('/', "-"); + let mut rewritten = entries.clone(); + for ((_, destination), (_, replacement)) in rewritten.iter_mut().take(UPDATE_COUNT).zip(updates) { + *destination = replacement.clone(); + } + comparison_runs(&format!("v2 {label} full rewrite"), || { + store_comparison_v2(&directory.path().join(format!("v2-{filename_label}.db")), &rewritten) + })?; + comparison_runs(&format!("v3 {label} 1000-item batch"), || { + let path = directory.path().join(format!("v3-{filename_label}.db")); + std::fs::copy(&v3_path, &path)?; + let mut updater = super::v3::BPlusTreeUpdate::::try_new(&path)?; + let references = updates.iter().map(|(key, value)| (key, value)).collect::>(); + updater.update_batch(&references).map_err(super::BPlusTreeError::to_io)?; + updater.commit() + })?; + } + + let compact_path = directory.path().join("v3-compact.db"); + std::fs::copy(&v3_path, &compact_path)?; + let mut updater = super::v3::BPlusTreeUpdate::::try_new(&compact_path)?; + let growing_refs = growing.iter().map(|(key, value)| (key, value)).collect::>(); + updater.update_batch(&growing_refs).map_err(super::BPlusTreeError::to_io)?; + updater.commit()?; + let grown_size = std::fs::metadata(&compact_path)?.len(); + let compact_runs = comparison_runs("v3 compaction", || updater.compact())?; + let compacted_size = std::fs::metadata(&compact_path)?.len(); + println!("v3 growth/compaction: grown={grown_size}, compacted={compacted_size}, median={:?}", comparison_median(&compact_runs)); + + let point_regression = comparison_median(&v3_point).as_secs_f64() / comparison_median(&v2_point).as_secs_f64() - 1.0; + let scan_regression = comparison_median(&v3_scan).as_secs_f64() / comparison_median(&v2_scan).as_secs_f64() - 1.0; + println!("gate point-latency regression: {:+.2}%", point_regression * 100.0); + println!("gate full-scan duration regression: {:+.2}%", scan_regression * 100.0); + Ok(()) +} diff --git a/backend/src/repository/bplustree/v2.rs b/backend/src/repository/bplustree/v2.rs new file mode 100644 index 000000000..5121af090 --- /dev/null +++ b/backend/src/repository/bplustree/v2.rs @@ -0,0 +1,1891 @@ +//! Legacy-compatible B+Tree storage v2. +//! +//! This module remains the active compatibility branch for existing v2 files. +//! Stabilization work must preserve `STORAGE_VERSION = 2` and avoid requiring a +//! rewrite of existing repositories. The future v3 storage line is expected to +//! live behind an explicit version boundary and a typed migration path. + +use crate::{ + repository::bplustree::common::{mmap_with_advice, read_exact_at_offset}, + utils, + utils::binary_deserialize, +}; +pub(crate) use crate::repository::bplustree::common::BPlusTreeError; +#[cfg(test)] +use crate::utils::binary_serialize_into; +#[cfg(test)] +use log::error; +use memmap2::{Advice, Mmap}; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use shared::error::to_io_error; +#[cfg(test)] +use shared::error::string_to_io_error; +use smallvec::{smallvec, SmallVec}; +use std::{ + collections::HashSet, + fs::File, + io::{self, BufReader, Read, Seek, SeekFrom}, + marker::PhantomData, + ops::Bound, + path::Path, +}; +#[cfg(test)] +use std::{ + borrow::Cow, + io::Write, +}; +#[cfg(test)] +use tempfile::NamedTempFile; +#[cfg(windows)] +use windows_sys::Win32::Foundation::CloseHandle; +#[cfg(windows)] +use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION}; + +const PAGE_SIZE: u16 = 4096; +pub const PAGE_SIZE_USIZE: usize = PAGE_SIZE as usize; +const LEN_SIZE: usize = 4; +const FLAG_SIZE: usize = 1; +pub(crate) const MAGIC: &[u8; 4] = b"BTRE"; +pub(crate) const STORAGE_VERSION: u32 = 2; +const HEADER_SIZE: u64 = PAGE_SIZE as u64; +#[cfg(test)] +const ROOT_OFFSET_POS: u64 = 8; +const METADATA_DATA_START_POS: usize = 20; +// Reserve space for metadata (e.g. 4096 - 16 = 4080 bytes max, but let's be safe) +const METADATA_MAX_SIZE: usize = 4000; +const HEADER_FLAG_HAS_METADATA_FLAGS: u32 = 1 << 31; +const HEADER_FLAG_HAS_TOMBSTONES: u32 = 1 << 30; +const HEADER_METADATA_LEN_MASK: u32 = !(HEADER_FLAG_HAS_METADATA_FLAGS | HEADER_FLAG_HAS_TOMBSTONES); + +#[inline] +#[cfg(test)] +const fn encode_metadata_len_with_flags(metadata_len: u32, has_tombstones: bool) -> u32 { + let mut encoded = metadata_len | HEADER_FLAG_HAS_METADATA_FLAGS; + if has_tombstones { + encoded |= HEADER_FLAG_HAS_TOMBSTONES; + } + encoded +} + +#[inline] +const fn decode_metadata_len_and_flags(raw: u32) -> (u32, bool) { + let metadata_len = raw & HEADER_METADATA_LEN_MASK; + let has_metadata_flags = (raw & HEADER_FLAG_HAS_METADATA_FLAGS) != 0; + let has_tombstones = if has_metadata_flags { + (raw & HEADER_FLAG_HAS_TOMBSTONES) != 0 + } else { + // Legacy v2 files (without header flags) are treated conservatively: + // assume tombstones may exist until a rewrite/compact writes proper flags. + true + }; + (metadata_len, has_tombstones) +} + +// v2 uses conservative runtime fanout instead of pretending that size_of::() +// predicts serialized key size. Multi-block nodes keep existing files compatible. +#[cfg(test)] +const DEFAULT_INNER_ORDER: usize = 64; +#[cfg(test)] +const DEFAULT_LEAF_ORDER: usize = 64; + +// Value packing configuration +#[cfg(test)] +const SMALL_VALUE_THRESHOLD: usize = 256; +#[cfg(test)] +const PACK_BLOCK_HEADER_SIZE: usize = 4; +#[cfg(test)] +const PACK_VALUE_HEADER_SIZE: usize = 4; + +// LZ4 compression configuration +#[cfg(test)] +const COMPRESSION_MIN_SIZE: usize = 64; +#[cfg(test)] +const COMPRESSION_THRESHOLD_PERCENT: usize = 85; +#[cfg(test)] +const COMPRESSION_FLAG_NONE: u8 = 0x00; +pub const COMPRESSION_FLAG_LZ4: u8 = 0x01; + +#[cfg(test)] +const MAGIC_METADATA_TARGET_ID_MAPPING: u8 = 0x01; + +type TraversalStack = SmallVec<[(u64, usize); 8]>; + +/* + B+Tree File Layout + ================== + + ┌─────────────────────────────────────────────────────────────┐ + │ File Header (PAGE_SIZE bytes, currently 4096) │ + ├─────────────────────────────────────────────────────────────┤ + │ MAGIC [4B: "BTRE"] │ + │ VERSION [4B: u32] │ + │ ROOT_OFFSET [8B: u64] │ + │ METADATA_LEN_FLAGS [4B: u32] │ + │ bit31: metadata flags initialized │ + │ bit30: has_tombstones │ + │ bits0..29: metadata length │ + │ METADATA [variable, up to 4000B] │ + │ [padding to PAGE_SIZE] │ + └─────────────────────────────────────────────────────────────┘ + + Leaf Node Layout (single or multi-block) + ┌─────────────────────────────────────────────────────────────┐ + │ IS_LEAF [1B: 0x01] │ + │ KEYS_LEN [4B: u32] │ + │ KEYS [MessagePack serialized Vec] │ + │ VALUE_INFO_LEN [4B: u32] │ + │ VALUE_INFO [MessagePack serialized Vec] │ + │ [padding to block boundary] │ + └─────────────────────────────────────────────────────────────┘ + + Internal Node Layout (supports multi-block when content exceeds PAGE_SIZE) + ┌─────────────────────────────────────────────────────────────┐ + │ IS_LEAF [1B: 0x00] │ + │ KEYS_LEN [4B: u32] │ + │ KEYS [MessagePack serialized Vec] │ + │ POINTERS_LEN [4B: u32] │ + │ POINTERS [MessagePack serialized Vec] │ + │ [padding to block boundary] │ + └─────────────────────────────────────────────────────────────┘ + + Note: Internal nodes can span multiple PAGE_SIZE blocks when + keys + pointers exceed a single page. The order calculation + uses a 75% safety factor to minimize multi-block nodes. + + Value Storage Modes: + - Single: Large values stored at [offset] with optional LZ4 compression + Format: [FLAG:1B][payload...] where FLAG = 0x00 (raw) or 0x01 (LZ4) + - Packed: Small values (≤256B) packed into PAGE_SIZE blocks + Format: [COUNT:4B][LEN:4B][data...][LEN:4B][data...]... +*/ + +#[inline] +fn u32_from_bytes(bytes: &[u8]) -> io::Result { Ok(u32::from_le_bytes(bytes.try_into().map_err(to_io_error)?)) } + +#[inline] +fn node_flag_to_is_leaf(flag: u8) -> io::Result { + match flag { + 0 => Ok(false), + 1 => Ok(true), + _ => Err(io::Error::new(io::ErrorKind::InvalidData, format!("Invalid B+Tree node flag: {flag}"))), + } +} + +#[inline] +fn checked_slice_range(start: usize, len: usize, total_len: usize) -> io::Result> { + let end = start + .checked_add(len) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "B+Tree node slice offset overflow"))?; + if end > total_len { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "B+Tree node slice out of bounds")); + } + Ok(start..end) +} + +fn valid_internal_pointer_count(key_count: usize, pointer_count: usize) -> bool { + pointer_count != 0 + && (key_count.checked_add(1) == Some(pointer_count) || (key_count != 0 && key_count == pointer_count)) +} + +#[inline] +fn get_entry_index_upper_bound(keys: &[K], key: &K) -> usize +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, +{ + let mut left = 0; + let mut right = keys.len(); + while left < right { + let mid = left + ((right - left) >> 1); + if &keys[mid] <= key { + left = mid + 1; + } else { + right = mid; + } + } + left +} + +#[inline] +#[cfg(test)] +const fn msgpack_array_header_len(count: usize) -> usize { + if count <= 0x0f { + 1 + } else if count <= u16::MAX as usize { + 3 + } else { + 5 + } +} + +#[inline] +#[cfg(test)] +const fn msgpack_u64_array_upper_bound_len(count: usize) -> usize { + // Worst-case per u64: marker + 8 bytes payload. + msgpack_array_header_len(count) + count.saturating_mul(9) +} + +// Adaptively compress value bytes if beneficial. +// Returns borrowed raw bytes when compression is not useful to avoid an +// allocation on the common uncompressed path. +#[cfg(test)] +fn compress_if_beneficial(raw_bytes: &[u8]) -> (u8, Cow<'_, [u8]>) { + if raw_bytes.len() >= COMPRESSION_MIN_SIZE { + let compressed = lz4_flex::compress_prepend_size(raw_bytes); + let threshold = (raw_bytes.len() * COMPRESSION_THRESHOLD_PERCENT) / 100; + + if compressed.len() < threshold { + // Compression is effective + (COMPRESSION_FLAG_LZ4, Cow::Owned(compressed)) + } else { + // Compression not worth it - return borrowed raw bytes. + (COMPRESSION_FLAG_NONE, Cow::Borrowed(raw_bytes)) + } + } else { + // Too small to compress - return borrowed raw bytes. + (COMPRESSION_FLAG_NONE, Cow::Borrowed(raw_bytes)) + } +} + +/// Represents how a value is stored on disk +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +enum ValueStorageMode { + /// Multiple small values packed in one block + /// (`block_offset`, `value_index_in_block`) + Packed(u64, u16), + + /// Single value in dedicated block(s) + /// (`block_offset`) + Single(u64), + + /// Entry is logically deleted. + Tombstone, +} + +#[derive(Debug, Clone)] +enum CacheData { + Compressed(u8, Vec), +} + +/// Extended value info that includes storage mode and length +#[derive(Debug, Serialize, Deserialize)] +struct ValueInfo { + mode: ValueStorageMode, + length: u32, + #[serde(skip, default)] + cache: Mutex>, +} + +impl ValueInfo { + #[inline] + const fn is_tombstone(&self) -> bool { matches!(self.mode, ValueStorageMode::Tombstone) } +} + +impl Clone for ValueInfo { + fn clone(&self) -> Self { + Self { + mode: self.mode, + length: self.length, + cache: Mutex::new(None), // Don't clone cache + } + } +} + +#[derive(Debug, Clone)] +struct BPlusTreeNode { + keys: Vec, + #[cfg_attr(not(test), allow(dead_code))] + children: Vec>, + is_leaf: bool, + value_info: Vec, + #[cfg_attr(not(test), allow(dead_code))] + values: Vec, // only used in leaf nodes +} + +impl BPlusTreeNode +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + #[inline] + #[cfg(test)] + const fn new(is_leaf: bool) -> Self { + Self { is_leaf, keys: vec![], children: vec![], value_info: vec![], values: vec![] } + } + + #[inline] + #[cfg(test)] + fn is_overflow(&self, order: usize) -> bool { self.keys.len() > order } + + #[inline] + #[cfg(test)] + const fn get_median_index(order: usize) -> usize { order >> 1 } + + #[cfg(test)] + fn find_leaf_entry(node: &Self) -> Option<&K> { + if node.is_leaf { + node.keys.first() + } else if let Some(child) = node.children.first() { + Self::find_leaf_entry(child) + } else { + None + } + } + + #[cfg(test)] + fn get_entry_index_upper_bound(&self, key: &K) -> usize { get_entry_index_upper_bound::(&self.keys, key) } + + #[cfg(test)] + fn insert(&mut self, key: K, v: V, inner_order: usize, leaf_order: usize) -> Option { + if self.is_leaf { + // Use single binary search instead of redundant searches + match self.keys.binary_search(&key) { + Ok(pos) => { + // Key exists, update value + self.values[pos] = v; + return None; + } + Err(pos) => { + // Key doesn't exist, insert at the correct position + self.keys.insert(pos, key); + self.values.insert(pos, v); + if self.is_overflow(leaf_order) { + return Some(self.split(leaf_order)); + } + } + } + } else { + let pos = self.get_entry_index_upper_bound(&key); + let child = self.children.get_mut(pos)?; + let node = child.insert(key.clone(), v, inner_order, leaf_order); + if let Some(tree_node) = node { + if let Some(leaf_key) = Self::find_leaf_entry(&tree_node) { + let idx = self.get_entry_index_upper_bound(leaf_key); + if self.keys.binary_search(leaf_key).is_err() { + self.keys.insert(idx, leaf_key.clone()); + self.children.insert(idx + 1, tree_node); + if self.is_overflow(inner_order) { + return Some(self.split(inner_order)); + } + } + } + } + } + None + } + + #[cfg(test)] + fn split(&mut self, order: usize) -> Self { + let median = Self::get_median_index(order); + if self.is_leaf { + let mut node = Self::new(true); + node.keys = self.keys.split_off(median); + node.values = self.values.split_off(median); + node + } else { + let mut node = Self::new(false); + node.keys = self.keys.split_off(median + 1); + node.children = self.children.split_off(median + 1); + // Internal keys are separators for children[1..]. The median key + // separates the two split nodes and is represented in the parent + // by the first leaf key of the returned right node. + let _separator = self.keys.pop(); + node + } + } + + #[cfg(test)] + fn add_historical_fence_key(&mut self) -> bool { + if self.is_leaf { + return false; + } + for index in 0..self.children.len().saturating_sub(1) { + let Some(fence) = self.keys.get(index).cloned() else { + continue; + }; + let child = &mut self.children[index]; + if !child.is_leaf && child.children.len() == child.keys.len().saturating_add(1) { + child.keys.push(fence); + return true; + } + } + self.children.iter_mut().any(Self::add_historical_fence_key) + } + + /// Write a packed value block to disk + #[cfg(test)] + fn write_packed_block( + file: &mut W, + buffer: &mut [u8], + offset: u64, + values: &[(u16, &[u8])], + ) -> io::Result<()> { + file.seek(SeekFrom::Start(offset))?; + + // Write count + let count = u32::try_from(values.len()).map_err(to_io_error)?; + buffer[0..4].copy_from_slice(&count.to_le_bytes()); + let mut pos = 4; + + // Write each value: length + data + for (_, value_bytes) in values { + let len = u32::try_from(value_bytes.len()).map_err(to_io_error)?; + buffer[pos..pos + 4].copy_from_slice(&len.to_le_bytes()); + pos += 4; + buffer[pos..pos + value_bytes.len()].copy_from_slice(value_bytes); + pos += value_bytes.len(); + } + + // Zero remaining space + if pos < PAGE_SIZE_USIZE { + buffer[pos..PAGE_SIZE_USIZE].fill(0u8); + } + + file.write_all(&buffer[..PAGE_SIZE_USIZE])?; + Ok(()) + } + + /// Calculate the serialized size of this node in bytes (rounded up to block size) + #[cfg(test)] + fn calculate_serialized_size(&self, serial_buf: &mut Vec) -> io::Result { + serial_buf.clear(); + + // Header: is_leaf flag + let mut size = FLAG_SIZE; + + // Keys: length + serialized data + binary_serialize_into(&mut *serial_buf, &self.keys)?; + size += LEN_SIZE + serial_buf.len(); + + if self.is_leaf { + // Leaf nodes now store value_info instead of values + // value_info: length + Vec<(u64, u32)> + // Reuse buf + serial_buf.clear(); + binary_serialize_into(&mut *serial_buf, &self.value_info)?; + size += LEN_SIZE + serial_buf.len(); + } else { + // Internal node: pointer length + pointers + // Pointer encoding is variable-length. Using small placeholder values + // can underestimate node size and cause offset overlap. + size += LEN_SIZE + msgpack_u64_array_upper_bound_len(self.children.len()); + } + + // Round up to block size + let blocks = size.div_ceil(PAGE_SIZE_USIZE); + Ok((blocks * PAGE_SIZE_USIZE) as u64) + } + + #[cfg(test)] + fn serialize_to_block( + &self, + file: &mut W, + buffer: &mut Vec, + serial_buf: &mut Vec, + offset: u64, + ) -> io::Result { + serial_buf.clear(); + binary_serialize_into(&mut *serial_buf, &self.keys)?; + let keys_len = u32::try_from(serial_buf.len()).map_err(to_io_error)?; + + if self.is_leaf { + let keys_end = serial_buf.len(); + // Append info_encoded to serial_buf to avoid second allocation + binary_serialize_into(&mut *serial_buf, &self.value_info)?; + let info_len = u32::try_from(serial_buf.len() - keys_end).map_err(to_io_error)?; + let info_slice = &serial_buf[keys_end..]; + + let content_size = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + info_len as usize; + let blocks = content_size.div_ceil(PAGE_SIZE_USIZE); + let capacity = blocks * PAGE_SIZE_USIZE; + debug_assert!( + content_size <= capacity, + "Leaf node content ({content_size}B) exceeds allocated capacity ({capacity}B)" + ); + + file.seek(SeekFrom::Start(offset))?; + + if buffer.len() < capacity { + buffer.resize(capacity, 0); + } + buffer[..capacity].fill(0); + + let mut pos = 0; + buffer[pos] = 1u8; + pos += FLAG_SIZE; + + buffer[pos..pos + LEN_SIZE].copy_from_slice(&keys_len.to_le_bytes()); + pos += LEN_SIZE; + + buffer[pos..pos + keys_len as usize].copy_from_slice(&serial_buf[0..keys_len as usize]); + pos += keys_len as usize; + + buffer[pos..pos + LEN_SIZE].copy_from_slice(&info_len.to_le_bytes()); + pos += LEN_SIZE; + + buffer[pos..pos + info_len as usize].copy_from_slice(info_slice); + + file.write_all(&buffer[..capacity])?; + + Ok(offset + (blocks as u64 * PAGE_SIZE_USIZE as u64)) + } else { + let ptr_count = self.children.len(); + // Conservative upper bound for MessagePack-encoded Vec. + // Must not underestimate, otherwise child blocks can overlap. + let ptr_encoded_size = msgpack_u64_array_upper_bound_len(ptr_count); + + let content_size = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + ptr_encoded_size; + let blocks_needed = content_size.div_ceil(PAGE_SIZE_USIZE); + + let parent_start = offset; + let mut current_offset = parent_start + (blocks_needed as u64 * PAGE_SIZE_USIZE as u64); + + let mut pointers = Vec::with_capacity(ptr_count); + for child in &self.children { + pointers.push(current_offset); + let mut child_scratch = Vec::new(); // Separate scratch for recursion to protect our serial_buf + current_offset = child.serialize_to_block(file, buffer, &mut child_scratch, current_offset)?; + } + + // Append pointers to serial_buf + let keys_end = serial_buf.len(); + binary_serialize_into(&mut *serial_buf, &pointers)?; + let pointers_len = u32::try_from(serial_buf.len() - keys_end).map_err(to_io_error)?; + let pointers_slice = &serial_buf[keys_end..]; + + file.seek(SeekFrom::Start(parent_start))?; + + let total_capacity = blocks_needed * PAGE_SIZE_USIZE; + let actual_content = FLAG_SIZE + LEN_SIZE + keys_len as usize + LEN_SIZE + pointers_len as usize; + debug_assert!( + actual_content <= total_capacity, + "Internal node content ({actual_content}B) exceeds allocated capacity ({total_capacity}B)" + ); + if buffer.len() < total_capacity { + buffer.resize(total_capacity, 0); + } + buffer[..total_capacity].fill(0); + + let mut pos = 0; + // Is_leaf=0 + buffer[pos] = 0u8; + pos += FLAG_SIZE; + + buffer[pos..pos + LEN_SIZE].copy_from_slice(&keys_len.to_le_bytes()); + pos += LEN_SIZE; + buffer[pos..pos + keys_len as usize].copy_from_slice(&serial_buf[0..keys_len as usize]); + pos += keys_len as usize; + + buffer[pos..pos + LEN_SIZE].copy_from_slice(&pointers_len.to_le_bytes()); + pos += LEN_SIZE; + buffer[pos..pos + pointers_len as usize].copy_from_slice(pointers_slice); + + file.write_all(&buffer[..total_capacity])?; + + Ok(current_offset) + } + } + + /// Serialize the tree in breadth-first order for better disk locality + /// This improves query performance by keeping nodes at the same level contiguous + #[cfg(test)] + fn serialize_breadth_first( + &mut self, + file: &mut W, + buffer: &mut Vec, + start_offset: u64, + ) -> io::Result { + let mut serial_buf = Vec::with_capacity(PAGE_SIZE_USIZE); + + self.serialize_bfs_pass1_populate_value_info(&mut serial_buf)?; + let (node_offsets, child_ids_by_node, current_offset) = + self.serialize_bfs_pass2_calculate_offsets(&mut serial_buf, start_offset)?; + self.serialize_bfs_pass3_assign_value_offsets(current_offset); + self.serialize_bfs_pass4_write_nodes(file, buffer, &mut serial_buf, &node_offsets, &child_ids_by_node)?; + self.serialize_bfs_pass5_write_values(file, buffer, &mut serial_buf)?; + + Ok(start_offset) + } + + #[cfg(test)] + fn serialize_bfs_pass1_populate_value_info(&mut self, serial_buf: &mut Vec) -> io::Result<()> { + let mut current_level_mut = vec![self]; + while !current_level_mut.is_empty() { + let mut next_level_mut = Vec::new(); + for node in current_level_mut { + if node.is_leaf { + node.value_info.clear(); + let mut serialized_values: Vec> = Vec::new(); + for value in &node.values { + serial_buf.clear(); + binary_serialize_into(serial_buf, value)?; + serialized_values.push(serial_buf.clone()); + } + + let mut current_pack_index: u16 = 0; + let mut current_pack_size = PACK_BLOCK_HEADER_SIZE; + let mut pack_count = 0u32; + + for value_bytes in serialized_values { + let size = value_bytes.len(); + + if size <= SMALL_VALUE_THRESHOLD { + let entry_size = PACK_VALUE_HEADER_SIZE + size; + + if current_pack_size + entry_size <= PAGE_SIZE_USIZE { + node.value_info.push(ValueInfo { + mode: ValueStorageMode::Packed(u64::from(pack_count), current_pack_index), + length: u32::try_from(size).map_err(to_io_error)?, + cache: Mutex::new(None), + }); + current_pack_index += 1; + current_pack_size += entry_size; + } else { + pack_count += 1; + current_pack_index = 1; + current_pack_size = PACK_BLOCK_HEADER_SIZE + entry_size; + + node.value_info.push(ValueInfo { + mode: ValueStorageMode::Packed(u64::from(pack_count), 0), + length: u32::try_from(size).map_err(to_io_error)?, + cache: Mutex::new(None), + }); + } + } else { + let (flag, payload) = compress_if_beneficial(&value_bytes); + let stored_size = 1 + payload.len(); + + let cache = if flag == COMPRESSION_FLAG_LZ4 { + Some(CacheData::Compressed(flag, payload.into_owned())) + } else { + None + }; + + node.value_info.push(ValueInfo { + mode: ValueStorageMode::Single(u64::MAX), + length: u32::try_from(stored_size).map_err(to_io_error)?, + cache: Mutex::new(cache), + }); + } + } + } else { + for child in &mut node.children { + next_level_mut.push(child); + } + } + } + current_level_mut = next_level_mut; + } + Ok(()) + } + + #[cfg(test)] + fn serialize_bfs_pass2_calculate_offsets( + &self, + serial_buf: &mut Vec, + start_offset: u64, + ) -> io::Result<(Vec, Vec>, u64)> { + let mut node_refs: Vec<&BPlusTreeNode> = vec![self]; + let mut node_offsets: Vec = vec![start_offset]; + let mut child_ids_by_node: Vec> = vec![Vec::new()]; + let mut current_offset = start_offset + self.calculate_serialized_size(serial_buf)?; + let mut current_level = vec![0usize]; + + while !current_level.is_empty() { + let mut next_level = Vec::new(); + for node_id in current_level { + let node = node_refs[node_id]; + if !node.is_leaf { + for child in &node.children { + let child_id = node_refs.len(); + node_refs.push(child); + node_offsets.push(current_offset); + child_ids_by_node.push(Vec::new()); + child_ids_by_node[node_id].push(child_id); + current_offset += child.calculate_serialized_size(serial_buf)?; + next_level.push(child_id); + } + } + } + current_level = next_level; + } + + Ok((node_offsets, child_ids_by_node, current_offset)) + } + + #[cfg(test)] + fn serialize_bfs_pass3_assign_value_offsets(&mut self, mut current_offset: u64) { + use std::collections::HashMap; + let mut current_level_mut = vec![self]; + while !current_level_mut.is_empty() { + let mut next_level_mut = Vec::new(); + for node in current_level_mut { + if node.is_leaf { + let mut pack_block_offsets: HashMap = HashMap::new(); + + for info in &mut node.value_info { + match &mut info.mode { + ValueStorageMode::Packed(pack_idx, _index) => { + if !pack_block_offsets.contains_key(pack_idx) { + pack_block_offsets.insert(*pack_idx, current_offset); + current_offset += PAGE_SIZE_USIZE as u64; + } + } + ValueStorageMode::Single(offset) if *offset == u64::MAX => { + *offset = current_offset; + current_offset += u64::from(info.length); + } + ValueStorageMode::Single(_) | ValueStorageMode::Tombstone => {} + } + } + + for info in &mut node.value_info { + if let ValueStorageMode::Packed(pack_idx, _index) = &mut info.mode { + let actual_offset = pack_block_offsets[pack_idx]; + *pack_idx = actual_offset; + } + } + } else { + for child in &mut node.children { + next_level_mut.push(child); + } + } + } + current_level_mut = next_level_mut; + } + } + + #[cfg(test)] + fn serialize_bfs_pass4_write_nodes( + &self, + file: &mut W, + buffer: &mut Vec, + serial_buf: &mut Vec, + node_offsets: &[u64], + child_ids_by_node: &[Vec], + ) -> io::Result<()> { + let mut node_refs: Vec<&BPlusTreeNode> = vec![self]; + let mut node_cursor = 0; + while node_cursor < node_refs.len() { + let node = node_refs[node_cursor]; + if !node.is_leaf { + node_refs.extend(node.children.iter()); + } + node_cursor += 1; + } + + if node_refs.len() != node_offsets.len() || node_refs.len() != child_ids_by_node.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "B+Tree serialization produced inconsistent node offset table", + )); + } + + for (node_id, node) in node_refs.iter().enumerate() { + let node_offset = node_offsets[node_id]; + + if node.is_leaf { + node.serialize_to_block(file, buffer, serial_buf, node_offset)?; + } else { + let node_child_ids = child_ids_by_node.get(node_id).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "B+Tree serialization missing child id table entry", + ) + })?; + let mut child_offsets = Vec::with_capacity(node_child_ids.len()); + for child_id in node_child_ids { + let Some(child_offset) = node_offsets.get(*child_id) else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "B+Tree serialization child id has no offset", + )); + }; + child_offsets.push(*child_offset); + } + + node.serialize_internal_with_offsets( + file, + buffer, + serial_buf, + node_offset, + &child_offsets, + )?; + } + } + Ok(()) + } + + #[cfg(test)] + fn serialize_bfs_pass5_write_values( + &self, + file: &mut W, + buffer: &mut [u8], + serial_buf: &mut Vec, + ) -> io::Result<()> { + use std::collections::HashMap; + let mut current_level_values = vec![self]; + while !current_level_values.is_empty() { + let mut next_level = Vec::new(); + for node in current_level_values { + if node.is_leaf { + let mut pack_blocks: HashMap)>> = HashMap::new(); + + for (value, info) in node.values.iter().zip(node.value_info.iter()) { + serial_buf.clear(); + binary_serialize_into(serial_buf, value)?; + + match info.mode { + ValueStorageMode::Packed(block_offset, index) => { + pack_blocks.entry(block_offset).or_default().push((index, serial_buf.clone())); + } + ValueStorageMode::Single(block_offset) => { + file.seek(SeekFrom::Start(block_offset))?; + + let cache_guard = info.cache.lock(); + let (flag, payload_ref) = + if let Some(CacheData::Compressed(c_flag, c_payload)) = cache_guard.as_ref() { + (*c_flag, c_payload.as_slice()) + } else { + (COMPRESSION_FLAG_NONE, serial_buf.as_slice()) + }; + + file.write_all(&[flag])?; + file.write_all(payload_ref)?; + } + ValueStorageMode::Tombstone => {} + } + } + + for (block_offset, mut values) in pack_blocks { + values.sort_by_key(|(idx, _)| *idx); + let value_refs: Vec<(u16, &[u8])> = + values.iter().map(|(idx, bytes)| (*idx, bytes.as_slice())).collect(); + Self::write_packed_block(file, buffer, block_offset, &value_refs)?; + } + } else { + for child in &node.children { + next_level.push(child); + } + } + } + current_level_values = next_level; + } + Ok(()) + } + + /// Serialize an internal node with pre-calculated child offsets + /// Supports multi-block internal nodes when keys + pointers exceed a single page + #[cfg(test)] + fn serialize_internal_with_offsets( + &self, + file: &mut W, + buffer: &mut Vec, + serial_buf: &mut Vec, + offset: u64, + child_offsets: &[u64], + ) -> io::Result { + // Similar to serialize_to_block but for internal nodes with known child offsets + serial_buf.clear(); + binary_serialize_into(&mut *serial_buf, &self.keys)?; + let keys_len = serial_buf.len(); + let keys_end = keys_len; + + binary_serialize_into(&mut *serial_buf, child_offsets)?; + let pointer_len = serial_buf.len() - keys_end; + + // Calculate total content size + let total_content_size = FLAG_SIZE + LEN_SIZE + keys_len + LEN_SIZE + pointer_len; + let blocks_needed = total_content_size.div_ceil(PAGE_SIZE_USIZE); + + let total_buffer_size = blocks_needed * PAGE_SIZE_USIZE; + if buffer.len() < total_buffer_size { + buffer.resize(total_buffer_size, 0); + } + buffer[..total_buffer_size].fill(0); + + let mut write_pos = 0; + + // Write is_leaf flag (0 for internal node) + buffer[write_pos] = u8::from(self.is_leaf); + write_pos += FLAG_SIZE; + + // Write keys length and data + buffer[write_pos..write_pos + LEN_SIZE] + .copy_from_slice(&u32::try_from(keys_len).map_err(to_io_error)?.to_le_bytes()); + write_pos += LEN_SIZE; + buffer[write_pos..write_pos + keys_len].copy_from_slice(&serial_buf[0..keys_end]); + write_pos += keys_len; + + // Write pointers length and data + buffer[write_pos..write_pos + LEN_SIZE] + .copy_from_slice(&u32::try_from(pointer_len).map_err(to_io_error)?.to_le_bytes()); + write_pos += LEN_SIZE; + buffer[write_pos..write_pos + pointer_len].copy_from_slice(&serial_buf[keys_end..]); + + // Write all blocks to file + file.seek(SeekFrom::Start(offset))?; + file.write_all(&buffer[..total_buffer_size])?; + + Ok(offset + total_buffer_size as u64) + } + + fn deserialize_from_block( + file: &mut R, + buffer: &mut Vec, + offset: u64, + nested: bool, + ) -> io::Result<(Self, Option>)> { + file.seek(SeekFrom::Start(offset))?; + + let header_required = FLAG_SIZE + LEN_SIZE; + if buffer.len() < header_required { + buffer.resize(header_required, 0); + } + + file.read_exact(&mut buffer[0..header_required])?; + + let is_leaf = node_flag_to_is_leaf(buffer[0])?; + #[allow(clippy::range_plus_one)] + let keys_len = u32_from_bytes(&buffer[FLAG_SIZE..FLAG_SIZE + LEN_SIZE])? as usize; + + let min_required = header_required + keys_len + LEN_SIZE; + if buffer.len() < min_required { + buffer.resize(min_required, 0); + } + + file.read_exact(&mut buffer[header_required..min_required])?; + + let mut read_pos = header_required; + let mut keys: Vec = binary_deserialize(&buffer[read_pos..read_pos + keys_len])?; + read_pos += keys_len; + + let payload_len = u32_from_bytes(&buffer[read_pos..read_pos + LEN_SIZE])? as usize; + read_pos += LEN_SIZE; + + let total_required = min_required + payload_len; + if buffer.len() < total_required { + buffer.resize(total_required, 0); + } + + file.read_exact(&mut buffer[min_required..total_required])?; + + let (value_info, values, children, children_pointer) = if is_leaf { + let mut info: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; + if info.len() != keys.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid leaf node: {} keys but {} value descriptors", keys.len(), info.len()), + )); + } + let vals = if nested { + let mut filtered_keys: Vec = Vec::with_capacity(keys.len()); + let mut filtered_info: Vec = Vec::with_capacity(info.len()); + let mut v = Vec::with_capacity(info.len()); + + let original_keys = std::mem::take(&mut keys); + for (entry_key, entry_info) in original_keys.into_iter().zip(info) { + if entry_info.is_tombstone() { + continue; + } + v.push(Self::load_value_from_info(file, &entry_info)?); + filtered_keys.push(entry_key); + filtered_info.push(entry_info); + } + + keys = filtered_keys; + info = filtered_info; + v + } else { + Vec::new() + }; + (info, vals, Vec::new(), None) + } else { + let pointers: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; + if !valid_internal_pointer_count(keys.len(), pointers.len()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid internal node: {} keys but {} child pointers", keys.len(), pointers.len()), + )); + } + let nodes = if nested { + let mut n = Vec::with_capacity(pointers.len()); + let mut child_buf = Vec::with_capacity(PAGE_SIZE_USIZE); + for &ptr in &pointers { + let (child, _) = Self::deserialize_from_block(file, &mut child_buf, ptr, nested)?; + n.push(child); + } + n + } else { + Vec::new() + }; + (Vec::new(), Vec::new(), nodes, Some(pointers)) + }; + + Ok((Self { keys, children, is_leaf, value_info, values }, children_pointer)) + } + + fn deserialize_from_mmap( + mmap: &[u8], + file: &mut R, + offset: u64, + nested: bool, + ) -> io::Result<(Self, Option>)> { + let start = usize::try_from(offset).map_err(to_io_error)?; + let header_end = start + .checked_add(FLAG_SIZE + LEN_SIZE) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; + // Basic safety check for mmap bounds + if header_end > mmap.len() { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); + } + + let keys_len = u32_from_bytes(&mmap[start + FLAG_SIZE..start + FLAG_SIZE + LEN_SIZE])? as usize; + let keys_start = header_end; + let len_pos = keys_start + .checked_add(keys_len) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; + + if len_pos + LEN_SIZE > mmap.len() { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); + } + let payload_len = u32_from_bytes(&mmap[len_pos..len_pos + LEN_SIZE])? as usize; + let total = len_pos + .checked_add(LEN_SIZE + payload_len) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap offset overflow"))?; + + if total > mmap.len() { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Mmap access out of bounds")); + } + + // We need to know the total size of the node to slice the mmap + // For simplicity, we can just slice a PAGE_SIZE or slightly more if we know it overflows. + // Actually, our serialize_to_block uses PAGE_SIZE blocks. + + //let slice = &mmap[start..]; + let slice = &mmap[start..total]; + Self::deserialize_from_block_slice(slice, Some(mmap), file, nested) + } + + fn deserialize_from_block_slice( + slice: &[u8], + mmap: Option<&[u8]>, + file: &mut R, + nested: bool, + ) -> io::Result<(Self, Option>)> { + let header_range = checked_slice_range(0, FLAG_SIZE + LEN_SIZE, slice.len())?; + // Node type + let is_leaf = node_flag_to_is_leaf(slice[0])?; + let mut read_pos = FLAG_SIZE; + + // ---- Keys ---- + let keys_length = u32_from_bytes(&slice[read_pos..header_range.end])? as usize; + read_pos += LEN_SIZE; + let keys_range = checked_slice_range(read_pos, keys_length, slice.len())?; + let mut keys: Vec = binary_deserialize(&slice[keys_range.clone()])?; + read_pos = keys_range.end; + + // ---- Value info (offset, length) for leaf nodes ---- + let (value_info, values): (Vec, Vec) = if is_leaf { + // Read value_info + let info_len_range = checked_slice_range(read_pos, LEN_SIZE, slice.len())?; + let info_length = u32_from_bytes(&slice[info_len_range.clone()])? as usize; + read_pos = info_len_range.end; + let info_range = checked_slice_range(read_pos, info_length, slice.len())?; + let mut info: Vec = binary_deserialize(&slice[info_range])?; + if info.len() != keys.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid leaf node: {} keys but {} value descriptors", keys.len(), info.len()), + )); + } + + // Values are loaded on-demand when nested=true + if nested { + let mut vals = Vec::with_capacity(info.len()); + let mut filtered_keys: Vec = Vec::with_capacity(keys.len()); + let mut filtered_info: Vec = Vec::with_capacity(info.len()); + let mut last_packed_block: Option<(u64, Vec)> = None; + let original_keys = std::mem::take(&mut keys); + for (entry_key, entry_info) in original_keys.into_iter().zip(info) { + if entry_info.is_tombstone() { + continue; + } + match entry_info.mode { + ValueStorageMode::Packed(block_offset, index) => { + // Packed loading optimization: reuse block if it's the same + if let Some((offset, ref block)) = last_packed_block { + if offset == block_offset { + vals.push(Self::extract_value_from_packed_block(block, index)?); + filtered_keys.push(entry_key); + filtered_info.push(entry_info); + continue; + } + } + + // Load new block + let mut block = vec![0u8; PAGE_SIZE_USIZE]; + file.seek(SeekFrom::Start(block_offset))?; + file.read_exact(&mut block)?; + vals.push(Self::extract_value_from_packed_block(&block, index)?); + last_packed_block = Some((block_offset, block)); + filtered_keys.push(entry_key); + filtered_info.push(entry_info); + } + ValueStorageMode::Single(_) => { + last_packed_block = None; + vals.push(Self::load_value_from_info(file, &entry_info)?); + filtered_keys.push(entry_key); + filtered_info.push(entry_info); + } + ValueStorageMode::Tombstone => {} + } + } + keys = filtered_keys; + info = filtered_info; + (info, vals) + } else { + (info, Vec::new()) + } + } else { + (Vec::new(), Vec::new()) + }; + + // ---- Pointers for internal nodes ---- + let (children, children_pointer): (Vec, Option>) = if is_leaf { + (Vec::new(), None) + } else { + let pointers_len_range = checked_slice_range(read_pos, LEN_SIZE, slice.len())?; + let pointers_length = u32_from_bytes(&slice[pointers_len_range.clone()])? as usize; + read_pos = pointers_len_range.end; + let pointers_range = checked_slice_range(read_pos, pointers_length, slice.len())?; + let pointers: Vec = binary_deserialize(&slice[pointers_range])?; + if !valid_internal_pointer_count(keys.len(), pointers.len()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid internal node: {} keys but {} child pointers", keys.len(), pointers.len()), + )); + } + if nested { + let mut nodes = Vec::with_capacity(pointers.len()); + let mut child_buffer = vec![0u8; PAGE_SIZE_USIZE]; + for &ptr in &pointers { + let (child, _) = if let Some(m) = mmap { + Self::deserialize_from_mmap(m, file, ptr, nested)? + } else { + Self::deserialize_from_block(file, &mut child_buffer, ptr, nested)? + }; + nodes.push(child); + } + (nodes, None) + } else { + (Vec::new(), Some(pointers)) + } + }; + + Ok((Self { keys, children, is_leaf, value_info, values }, children_pointer)) + } + + /// Load a value based on its storage info + fn load_value_from_info(file: &mut R, info: &ValueInfo) -> io::Result { + // Fast path: Check cache for Single mode + if let ValueStorageMode::Single(_) = info.mode { + let cache_guard = info.cache.lock(); + if let Some(CacheData::Compressed(flag, payload)) = cache_guard.as_ref() { + if *flag == COMPRESSION_FLAG_LZ4 { + let decompressed = lz4_flex::decompress_size_prepended(payload).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 cache decompression failed: {e}")) + })?; + return binary_deserialize(&decompressed); + } + return binary_deserialize(payload); + } + } + + match info.mode { + ValueStorageMode::Single(offset) => { + let stored_len = info.length as usize; + if stored_len < 1 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid value length")); + } + + // Read everything: flag + payload + file.seek(SeekFrom::Start(offset))?; + let mut buffer = vec![0u8; stored_len]; + file.read_exact(&mut buffer)?; + + let flag = buffer[0]; + // Split payload without re-allocating if possible? Vec::split_off allocates new vec for tail. + // We want payload as Vec for cache. + let payload = buffer[1..].to_vec(); + + // Decompress for result + let data = if flag == COMPRESSION_FLAG_LZ4 { + lz4_flex::decompress_size_prepended(&payload).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 decompression failed: {e}")) + })? + } else { + payload.clone() + }; + + // Update cache + *info.cache.lock() = Some(CacheData::Compressed(flag, payload)); + + binary_deserialize(&data) + } + ValueStorageMode::Packed(block_offset, index) => { + Self::load_value_from_packed_block(file, block_offset, index, info.length) + } + ValueStorageMode::Tombstone => { + Err(io::Error::new(io::ErrorKind::NotFound, "value was deleted (tombstone)")) + } + } + } + + /// Load a value from a packed block + fn load_value_from_packed_block( + file: &mut R, + block_offset: u64, + value_index: u16, + _expected_length: u32, + ) -> io::Result { + file.seek(SeekFrom::Start(block_offset))?; + + let mut block_buffer = vec![0u8; PAGE_SIZE_USIZE]; + file.read_exact(&mut block_buffer)?; + + Self::extract_value_from_packed_block(&block_buffer, value_index) + } + + /// Helper to extract value from a packed block that is already in memory + fn extract_value_from_packed_block(block_buffer: &[u8], value_index: u16) -> io::Result { + // Read count + if block_buffer.len() < 4 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block too small")); + } + let mut pos = 4; + + // Skip to target value + for i in 0..=value_index { + if pos + 4 > block_buffer.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Packed block corrupted: position {pos} exceeds block size"), + )); + } + + let len = u32::from_le_bytes(block_buffer[pos..pos + 4].try_into().map_err(to_io_error)?) as usize; + pos += 4; + + if i == value_index { + // Found target value + if pos + len > block_buffer.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Packed value corrupted: length {len} at position {pos} exceeds block size"), + )); + } + + let value_data = &block_buffer[pos..pos + len]; + return binary_deserialize(value_data); + } + + pos += len; + } + + Err(io::Error::new(io::ErrorKind::InvalidData, format!("Value index {value_index} not found in packed block"))) + } +} + +// ----------------------------------------------------------------------------- +// Metadata Enum +// ----------------------------------------------------------------------------- +#[cfg(test)] +#[derive(Clone, Debug, PartialEq)] +pub enum BPlusTreeMetadata { + Empty, + TargetIdMapping(u32), +} + +#[cfg(test)] +impl BPlusTreeMetadata { + pub fn to_bytes(&self) -> Vec { + match self { + Self::Empty => Vec::new(), + Self::TargetIdMapping(val) => { + let mut bytes = vec![MAGIC_METADATA_TARGET_ID_MAPPING]; // Type tag + bytes.extend_from_slice(&val.to_le_bytes()); + bytes + } + } + } + +} + +#[cfg(test)] +#[derive(Debug, Clone)] +pub struct BPlusTree { + root: BPlusTreeNode, + inner_order: usize, + leaf_order: usize, + metadata: BPlusTreeMetadata, + dirty: bool, +} + +#[cfg(test)] +const fn sanitize_order(order: usize) -> usize { + if order < 2 { + 2 + } else { + order + } +} + +#[cfg(test)] +const fn default_orders() -> (usize, usize) { (DEFAULT_INNER_ORDER, DEFAULT_LEAF_ORDER) } + +#[cfg(test)] +impl Default for BPlusTree +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + fn default() -> Self { Self::new() } +} + +#[cfg(test)] +impl BPlusTree +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + pub const fn new() -> Self { + let (inner_order, leaf_order) = default_orders(); + Self::new_with_orders(inner_order, leaf_order) + } + + /// Create a v2 tree with explicit in-memory fanout. + /// + /// This does not change the on-disk v2 format. Orders below 2 are clamped + /// because B+Tree split logic requires at least two keys per node. + pub const fn new_with_orders(inner_order: usize, leaf_order: usize) -> Self { + Self { + root: BPlusTreeNode::::new(true), + inner_order: sanitize_order(inner_order), + leaf_order: sanitize_order(leaf_order), + metadata: BPlusTreeMetadata::Empty, + dirty: true, // an empty tree is stored! + } + } + + /// Helper to set metadata + pub fn set_metadata(&mut self, data: BPlusTreeMetadata) { + self.metadata = data; + self.dirty = true; + } + + pub fn insert(&mut self, key: K, value: V) { + self.dirty = true; + if self.root.keys.is_empty() { + self.root.keys.push(key); + self.root.values.push(value); + return; + } + + if let Some(node) = self.root.insert(key, value, self.inner_order, self.leaf_order) { + let child_key_opt = + if node.is_leaf { node.keys.first() } else { BPlusTreeNode::::find_leaf_entry(&node) }; + + if let Some(child_key) = child_key_opt { + let mut new_root = BPlusTreeNode::::new(false); + new_root.keys.push(child_key.clone()); + new_root.children.push(std::mem::replace(&mut self.root, BPlusTreeNode::new(true))); + new_root.children.push(node); + + self.root = new_root; + } else { + error!("Failed to insert child key"); + } + } + } + + pub(crate) fn add_historical_fence_key(&mut self) -> bool { self.root.add_historical_fence_key() } + + pub(crate) fn remove_last_root_child(&mut self) -> bool { + !self.root.is_leaf && self.root.children.pop().is_some() + } + + pub fn store(&mut self, filepath: &Path) -> io::Result { + if self.dirty { + self.store_internal(filepath) + } else { + Ok(0) + } + } + + /// Internal store without locking, used for compaction or initial save. + fn store_internal(&mut self, filepath: &Path) -> io::Result { + let tempfile = if let Some(parent_dir) = filepath.parent() { + if let Ok(file) = NamedTempFile::new_in(parent_dir) { + file + } else { + let temp_dir = tempfile::env::temp_dir(); + NamedTempFile::new_in(&temp_dir)? + } + } else { + let temp_dir = tempfile::env::temp_dir(); + NamedTempFile::new_in(&temp_dir)? + }; + let mut file = utils::file_writer(&tempfile); + let mut buffer = vec![0u8; PAGE_SIZE_USIZE]; + + // Write header block 0 + let mut header = [0u8; PAGE_SIZE_USIZE]; + header[0..4].copy_from_slice(MAGIC); + header[4..8].copy_from_slice(&STORAGE_VERSION.to_le_bytes()); + // Placeholder for root offset, will be updated after serialization + header[8..16].copy_from_slice(&HEADER_SIZE.to_le_bytes()); + + let meta_bytes = self.metadata.to_bytes(); + if meta_bytes.len() > METADATA_MAX_SIZE || METADATA_DATA_START_POS + meta_bytes.len() > PAGE_SIZE_USIZE { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "Metadata too large for header page")); + } + let metadata_len = + u32::try_from(meta_bytes.len()).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?; + let metadata_len_with_flags = encode_metadata_len_with_flags(metadata_len, false); + header[16..20].copy_from_slice(&metadata_len_with_flags.to_le_bytes()); + if !meta_bytes.is_empty() { + header[METADATA_DATA_START_POS..METADATA_DATA_START_POS + meta_bytes.len()].copy_from_slice(&meta_bytes); + } + + file.write_all(&header)?; + + // We need to ensure we pad to PAGE_SIZE before continuing + file.seek(SeekFrom::Start(HEADER_SIZE))?; + + // Use breadth-first serialization for better disk locality + match self.root.serialize_breadth_first(&mut file, &mut buffer, HEADER_SIZE) { + Ok(root_offset) => { + // Update root offset in header + file.seek(SeekFrom::Start(ROOT_OFFSET_POS))?; + file.write_all(&root_offset.to_le_bytes())?; + + file.flush()?; + drop(file); + if let Err(err) = utils::rename_or_copy(tempfile.path(), filepath, false) { + return Err(string_to_io_error(format!( + "Temp file rename/copy did not work {} {err}", + tempfile.path().to_string_lossy() + ))); + } + self.dirty = false; + Ok(root_offset) + } + Err(err) => Err(err), + } + } + +} + +fn validate_legacy_tree( + mmap: Option<&Mmap>, + file: &mut BufReader, + root_offset: u64, + file_len: u64, +) -> io::Result<()> +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + let mut pending = vec![(root_offset, None::, None::, true)]; + let mut visited = HashSet::new(); + while let Some((offset, lower, upper, is_root)) = pending.pop() { + if offset < HEADER_SIZE || offset >= file_len { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("legacy B+Tree child offset {offset} is outside file length {file_len}"), + )); + } + if !visited.insert(offset) { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy B+Tree contains a cycle or shared child")); + } + + let (node, pointers) = if let Some(mapped) = mmap { + let mut cursor = io::Cursor::new(mapped.as_ref()); + BPlusTreeNode::::deserialize_from_mmap(mapped, &mut cursor, offset, false)? + } else { + let mut buffer = Vec::with_capacity(PAGE_SIZE_USIZE); + BPlusTreeNode::::deserialize_from_block(file, &mut buffer, offset, false)? + }; + if node.keys.windows(2).any(|keys| keys[0] >= keys[1]) { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy B+Tree node keys are not strictly increasing")); + } + + if node.is_leaf { + if let Some(expected) = lower.as_ref() { + if node.keys.first() != Some(expected) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "legacy B+Tree separator differs from child minimum", + )); + } + } + if upper.as_ref().is_some_and(|bound| node.keys.last().is_some_and(|key| key >= bound)) { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy B+Tree leaf exceeds its upper bound")); + } + continue; + } + + let pointers = pointers.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "internal node has no pointers"))?; + let has_fence_key = pointers.len() == node.keys.len(); + if has_fence_key && (is_root || node.keys.last() != upper.as_ref()) { + return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid historical internal-node fence key")); + } + + for (index, child) in pointers.into_iter().enumerate().rev() { + let child_lower = if index == 0 { lower.clone() } else { node.keys.get(index - 1).cloned() }; + let child_upper = node.keys.get(index).cloned().or_else(|| upper.clone()); + pending.push((child, child_lower, child_upper, false)); + } + } + Ok(()) +} + + +/// `BPlusTreeQuery` performs on-disk queries without loading the entire tree into memory. +/// For frequent queries, consider using `BPlusTree::load()` instead, which loads the full tree into memory +/// at the cost of higher memory usage. +pub struct BPlusTreeQuery { + file: Option>, + mmap: Option, + has_tombstones: bool, + buffer: Vec, + root_offset: u64, + _marker_k: PhantomData, + _marker_v: PhantomData, +} + +impl BPlusTreeQuery +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + pub fn try_from_file(file: File) -> io::Result { + let metadata = file.metadata()?; + let file_len = metadata.len(); + if file_len < HEADER_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "File too small")); + } + + // Try Mmap + let mmap = mmap_with_advice(&file, Advice::Normal, "B+Tree query"); + + // Verify Header + let mut header = [0u8; METADATA_DATA_START_POS]; + read_exact_at_offset(&file, &mut header, 0)?; + + if &header[0..4] != MAGIC { + return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid magic number")); + } + let version = u32::from_le_bytes( + header[4..8].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid version slice"))?, + ); + if version != STORAGE_VERSION { + return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Unsupported storage version: {version}"))); + } + let root_offset = u64::from_le_bytes( + header[8..16] + .try_into() + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid root offset slice"))?, + ); + let metadata_len_raw = u32::from_le_bytes( + header[16..20] + .try_into() + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid metadata length slice"))?, + ); + let (metadata_len, has_tombstones) = decode_metadata_len_and_flags(metadata_len_raw); + if metadata_len as usize > METADATA_MAX_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Metadata too large: {metadata_len}"))); + } + + let mut validation_file = utils::file_reader(file.try_clone()?); + validate_legacy_tree::(mmap.as_ref(), &mut validation_file, root_offset, file_len)?; + + Ok(Self { + file: if mmap.is_some() { None } else { Some(utils::file_reader(file)) }, + mmap, + has_tombstones, + buffer: vec![0u8; PAGE_SIZE_USIZE], + root_offset, + _marker_k: PhantomData, + _marker_v: PhantomData, + }) + } + + pub fn try_new(filepath: &Path) -> io::Result { + Self::try_from_file(File::open(filepath)?) + } + +} + + +impl BPlusTreeQuery +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + /// Iterates over key-value pairs within a given range using `right_sibling` pointers. + /// + /// This is more efficient than iterating the full tree and filtering when you only + /// need a subset of keys. + /// + /// Tombstones are skipped automatically. + pub fn range_iter( + &mut self, + start: Bound<&K>, + end: Bound<&K>, + ) -> impl Iterator> + '_ { + let start_cloned = match start { + Bound::Included(k) => Bound::Included(k.clone()), + Bound::Excluded(k) => Bound::Excluded(k.clone()), + Bound::Unbounded => Bound::Unbounded, + }; + let end_cloned = match end { + Bound::Included(k) => Bound::Included(k.clone()), + Bound::Excluded(k) => Bound::Excluded(k.clone()), + Bound::Unbounded => Bound::Unbounded, + }; + RangeLeafIterator::new(self, start_cloned, end_cloned) + } + +} + +/// Range scan iterator that seeks into the tree and then walks in-order +/// without scanning from the root for every entry. +struct RangeLeafIterator<'a, K, V> { + tree: &'a mut BPlusTreeQuery, + start_bound: Bound, + end_bound: Bound, + stack: TraversalStack, + current_leaf: Option>, + leaf_idx: usize, + initialized: bool, + exhausted: bool, +} + +/// Shared range-iterator helpers retained by the frozen v2 migration reader. +macro_rules! impl_range_leaf_common { + ($tree_ty:ty, $lt:tt) => { + fn new(tree: &$lt mut $tree_ty, start: Bound, end: Bound) -> Self { + Self { + tree, + start_bound: start, + end_bound: end, + stack: smallvec![], + current_leaf: None, + leaf_idx: 0, + initialized: false, + exhausted: false, + } + } + + fn load_leaf_from_node(&mut self, node: BPlusTreeNode, start_idx: usize) { + self.current_leaf = Some(node); + self.leaf_idx = start_idx; + } + + fn key_past_end(&self, key: &K) -> bool { + match &self.end_bound { + Bound::Included(end) => key > end, + Bound::Excluded(end) => key >= end, + Bound::Unbounded => false, + } + } + + }; +} + +impl<'a, K, V> RangeLeafIterator<'a, K, V> +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + impl_range_leaf_common!(BPlusTreeQuery, 'a); + + fn descend_to_leaf(&mut self, mut offset: u64, mut start_key: Option<&K>) -> io::Result<()> { + loop { + let (node, pointers) = if let Some(mmap) = &self.tree.mmap { + let mut cursor = io::Cursor::new(mmap.as_ref()); + BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? + } else if let Some(file) = &mut self.tree.file { + BPlusTreeNode::::deserialize_from_block(file, &mut self.tree.buffer, offset, false)? + } else { + return Err(io::Error::other("No data source available")); + }; + + if node.is_leaf { + let start_idx = if let Some(key) = start_key { + match self.start_bound { + Bound::Included(_) => node.keys.partition_point(|candidate| candidate < key), + Bound::Excluded(_) => node.keys.partition_point(|candidate| candidate <= key), + Bound::Unbounded => 0, + } + } else { + 0 + }; + self.load_leaf_from_node(node, start_idx); + return Ok(()); + } + + let child_idx = if let Some(key) = start_key { + get_entry_index_upper_bound(&node.keys, key) + } else { + 0 + }; + + let Some(ptrs) = pointers else { + self.exhausted = true; + return Ok(()); + }; + let Some(&next_offset) = ptrs.get(child_idx) else { + self.exhausted = true; + return Ok(()); + }; + self.stack.push((offset, child_idx.saturating_add(1))); + offset = next_offset; + start_key = None.or(start_key); + } + } + + fn initialize(&mut self) -> io::Result<()> { + if self.initialized { + return Ok(()); + } + self.initialized = true; + let start_key = match self.start_bound.clone() { + Bound::Included(key) | Bound::Excluded(key) => Some(key), + Bound::Unbounded => None, + }; + self.descend_to_leaf(self.tree.root_offset, start_key.as_ref()) + } + + fn advance_leaf(&mut self) -> io::Result<()> { + while let Some((offset, child_idx)) = self.stack.pop() { + let (_node, pointers) = if let Some(mmap) = &self.tree.mmap { + let mut cursor = io::Cursor::new(mmap.as_ref()); + BPlusTreeNode::::deserialize_from_mmap(mmap, &mut cursor, offset, false)? + } else if let Some(file) = &mut self.tree.file { + BPlusTreeNode::::deserialize_from_block(file, &mut self.tree.buffer, offset, false)? + } else { + return Err(io::Error::other("No data source available")); + }; + + let Some(ptrs) = pointers else { + continue; + }; + let Some(&next_offset) = ptrs.get(child_idx) else { + continue; + }; + if child_idx + 1 < ptrs.len() { + self.stack.push((offset, child_idx + 1)); + } + self.descend_to_leaf(next_offset, None)?; + return Ok(()); + } + + self.exhausted = true; + Ok(()) + } +} + +impl Iterator for RangeLeafIterator<'_, K, V> +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + type Item = Result<(K, V), BPlusTreeError>; + + fn next(&mut self) -> Option { + if self.exhausted { + return None; + } + + // Lazy initialization + if !self.initialized { + if let Err(e) = self.initialize() { + self.exhausted = true; + return Some(Err(BPlusTreeError::Io(e))); + } + } + + loop { + let Some(node) = self.current_leaf.as_ref() else { + if self.exhausted { + return None; + } + if let Err(e) = self.advance_leaf() { + self.exhausted = true; + return Some(Err(BPlusTreeError::Io(e))); + } + if self.exhausted { + return None; + } + continue; + }; + + if self.leaf_idx >= node.keys.len() { + self.current_leaf = None; + if self.exhausted { + return None; + } + if let Err(e) = self.advance_leaf() { + self.exhausted = true; + return Some(Err(BPlusTreeError::Io(e))); + } + if self.exhausted { + return None; + } + continue; + } + + let idx = self.leaf_idx; + self.leaf_idx += 1; + let key = node.keys[idx].clone(); + + if self.key_past_end(&key) { + self.current_leaf = None; + self.exhausted = true; + return None; + } + + let info = node.value_info[idx].clone(); + if self.tree.has_tombstones && info.is_tombstone() { + continue; + } + + let value = if let Some(mmap) = &self.tree.mmap { + let mut cursor = io::Cursor::new(mmap.as_ref()); + match BPlusTreeNode::::load_value_from_info(&mut cursor, &info) { + Ok(value) => value, + Err(err) => { + self.exhausted = true; + return Some(Err(BPlusTreeError::Io(err))); + } + } + } else if let Some(file) = &mut self.tree.file { + match BPlusTreeNode::::load_value_from_info(file, &info) { + Ok(value) => value, + Err(err) => { + self.exhausted = true; + return Some(Err(BPlusTreeError::Io(err))); + } + } + } else { + self.exhausted = true; + return Some(Err(BPlusTreeError::InvalidStructure("No data source available".into()))); + }; + + return Some(Ok((key, value))); + } + } +} diff --git a/backend/src/repository/bplustree/v3/format.rs b/backend/src/repository/bplustree/v3/format.rs new file mode 100644 index 000000000..2586eac3e --- /dev/null +++ b/backend/src/repository/bplustree/v3/format.rs @@ -0,0 +1,1524 @@ +use super::BPlusTreeMetadata; +use std::io; + +pub(crate) const PAGE_SIZE: usize = 4096; +const PAGE_SIZE_U32: u32 = 4096; +#[cfg(test)] +pub(crate) const STORAGE_VERSION_V1: u32 = 1; +#[cfg(test)] +pub(crate) const STORAGE_VERSION_V2: u32 = 2; +pub(crate) const STORAGE_VERSION_V3: u32 = 3; +pub(crate) const MAX_ENCODED_KEY_LEN: usize = 2004; +pub(crate) const MAX_CELL_FOOTPRINT: usize = 2032; +pub(crate) const MAX_INLINE_STORED_VALUE: usize = 512; +pub(crate) const OVERFLOW_PAYLOAD_LEN: usize = 4056; +pub(crate) const MAGIC: &[u8; 4] = b"BTRE"; + +const DATABASE_CHECKSUM_OFFSET: usize = 72; +const PAGE_CHECKSUM_OFFSET: usize = 24; +const DATABASE_METADATA_OFFSET: usize = 76; +pub(crate) const PAGE_HEADER_LEN: usize = 32; +pub(crate) const INTERNAL_PREAMBLE_LEN: usize = 8; +pub(crate) const SLOT_LEN: usize = 4; +pub(crate) const INTERNAL_CELL_PREFIX_LEN: usize = 12; +pub(crate) const LEAF_CELL_PREFIX_LEN: usize = 24; +#[cfg(test)] +pub(crate) const OVERFLOW_HEADER_LEN: usize = 8; +const COMPRESSION_MIN_LENGTH: usize = 64; +const COMPRESSION_PERCENT: usize = 85; + +fn invalid_data(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidData, message) } + +fn invalid_input(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message) } + +fn checked_end(offset: usize, length: usize) -> io::Result { + offset.checked_add(length).ok_or_else(|| invalid_data("format offset overflow")) +} + +fn bytes_at(bytes: &[u8], offset: usize) -> io::Result<[u8; N]> { + let end = checked_end(offset, N)?; + bytes + .get(offset..end) + .ok_or_else(|| invalid_data("truncated format field"))? + .try_into() + .map_err(|_| invalid_data("invalid format field length")) +} + +fn write_at(bytes: &mut [u8], offset: usize, value: &[u8]) -> io::Result<()> { + let end = checked_end(offset, value.len())?; + bytes + .get_mut(offset..end) + .ok_or_else(|| invalid_data("truncated format destination"))? + .copy_from_slice(value); + Ok(()) +} + +fn read_u16(bytes: &[u8], offset: usize) -> io::Result { Ok(u16::from_le_bytes(bytes_at(bytes, offset)?)) } + +fn read_u32(bytes: &[u8], offset: usize) -> io::Result { Ok(u32::from_le_bytes(bytes_at(bytes, offset)?)) } + +fn read_u64(bytes: &[u8], offset: usize) -> io::Result { Ok(u64::from_le_bytes(bytes_at(bytes, offset)?)) } + +fn read_u8(bytes: &[u8], offset: usize) -> io::Result { + let [value] = bytes_at(bytes, offset)?; + Ok(value) +} + +fn require_zero(bytes: &[u8], message: &'static str) -> io::Result<()> { + if bytes.iter().all(|byte| *byte == 0) { + Ok(()) + } else { + Err(invalid_data(message)) + } +} + +fn exact_page(bytes: &[u8]) -> io::Result<()> { + if bytes.len() == PAGE_SIZE { + Ok(()) + } else { + Err(invalid_data("page must be exactly 4096 bytes")) + } +} + +fn checksum_with_zeroed_field(page: &[u8], offset: usize) -> io::Result { + exact_page(page)?; + let checksum_end = checked_end(offset, 4)?; + let before = page.get(..offset).ok_or_else(|| invalid_data("missing checksum prefix"))?; + let after = page.get(checksum_end..).ok_or_else(|| invalid_data("missing checksum suffix"))?; + let mut hasher = crc32fast::Hasher::new(); + hasher.update(before); + hasher.update(&[0; 4]); + hasher.update(after); + Ok(hasher.finalize()) +} + +fn write_checksum(page: &mut [u8], offset: usize) -> io::Result<()> { + exact_page(page)?; + write_at(page, offset, &[0; 4])?; + let checksum = checksum_with_zeroed_field(page, offset)?; + write_at(page, offset, &checksum.to_le_bytes()) +} + +fn verify_checksum(page: &[u8], offset: usize) -> io::Result<()> { + let stored = read_u32(page, offset)?; + if stored == checksum_with_zeroed_field(page, offset)? { + Ok(()) + } else { + Err(invalid_data("page checksum mismatch")) + } +} + +fn validate_database_header(header: &DatabaseHeader) -> io::Result<()> { + if header.root_page_id == 0 || header.root_page_id >= header.next_page_id { + return Err(invalid_data("invalid root page id")); + } + if header.free_page_head != 0 + && (header.free_page_head >= header.next_page_id || header.free_page_head == header.root_page_id) + { + return Err(invalid_data("invalid free page head")); + } + if header.generation == 0 { + return Err(invalid_data("generation must be nonzero")); + } + if header.database_id.iter().all(|byte| *byte == 0) { + return Err(invalid_data("database identity must be nonzero")); + } + Ok(()) +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct DatabaseHeader { + pub(crate) root_page_id: u64, + pub(crate) next_page_id: u64, + pub(crate) free_page_head: u64, + pub(crate) generation: u64, + pub(crate) database_id: [u8; 16], + pub(crate) metadata: BPlusTreeMetadata, +} + +impl DatabaseHeader { + pub(crate) fn encode(&self) -> io::Result<[u8; PAGE_SIZE]> { + validate_database_header(self)?; + let mut page = [0u8; PAGE_SIZE]; + write_at(&mut page, 0, MAGIC)?; + write_at(&mut page, 4, &STORAGE_VERSION_V3.to_le_bytes())?; + write_at(&mut page, 8, &PAGE_SIZE_U32.to_le_bytes())?; + write_at(&mut page, 16, &self.root_page_id.to_le_bytes())?; + write_at(&mut page, 24, &self.next_page_id.to_le_bytes())?; + write_at(&mut page, 32, &self.free_page_head.to_le_bytes())?; + write_at(&mut page, 40, &self.generation.to_le_bytes())?; + write_at(&mut page, 48, &self.database_id)?; + match self.metadata { + BPlusTreeMetadata::Empty => {} + BPlusTreeMetadata::TargetIdMapping(value) => { + write_at(&mut page, 64, &5u32.to_le_bytes())?; + write_at(&mut page, DATABASE_METADATA_OFFSET, &[1])?; + write_at(&mut page, DATABASE_METADATA_OFFSET + 1, &value.to_le_bytes())?; + } + } + write_checksum(&mut page, DATABASE_CHECKSUM_OFFSET)?; + Ok(page) + } + + pub(crate) fn decode(page: &[u8]) -> io::Result { + exact_page(page)?; + if bytes_at::<4>(page, 0)? != *MAGIC { + return Err(invalid_data("invalid database magic")); + } + if read_u32(page, 4)? != STORAGE_VERSION_V3 { + return Err(invalid_data("unsupported storage version")); + } + if read_u32(page, 8)? != PAGE_SIZE_U32 { + return Err(invalid_data("invalid page size")); + } + if read_u32(page, 12)? != 0 { + return Err(invalid_data("unknown database feature flags")); + } + require_zero( + page.get(68..72).ok_or_else(|| invalid_data("missing database reserved bytes"))?, + "database reserved bytes must be zero", + )?; + verify_checksum(page, DATABASE_CHECKSUM_OFFSET)?; + + let metadata_length = read_u32(page, 64)?; + let metadata = match metadata_length { + 0 => BPlusTreeMetadata::Empty, + 5 => { + let encoded = bytes_at::<5>(page, DATABASE_METADATA_OFFSET)?; + let [tag, value0, value1, value2, value3] = encoded; + if tag != 1 { + return Err(invalid_data("unknown metadata tag")); + } + BPlusTreeMetadata::TargetIdMapping(u32::from_le_bytes([value0, value1, value2, value3])) + } + _ => return Err(invalid_data("invalid metadata length")), + }; + let tail_start = checked_end( + DATABASE_METADATA_OFFSET, + usize::try_from(metadata_length).map_err(|_| invalid_data("metadata length exceeds usize"))?, + )?; + require_zero( + page.get(tail_start..).ok_or_else(|| invalid_data("metadata extends beyond header"))?, + "database header tail must be zero", + )?; + + let header = Self { + root_page_id: read_u64(page, 16)?, + next_page_id: read_u64(page, 24)?, + free_page_head: read_u64(page, 32)?, + generation: read_u64(page, 40)?, + database_id: bytes_at(page, 48)?, + metadata, + }; + validate_database_header(&header)?; + Ok(header) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum PageType { + Leaf = 1, + Internal = 2, + Overflow = 3, + Free = 4, +} + +impl TryFrom for PageType { + type Error = io::Error; + + fn try_from(value: u8) -> Result { + match value { + 1 => Ok(Self::Leaf), + 2 => Ok(Self::Internal), + 3 => Ok(Self::Overflow), + 4 => Ok(Self::Free), + _ => Err(invalid_data("unknown page type")), + } + } +} + +fn validate_page_id(page_id: u64, next_page_id: u64, kind: io::ErrorKind) -> io::Result<()> { + if page_id != 0 && page_id < next_page_id { + Ok(()) + } else { + Err(io::Error::new(kind, "invalid page id bounds")) + } +} + +fn validate_reference(reference: u64, page_id: u64, next_page_id: u64, kind: io::ErrorKind) -> io::Result<()> { + if reference == 0 || (reference < next_page_id && reference != page_id) { + Ok(()) + } else { + Err(io::Error::new(kind, "invalid page reference")) + } +} + +fn expected_slot_end(base: usize, cell_count: u16, kind: io::ErrorKind) -> io::Result { + let slots = usize::from(cell_count) + .checked_mul(SLOT_LEN) + .ok_or_else(|| io::Error::new(kind, "slot directory size overflow"))?; + let end = base.checked_add(slots).ok_or_else(|| io::Error::new(kind, "slot directory offset overflow"))?; + u16::try_from(end).map_err(|_| io::Error::new(kind, "slot directory exceeds page")) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct PageHeader { + pub(crate) page_type: PageType, + pub(crate) cell_count: u16, + pub(crate) free_start: u16, + pub(crate) free_end: u16, + pub(crate) left: u64, + pub(crate) right: u64, +} + +impl PageHeader { + fn validate(&self, page_id: u64, next_page_id: u64, kind: io::ErrorKind) -> io::Result<()> { + validate_page_id(page_id, next_page_id, kind)?; + match self.page_type { + PageType::Leaf => { + if self.free_start != expected_slot_end(PAGE_HEADER_LEN, self.cell_count, kind)? { + return Err(io::Error::new(kind, "invalid leaf free_start")); + } + if self.cell_count == 0 { + if usize::from(self.free_end) != PAGE_SIZE { + return Err(io::Error::new(kind, "invalid empty leaf free_end")); + } + } else if self.free_end < self.free_start || usize::from(self.free_end) >= PAGE_SIZE { + return Err(io::Error::new(kind, "invalid leaf free_end")); + } + validate_reference(self.left, page_id, next_page_id, kind)?; + validate_reference(self.right, page_id, next_page_id, kind) + } + PageType::Internal => { + if self.cell_count == 0 + || self.free_start != expected_slot_end(PAGE_HEADER_LEN + INTERNAL_PREAMBLE_LEN, self.cell_count, kind)? + || self.free_end < self.free_start + || usize::from(self.free_end) >= PAGE_SIZE + || self.left != 0 + || self.right != 0 + { + return Err(io::Error::new(kind, "invalid internal page header")); + } + Ok(()) + } + PageType::Overflow | PageType::Free => { + if self.cell_count != 0 || self.free_start != 0 || self.free_end != 0 || self.left != 0 { + return Err(io::Error::new(kind, "invalid chain page header")); + } + validate_reference(self.right, page_id, next_page_id, kind) + } + } + } + + pub(crate) fn encode_into(&self, page: &mut [u8], page_id: u64, next_page_id: u64) -> io::Result<()> { + exact_page(page)?; + self.validate(page_id, next_page_id, io::ErrorKind::InvalidInput)?; + write_at(page, 0, &[self.page_type as u8])?; + write_at(page, 1, &[0])?; + write_at(page, 2, &self.cell_count.to_le_bytes())?; + write_at(page, 4, &self.free_start.to_le_bytes())?; + write_at(page, 6, &self.free_end.to_le_bytes())?; + write_at(page, 8, &self.left.to_le_bytes())?; + write_at(page, 16, &self.right.to_le_bytes())?; + write_at(page, 28, &[0; 4])?; + write_page_checksum(page) + } + + pub(crate) fn decode(page: &[u8], page_id: u64, next_page_id: u64) -> io::Result { + exact_page(page)?; + verify_page_checksum(page)?; + let page_type = PageType::try_from(read_u8(page, 0)?)?; + if read_u8(page, 1)? != 0 { + return Err(invalid_data("unknown page flags")); + } + require_zero( + page.get(28..32).ok_or_else(|| invalid_data("missing page reserved bytes"))?, + "page reserved bytes must be zero", + )?; + let header = Self { + page_type, + cell_count: read_u16(page, 2)?, + free_start: read_u16(page, 4)?, + free_end: read_u16(page, 6)?, + left: read_u64(page, 8)?, + right: read_u64(page, 16)?, + }; + header.validate(page_id, next_page_id, io::ErrorKind::InvalidData)?; + Ok(header) + } +} + +#[cfg(test)] +pub(crate) fn page_checksum(page: &[u8]) -> io::Result { checksum_with_zeroed_field(page, PAGE_CHECKSUM_OFFSET) } + +pub(crate) fn write_page_checksum(page: &mut [u8]) -> io::Result<()> { write_checksum(page, PAGE_CHECKSUM_OFFSET) } + +pub(crate) fn verify_page_checksum(page: &[u8]) -> io::Result<()> { verify_checksum(page, PAGE_CHECKSUM_OFFSET) } + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct Slot { + pub(crate) offset: u16, + pub(crate) length: u16, +} + +impl Slot { + pub(crate) fn encode(self) -> [u8; SLOT_LEN] { + let [offset0, offset1] = self.offset.to_le_bytes(); + let [length0, length1] = self.length.to_le_bytes(); + [offset0, offset1, length0, length1] + } + + pub(crate) fn decode(bytes: &[u8]) -> io::Result { + if bytes.len() != SLOT_LEN { + return Err(invalid_data("invalid slot length")); + } + Ok(Self { offset: read_u16(bytes, 0)?, length: read_u16(bytes, 2)? }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct InternalPreamble { + pub(crate) leftmost_child: u64, +} + +impl InternalPreamble { + pub(crate) fn encode_into(self, page: &mut [u8], page_id: u64, next_page_id: u64) -> io::Result<()> { + exact_page(page)?; + if self.leftmost_child == 0 { + return Err(invalid_input("internal leftmost child must be nonzero")); + } + validate_reference(self.leftmost_child, page_id, next_page_id, io::ErrorKind::InvalidInput)?; + write_at(page, PAGE_HEADER_LEN, &self.leftmost_child.to_le_bytes()) + } + + pub(crate) fn decode(page: &[u8], page_id: u64, next_page_id: u64) -> io::Result { + exact_page(page)?; + let leftmost_child = read_u64(page, PAGE_HEADER_LEN)?; + if leftmost_child == 0 { + return Err(invalid_data("internal leftmost child must be nonzero")); + } + validate_reference(leftmost_child, page_id, next_page_id, io::ErrorKind::InvalidData)?; + Ok(Self { leftmost_child }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct InternalCellPrefix { + pub(crate) key_length: u16, + pub(crate) right_child: u64, +} + +impl InternalCellPrefix { + fn validate(&self, page_id: u64, next_page_id: u64, kind: io::ErrorKind) -> io::Result<()> { + if self.key_length == 0 || usize::from(self.key_length) > MAX_ENCODED_KEY_LEN || self.right_child == 0 { + return Err(io::Error::new(kind, "invalid internal cell prefix")); + } + validate_reference(self.right_child, page_id, next_page_id, kind) + } + + pub(crate) fn encode(&self, page_id: u64, next_page_id: u64) -> io::Result<[u8; INTERNAL_CELL_PREFIX_LEN]> { + self.validate(page_id, next_page_id, io::ErrorKind::InvalidInput)?; + let mut bytes = [0u8; INTERNAL_CELL_PREFIX_LEN]; + write_at(&mut bytes, 0, &self.key_length.to_le_bytes())?; + write_at(&mut bytes, 4, &self.right_child.to_le_bytes())?; + Ok(bytes) + } + + pub(crate) fn decode(bytes: &[u8], page_id: u64, next_page_id: u64) -> io::Result { + let prefix = bytes + .get(..INTERNAL_CELL_PREFIX_LEN) + .ok_or_else(|| invalid_data("truncated internal cell prefix"))?; + require_zero( + prefix.get(2..4).ok_or_else(|| invalid_data("missing internal cell reserved bytes"))?, + "internal cell reserved bytes must be zero", + )?; + let decoded = Self { key_length: read_u16(prefix, 0)?, right_child: read_u64(prefix, 4)? }; + decoded.validate(page_id, next_page_id, io::ErrorKind::InvalidData)?; + Ok(decoded) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum ValueKind { + Inline = 0, + Overflow = 1, + Tombstone = 2, +} + +impl TryFrom for ValueKind { + type Error = io::Error; + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(Self::Inline), + 1 => Ok(Self::Overflow), + 2 => Ok(Self::Tombstone), + _ => Err(invalid_data("unknown value kind")), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum Compression { + None = 0, + Lz4 = 1, +} + +impl TryFrom for Compression { + type Error = io::Error; + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(Self::None), + 1 => Ok(Self::Lz4), + _ => Err(invalid_data("unknown compression mode")), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct LeafCellPrefix { + pub(crate) key_length: u16, + pub(crate) value_kind: ValueKind, + pub(crate) compression: Compression, + pub(crate) logical_length: u32, + pub(crate) stored_length: u32, + pub(crate) overflow_head: u64, + pub(crate) stored_crc32: u32, +} + +impl LeafCellPrefix { + fn validate(&self, page_id: u64, next_page_id: u64, kind: io::ErrorKind) -> io::Result<()> { + if self.key_length == 0 || usize::from(self.key_length) > MAX_ENCODED_KEY_LEN { + return Err(io::Error::new(kind, "invalid leaf key length")); + } + let stored_on_page = match self.value_kind { + ValueKind::Inline => usize::try_from(self.stored_length) + .map_err(|_| io::Error::new(kind, "inline stored length exceeds usize"))?, + ValueKind::Overflow | ValueKind::Tombstone => 0, + }; + let footprint = SLOT_LEN + .checked_add(LEAF_CELL_PREFIX_LEN) + .and_then(|size| size.checked_add(usize::from(self.key_length))) + .and_then(|size| size.checked_add(stored_on_page)) + .ok_or_else(|| io::Error::new(kind, "leaf cell footprint overflow"))?; + if footprint > MAX_CELL_FOOTPRINT { + return Err(io::Error::new(kind, "leaf cell footprint exceeds limit")); + } + match self.value_kind { + ValueKind::Inline => { + if self.logical_length == 0 + || self.stored_length == 0 + || stored_on_page > MAX_INLINE_STORED_VALUE + || self.overflow_head != 0 + || (self.compression == Compression::None && self.logical_length != self.stored_length) + { + return Err(io::Error::new(kind, "invalid inline value descriptor")); + } + } + ValueKind::Overflow => { + if self.logical_length == 0 || self.stored_length == 0 || self.overflow_head == 0 { + return Err(io::Error::new(kind, "invalid overflow value descriptor")); + } + if self.compression == Compression::None && self.logical_length != self.stored_length { + return Err(io::Error::new(kind, "invalid uncompressed overflow lengths")); + } + validate_reference(self.overflow_head, page_id, next_page_id, kind)?; + } + ValueKind::Tombstone => { + if self.compression != Compression::None + || self.logical_length != 0 + || self.stored_length != 0 + || self.overflow_head != 0 + || self.stored_crc32 != 0 + { + return Err(io::Error::new(kind, "invalid tombstone descriptor")); + } + } + } + if self.compression == Compression::Lz4 { + let threshold = u64::from(self.logical_length) * COMPRESSION_PERCENT as u64 / 100; + if self.logical_length < 64 || self.stored_length < 4 || u64::from(self.stored_length) >= threshold { + return Err(io::Error::new(kind, "invalid compressed value lengths")); + } + } + Ok(()) + } + + pub(crate) fn encode(&self, page_id: u64, next_page_id: u64) -> io::Result<[u8; LEAF_CELL_PREFIX_LEN]> { + self.validate(page_id, next_page_id, io::ErrorKind::InvalidInput)?; + let mut bytes = [0u8; LEAF_CELL_PREFIX_LEN]; + write_at(&mut bytes, 0, &self.key_length.to_le_bytes())?; + write_at(&mut bytes, 2, &[self.value_kind as u8])?; + write_at(&mut bytes, 3, &[self.compression as u8])?; + write_at(&mut bytes, 4, &self.logical_length.to_le_bytes())?; + write_at(&mut bytes, 8, &self.stored_length.to_le_bytes())?; + write_at(&mut bytes, 12, &self.overflow_head.to_le_bytes())?; + write_at(&mut bytes, 20, &self.stored_crc32.to_le_bytes())?; + Ok(bytes) + } + + pub(crate) fn decode(bytes: &[u8], page_id: u64, next_page_id: u64) -> io::Result { + let prefix = bytes + .get(..LEAF_CELL_PREFIX_LEN) + .ok_or_else(|| invalid_data("truncated leaf cell prefix"))?; + let decoded = Self { + key_length: read_u16(prefix, 0)?, + value_kind: ValueKind::try_from(read_u8(prefix, 2)?)?, + compression: Compression::try_from(read_u8(prefix, 3)?)?, + logical_length: read_u32(prefix, 4)?, + stored_length: read_u32(prefix, 8)?, + overflow_head: read_u64(prefix, 12)?, + stored_crc32: read_u32(prefix, 20)?, + }; + decoded.validate(page_id, next_page_id, io::ErrorKind::InvalidData)?; + Ok(decoded) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct OverflowHeader { + pub(crate) payload_length: u16, +} + +impl OverflowHeader { + pub(crate) fn encode_into(self, page: &mut [u8]) -> io::Result<()> { + exact_page(page)?; + let payload_length = usize::from(self.payload_length); + if payload_length > OVERFLOW_PAYLOAD_LEN { + return Err(invalid_input("overflow payload is too large")); + } + write_at(page, PAGE_HEADER_LEN, &self.payload_length.to_le_bytes())?; + write_at(page, PAGE_HEADER_LEN + 2, &[0; 6])?; + let tail_start = checked_end(40, payload_length)?; + page.get_mut(tail_start..) + .ok_or_else(|| invalid_data("overflow payload exceeds page"))? + .fill(0); + Ok(()) + } + + pub(crate) fn decode(page: &[u8]) -> io::Result { + exact_page(page)?; + let payload_length = read_u16(page, PAGE_HEADER_LEN)?; + if usize::from(payload_length) > OVERFLOW_PAYLOAD_LEN { + return Err(invalid_data("overflow payload is too large")); + } + require_zero( + page.get(34..40).ok_or_else(|| invalid_data("missing overflow reserved bytes"))?, + "overflow reserved bytes must be zero", + )?; + let tail_start = checked_end(40, usize::from(payload_length))?; + require_zero( + page.get(tail_start..).ok_or_else(|| invalid_data("overflow payload exceeds page"))?, + "overflow page tail must be zero", + )?; + Ok(Self { payload_length }) + } +} + +pub(crate) fn encode_free_body(page: &mut [u8]) -> io::Result<()> { + exact_page(page)?; + page.get_mut(PAGE_HEADER_LEN..) + .ok_or_else(|| invalid_data("missing free page body"))? + .fill(0); + Ok(()) +} + +pub(crate) fn validate_free_body(page: &[u8]) -> io::Result<()> { + exact_page(page)?; + require_zero( + page.get(PAGE_HEADER_LEN..).ok_or_else(|| invalid_data("missing free page body"))?, + "free page body must be zero", + ) +} + +pub(crate) fn stored_value_checksum(stored: &[u8]) -> u32 { crc32fast::hash(stored) } + +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum StoredValue<'a> { + BorrowedRaw(&'a [u8]), + Compressed(&'a [u8]), +} + +impl StoredValue<'_> { + pub(crate) const fn compression(&self) -> Compression { + match self { + Self::BorrowedRaw(_) => Compression::None, + Self::Compressed(_) => Compression::Lz4, + } + } + + pub(crate) const fn as_slice(&self) -> &[u8] { + match self { + Self::BorrowedRaw(bytes) | Self::Compressed(bytes) => bytes, + } + } +} + +fn compression_is_beneficial(raw_length: usize, stored_length: usize) -> io::Result { + let threshold = raw_length + .checked_mul(COMPRESSION_PERCENT) + .ok_or_else(|| invalid_input("compression threshold overflow"))? + / 100; + Ok(stored_length < threshold) +} + +pub(crate) fn encode_value<'a>(raw: &'a [u8], scratch: &'a mut Vec) -> io::Result> { + scratch.clear(); + if raw.len() < COMPRESSION_MIN_LENGTH { + return Ok(StoredValue::BorrowedRaw(raw)); + } + let raw_length = u32::try_from(raw.len()).map_err(|_| invalid_input("value exceeds u32 length"))?; + let maximum = lz4_flex::block::get_maximum_output_size(raw.len()) + .checked_add(4) + .ok_or_else(|| invalid_input("compressed value size overflow"))?; + scratch + .try_reserve(maximum) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + scratch.resize(maximum, 0); + write_at(scratch, 0, &raw_length.to_le_bytes())?; + let output = scratch.get_mut(4..).ok_or_else(|| invalid_data("missing compression output"))?; + let compressed_length = lz4_flex::block::compress_into(raw, output) + .map_err(|err| io::Error::other(format!("LZ4 compression failed: {err}")))?; + let stored_length = checked_end(4, compressed_length)?; + scratch.truncate(stored_length); + if compression_is_beneficial(raw.len(), scratch.len())? { + Ok(StoredValue::Compressed(scratch.as_slice())) + } else { + Ok(StoredValue::BorrowedRaw(raw)) + } +} + +pub(crate) fn decompress_value_into<'a>( + stored: &[u8], + logical_length: u32, + maximum_length: usize, + scratch: &'a mut Vec, +) -> io::Result<&'a [u8]> { + let encoded_length = read_u32(stored, 0)?; + if encoded_length != logical_length { + return Err(invalid_data("LZ4 logical length mismatch")); + } + let logical_length = usize::try_from(logical_length).map_err(|_| invalid_data("logical length exceeds usize"))?; + if logical_length > maximum_length { + return Err(invalid_data("logical length exceeds allocation limit")); + } + let payload = stored.get(4..).ok_or_else(|| invalid_data("missing LZ4 payload"))?; + scratch.clear(); + scratch + .try_reserve(logical_length) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + scratch.resize(logical_length, 0); + let decoded_length = lz4_flex::block::decompress_into(payload, scratch.as_mut_slice()) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 decompression failed: {err}")))?; + if decoded_length != logical_length { + return Err(invalid_data("decompressed value length mismatch")); + } + scratch.get(..decoded_length).ok_or_else(|| invalid_data("decompressed value exceeds scratch buffer")) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct LeafCellRef<'a> { + pub(crate) key_bytes: &'a [u8], + pub(crate) value: LeafValueRef<'a>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum LeafValueRef<'a> { + Inline { + compression: Compression, + logical_len: u32, + stored: &'a [u8], + crc32: u32, + }, + Overflow { + compression: Compression, + logical_len: u32, + stored_len: u32, + head: u64, + crc32: u32, + }, + Tombstone, +} + +impl<'a> LeafCellRef<'a> { + pub(crate) fn decode(cell: &'a [u8], page_id: u64, next_page_id: u64) -> io::Result { + let prefix = LeafCellPrefix::decode(cell, page_id, next_page_id)?; + let key_end = LEAF_CELL_PREFIX_LEN + .checked_add(usize::from(prefix.key_length)) + .ok_or_else(|| invalid_data("leaf key range overflow"))?; + let key_bytes = cell + .get(LEAF_CELL_PREFIX_LEN..key_end) + .ok_or_else(|| invalid_data("truncated leaf key"))?; + let value = match prefix.value_kind { + ValueKind::Inline => { + let cell_end = key_end + .checked_add( + usize::try_from(prefix.stored_length) + .map_err(|_| invalid_data("inline stored length exceeds usize"))?, + ) + .ok_or_else(|| invalid_data("inline value range overflow"))?; + if cell_end != cell.len() { + return Err(invalid_data("invalid inline leaf cell length")); + } + let stored = cell.get(key_end..cell_end).ok_or_else(|| invalid_data("truncated inline value"))?; + if stored_value_checksum(stored) != prefix.stored_crc32 { + return Err(invalid_data("stored value checksum mismatch")); + } + LeafValueRef::Inline { + compression: prefix.compression, + logical_len: prefix.logical_length, + stored, + crc32: prefix.stored_crc32, + } + } + ValueKind::Overflow => { + if key_end != cell.len() { + return Err(invalid_data("invalid overflow leaf cell length")); + } + LeafValueRef::Overflow { + compression: prefix.compression, + logical_len: prefix.logical_length, + stored_len: prefix.stored_length, + head: prefix.overflow_head, + crc32: prefix.stored_crc32, + } + } + ValueKind::Tombstone => { + if key_end != cell.len() { + return Err(invalid_data("invalid tombstone leaf cell length")); + } + LeafValueRef::Tombstone + } + }; + Ok(Self { key_bytes, value }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct InternalCellRef<'a> { + pub(crate) key_bytes: &'a [u8], + pub(crate) right_child: u64, +} + +impl<'a> InternalCellRef<'a> { + pub(crate) fn decode(cell: &'a [u8], page_id: u64, next_page_id: u64) -> io::Result { + let prefix = InternalCellPrefix::decode(cell, page_id, next_page_id)?; + let cell_end = INTERNAL_CELL_PREFIX_LEN + .checked_add(usize::from(prefix.key_length)) + .ok_or_else(|| invalid_data("internal key range overflow"))?; + if cell_end != cell.len() { + return Err(invalid_data("invalid internal cell length")); + } + let key_bytes = cell + .get(INTERNAL_CELL_PREFIX_LEN..cell_end) + .ok_or_else(|| invalid_data("truncated internal key"))?; + Ok(Self { key_bytes, right_child: prefix.right_child }) + } +} + +fn encoded_key_length(key: &[u8]) -> io::Result { + if key.is_empty() || key.len() > MAX_ENCODED_KEY_LEN { + return Err(invalid_input("invalid encoded key length")); + } + u16::try_from(key.len()).map_err(|_| invalid_input("encoded key length exceeds u16")) +} + +fn write_cell(output: &mut Vec, prefix: &[u8], key: &[u8], stored: &[u8]) -> io::Result<()> { + let length = prefix + .len() + .checked_add(key.len()) + .and_then(|value| value.checked_add(stored.len())) + .ok_or_else(|| invalid_input("encoded cell length overflow"))?; + output.clear(); + output + .try_reserve(length) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + output.extend_from_slice(prefix); + output.extend_from_slice(key); + output.extend_from_slice(stored); + Ok(()) +} + +pub(crate) fn encode_inline_leaf_cell( + key: &[u8], + logical_len: u32, + compression: Compression, + stored: &[u8], + output: &mut Vec, +) -> io::Result<()> { + let prefix = LeafCellPrefix { + key_length: encoded_key_length(key)?, + value_kind: ValueKind::Inline, + compression, + logical_length: logical_len, + stored_length: u32::try_from(stored.len()).map_err(|_| invalid_input("stored value exceeds u32"))?, + overflow_head: 0, + stored_crc32: stored_value_checksum(stored), + } + .encode(1, 2)?; + write_cell(output, &prefix, key, stored) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn encode_overflow_leaf_cell( + key: &[u8], + logical_len: u32, + compression: Compression, + stored_len: u32, + head: u64, + stored_crc32: u32, + page_id: u64, + next_page_id: u64, + output: &mut Vec, +) -> io::Result<()> { + let prefix = LeafCellPrefix { + key_length: encoded_key_length(key)?, + value_kind: ValueKind::Overflow, + compression, + logical_length: logical_len, + stored_length: stored_len, + overflow_head: head, + stored_crc32, + } + .encode(page_id, next_page_id)?; + write_cell(output, &prefix, key, &[]) +} + +pub(crate) fn encode_tombstone_leaf_cell(key: &[u8], output: &mut Vec) -> io::Result<()> { + let prefix = LeafCellPrefix { + key_length: encoded_key_length(key)?, + value_kind: ValueKind::Tombstone, + compression: Compression::None, + logical_length: 0, + stored_length: 0, + overflow_head: 0, + stored_crc32: 0, + } + .encode(1, 2)?; + write_cell(output, &prefix, key, &[]) +} + +pub(crate) fn encode_internal_cell( + key: &[u8], + right_child: u64, + page_id: u64, + next_page_id: u64, + output: &mut Vec, +) -> io::Result<()> { + let prefix = InternalCellPrefix { key_length: encoded_key_length(key)?, right_child } + .encode(page_id, next_page_id)?; + write_cell(output, &prefix, key, &[]) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct Locator { + pub(crate) leaf_page_id: u64, + pub(crate) slot_index: u16, + pub(crate) serialized_key_crc32: u32, +} + +impl Locator { + pub(crate) fn for_key(leaf_page_id: u64, slot_index: u16, serialized_key: &[u8]) -> io::Result { + if leaf_page_id == 0 { + return Err(invalid_input("locator leaf page id must be nonzero")); + } + Ok(Self { leaf_page_id, slot_index, serialized_key_crc32: crc32fast::hash(serialized_key) }) + } + + pub(crate) fn encode(self) -> [u8; 16] { + let mut encoded = [0; 16]; + encoded[0..8].copy_from_slice(&self.leaf_page_id.to_le_bytes()); + encoded[8..10].copy_from_slice(&self.slot_index.to_le_bytes()); + encoded[12..16].copy_from_slice(&self.serialized_key_crc32.to_le_bytes()); + encoded + } + + pub(crate) fn decode(encoded: &[u8]) -> io::Result { + if encoded.len() != 16 { + return Err(invalid_data("locator must be exactly 16 bytes")); + } + require_zero( + encoded.get(10..12).ok_or_else(|| invalid_data("missing locator reserved bytes"))?, + "locator reserved bytes must be zero", + )?; + let locator = Self { + leaf_page_id: read_u64(encoded, 0)?, + slot_index: read_u16(encoded, 8)?, + serialized_key_crc32: read_u32(encoded, 12)?, + }; + if locator.leaf_page_id == 0 { + return Err(invalid_data("locator leaf page id must be nonzero")); + } + Ok(locator) + } +} + +pub(crate) fn decompress_value_in_place( + scratch: &mut Vec, + logical_length: u32, + maximum_length: usize, +) -> io::Result<()> { + let encoded_length = read_u32(scratch, 0)?; + if encoded_length != logical_length { + return Err(invalid_data("LZ4 logical length mismatch")); + } + let logical_length = usize::try_from(logical_length).map_err(|_| invalid_data("logical length exceeds usize"))?; + if logical_length > maximum_length { + return Err(invalid_data("logical length exceeds allocation limit")); + } + let stored_length = scratch.len(); + let total_length = logical_length + .checked_add(stored_length) + .ok_or_else(|| invalid_data("in-place decompression size overflow"))?; + scratch + .try_reserve(total_length.saturating_sub(stored_length)) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + scratch.resize(total_length, 0); + scratch.copy_within(0..stored_length, logical_length); + let (output, encoded) = scratch.split_at_mut(logical_length); + let payload = encoded + .get(4..stored_length) + .ok_or_else(|| invalid_data("missing LZ4 payload"))?; + let decoded_length = lz4_flex::block::decompress_into(payload, output) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 decompression failed: {err}")))?; + if decoded_length != logical_length { + return Err(invalid_data("decompressed value length mismatch")); + } + scratch.truncate(logical_length); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::repository::bplustree::v3::BPlusTreeMetadata; + use std::io; + + const PAGE_ID: u64 = 7; + const NEXT_PAGE_ID: u64 = 19; + + fn invalid_data(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidData => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidData, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidData")), + } + } + + fn invalid_input(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidInput => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidInput, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidInput")), + } + } + + fn u32_at(bytes: &[u8], offset: usize) -> io::Result { + let value = bytes + .get(offset..offset + 4) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing u32"))? + .try_into() + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + Ok(u32::from_le_bytes(value)) + } + + fn write_crc(page: &mut [u8], checksum_offset: usize) -> io::Result<()> { + page.get_mut(checksum_offset..checksum_offset + 4) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing checksum"))? + .fill(0); + let checksum = crc32fast::hash(page); + page.get_mut(checksum_offset..checksum_offset + 4) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing checksum"))? + .copy_from_slice(&checksum.to_le_bytes()); + Ok(()) + } + + #[test] + fn write_checksum_rejects_short_buffer_without_mutating_it() -> io::Result<()> { + let mut short = [0x5a; DATABASE_CHECKSUM_OFFSET + 4]; + let original = short; + invalid_data(write_checksum(&mut short, DATABASE_CHECKSUM_OFFSET))?; + assert_eq!(short, original); + Ok(()) + } + + fn golden_database_header() -> [u8; PAGE_SIZE] { + let mut page = [0u8; PAGE_SIZE]; + page[0..4].copy_from_slice(b"BTRE"); + page[4..8].copy_from_slice(&3u32.to_le_bytes()); + page[8..12].copy_from_slice(&4096u32.to_le_bytes()); + page[16..24].copy_from_slice(&7u64.to_le_bytes()); + page[24..32].copy_from_slice(&19u64.to_le_bytes()); + page[32..40].copy_from_slice(&3u64.to_le_bytes()); + page[40..48].copy_from_slice(&5u64.to_le_bytes()); + page[48..64].fill(0x11); + page[64..68].copy_from_slice(&5u32.to_le_bytes()); + page[72..76].copy_from_slice(&0xaa07_e109u32.to_le_bytes()); + page[76..81].copy_from_slice(&[1, 42, 0, 0, 0]); + page + } + + fn page_fixture( + page_type: u8, + cell_count: u16, + free_start: u16, + free_end: u16, + left: u64, + right: u64, + checksum: u32, + ) -> [u8; PAGE_SIZE] { + let mut page = [0u8; PAGE_SIZE]; + page[0] = page_type; + page[2..4].copy_from_slice(&cell_count.to_le_bytes()); + page[4..6].copy_from_slice(&free_start.to_le_bytes()); + page[6..8].copy_from_slice(&free_end.to_le_bytes()); + page[8..16].copy_from_slice(&left.to_le_bytes()); + page[16..24].copy_from_slice(&right.to_le_bytes()); + page[24..28].copy_from_slice(&checksum.to_le_bytes()); + page + } + + #[test] + fn format_constants_are_frozen() { + assert_eq!(PAGE_SIZE, 4096); + assert_eq!(STORAGE_VERSION_V1, 1); + assert_eq!(STORAGE_VERSION_V2, 2); + assert_eq!(STORAGE_VERSION_V3, 3); + assert_eq!(MAX_ENCODED_KEY_LEN, 2004); + assert_eq!(MAX_CELL_FOOTPRINT, 2032); + assert_eq!(MAX_INLINE_STORED_VALUE, 512); + assert_eq!(OVERFLOW_PAYLOAD_LEN, 4056); + assert_eq!(PAGE_HEADER_LEN, 32); + assert_eq!(SLOT_LEN, 4); + assert_eq!(INTERNAL_PREAMBLE_LEN, 8); + assert_eq!(INTERNAL_CELL_PREFIX_LEN, 12); + assert_eq!(LEAF_CELL_PREFIX_LEN, 24); + assert_eq!(OVERFLOW_HEADER_LEN, 8); + assert_eq!(MAGIC, b"BTRE"); + } + + #[test] + fn database_header_golden_bytes_round_trip() -> io::Result<()> { + let header = DatabaseHeader { + root_page_id: 7, + next_page_id: 19, + free_page_head: 3, + generation: 5, + database_id: [0x11; 16], + metadata: BPlusTreeMetadata::TargetIdMapping(42), + }; + + let encoded = header.encode()?; + assert_eq!(encoded, golden_database_header()); + assert_eq!(&encoded[0..4], b"BTRE"); + assert_eq!(u32_at(&encoded, 4)?, STORAGE_VERSION_V3); + assert_eq!(u32_at(&encoded, 8)?, PAGE_SIZE_U32); + assert_eq!(u32_at(&encoded, 72)?, 0xaa07_e109); + assert_eq!(DatabaseHeader::decode(&encoded)?, header); + Ok(()) + } + + #[test] + fn database_header_empty_metadata_round_trip() -> io::Result<()> { + let header = DatabaseHeader { + root_page_id: 1, + next_page_id: 2, + free_page_head: 0, + generation: 1, + database_id: [0x22; 16], + metadata: BPlusTreeMetadata::Empty, + }; + let encoded = header.encode()?; + assert_eq!(u32_at(&encoded, 64)?, 0); + assert!(encoded[76..].iter().all(|byte| *byte == 0)); + assert_eq!(DatabaseHeader::decode(&encoded)?, header); + Ok(()) + } + + #[test] + fn database_header_rejects_corrupt_fields() -> io::Result<()> { + let mut page = golden_database_header(); + page[0] = b'X'; + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[4..8].copy_from_slice(&2u32.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[8..12].copy_from_slice(&8192u32.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[12] = 1; + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[68] = 1; + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[76] = 2; + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[64..68].copy_from_slice(&4u32.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[81] = 1; + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[40] ^= 1; + invalid_data(DatabaseHeader::decode(&page))?; + Ok(()) + } + + #[test] + fn database_header_rejects_invalid_structural_state() -> io::Result<()> { + for range in [16..24, 40..48, 48..64] { + let mut page = golden_database_header(); + page[range].fill(0); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + } + + let mut page = golden_database_header(); + page[24..32].copy_from_slice(&7u64.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let mut page = golden_database_header(); + page[32..40].copy_from_slice(&19u64.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + + let header = DatabaseHeader { + root_page_id: 7, + next_page_id: 19, + free_page_head: 7, + generation: 5, + database_id: [0x11; 16], + metadata: BPlusTreeMetadata::Empty, + }; + invalid_data(header.encode())?; + + let mut page = golden_database_header(); + page[32..40].copy_from_slice(&7u64.to_le_bytes()); + write_crc(&mut page, 72)?; + invalid_data(DatabaseHeader::decode(&page))?; + Ok(()) + } + + #[test] + fn leaf_page_header_golden_bytes_round_trip() -> io::Result<()> { + let expected = page_fixture(1, 2, 40, 4000, 6, 8, 0x88c9_dd55); + let header = PageHeader { + page_type: PageType::Leaf, + cell_count: 2, + free_start: 40, + free_end: 4000, + left: 6, + right: 8, + }; + let mut encoded = [0u8; PAGE_SIZE]; + header.encode_into(&mut encoded, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(encoded, expected); + assert_eq!(page_checksum(&encoded)?, 0x88c9_dd55); + assert_eq!(PageHeader::decode(&encoded, PAGE_ID, NEXT_PAGE_ID)?, header); + Ok(()) + } + + #[test] + fn internal_page_header_and_preamble_golden_bytes_round_trip() -> io::Result<()> { + let mut expected = page_fixture(2, 1, 44, 4000, 0, 0, 0x82b3_f78a); + expected[32..40].copy_from_slice(&3u64.to_le_bytes()); + let header = PageHeader { + page_type: PageType::Internal, + cell_count: 1, + free_start: 44, + free_end: 4000, + left: 0, + right: 0, + }; + let preamble = InternalPreamble { leftmost_child: 3 }; + let mut encoded = [0u8; PAGE_SIZE]; + preamble.encode_into(&mut encoded, PAGE_ID, NEXT_PAGE_ID)?; + header.encode_into(&mut encoded, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(encoded, expected); + assert_eq!(PageHeader::decode(&encoded, PAGE_ID, NEXT_PAGE_ID)?, header); + assert_eq!(InternalPreamble::decode(&encoded, PAGE_ID, NEXT_PAGE_ID)?, preamble); + Ok(()) + } + + #[test] + fn overflow_page_header_golden_bytes_round_trip() -> io::Result<()> { + let mut expected = page_fixture(3, 0, 0, 0, 0, 8, 0x8731_d58d); + expected[32..34].copy_from_slice(&3u16.to_le_bytes()); + expected[40..43].copy_from_slice(b"abc"); + let header = PageHeader { + page_type: PageType::Overflow, + cell_count: 0, + free_start: 0, + free_end: 0, + left: 0, + right: 8, + }; + let overflow = OverflowHeader { payload_length: 3 }; + let mut encoded = [0u8; PAGE_SIZE]; + encoded[40..43].copy_from_slice(b"abc"); + overflow.encode_into(&mut encoded)?; + header.encode_into(&mut encoded, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(encoded, expected); + assert_eq!(PageHeader::decode(&encoded, PAGE_ID, NEXT_PAGE_ID)?, header); + assert_eq!(OverflowHeader::decode(&encoded)?, overflow); + Ok(()) + } + + #[test] + fn free_page_header_golden_bytes_round_trip() -> io::Result<()> { + let expected = page_fixture(4, 0, 0, 0, 0, 8, 0x2864_f2d3); + let header = PageHeader { + page_type: PageType::Free, + cell_count: 0, + free_start: 0, + free_end: 0, + left: 0, + right: 8, + }; + let mut encoded = [0u8; PAGE_SIZE]; + encode_free_body(&mut encoded)?; + header.encode_into(&mut encoded, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(encoded, expected); + assert_eq!(PageHeader::decode(&encoded, PAGE_ID, NEXT_PAGE_ID)?, header); + validate_free_body(&encoded)?; + Ok(()) + } + + #[test] + fn page_header_rejects_unknown_type_flags_reserved_checksum_and_references() -> io::Result<()> { + let expected = page_fixture(1, 2, 40, 4000, 6, 8, 0x88c9_dd55); + + for (offset, value) in [(0, 9), (1, 1), (28, 1)] { + let mut page = expected; + page[offset] = value; + write_crc(&mut page, 24)?; + invalid_data(PageHeader::decode(&page, PAGE_ID, NEXT_PAGE_ID))?; + } + + let mut page = expected; + page[16..24].copy_from_slice(&PAGE_ID.to_le_bytes()); + write_crc(&mut page, 24)?; + invalid_data(PageHeader::decode(&page, PAGE_ID, NEXT_PAGE_ID))?; + + let mut page = expected; + page[2] ^= 1; + invalid_data(PageHeader::decode(&page, PAGE_ID, NEXT_PAGE_ID))?; + Ok(()) + } + + #[test] + fn overflow_and_free_bodies_reject_reserved_or_nonzero_tail_bytes() -> io::Result<()> { + let mut overflow = [0u8; PAGE_SIZE]; + overflow[32..34].copy_from_slice(&1u16.to_le_bytes()); + overflow[40] = 7; + overflow[34] = 1; + invalid_data(OverflowHeader::decode(&overflow))?; + + let mut overflow = [0u8; PAGE_SIZE]; + overflow[32..34].copy_from_slice(&1u16.to_le_bytes()); + overflow[40] = 7; + overflow[41] = 1; + invalid_data(OverflowHeader::decode(&overflow))?; + + let mut free = [0u8; PAGE_SIZE]; + free[32] = 1; + invalid_data(validate_free_body(&free))?; + Ok(()) + } + + #[test] + fn overflow_payload_accepts_4056_and_rejects_4057() -> io::Result<()> { + let mut page = [0u8; PAGE_SIZE]; + OverflowHeader { payload_length: 4056 }.encode_into(&mut page)?; + assert_eq!(OverflowHeader::decode(&page)?.payload_length, 4056); + + invalid_input(OverflowHeader { payload_length: 4057 }.encode_into(&mut page))?; + page[32..34].copy_from_slice(&4057u16.to_le_bytes()); + invalid_data(OverflowHeader::decode(&page))?; + Ok(()) + } + + #[test] + fn slot_and_cell_prefixes_have_golden_bytes() -> io::Result<()> { + let slot = Slot { offset: 0x1234, length: 0x5678 }; + let slot_bytes = [0x34, 0x12, 0x78, 0x56]; + assert_eq!(slot.encode(), slot_bytes); + assert_eq!(Slot::decode(&slot_bytes)?, slot); + + let internal = InternalCellPrefix { key_length: 3, right_child: 9 }; + let internal_bytes = [3, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0]; + assert_eq!(internal.encode(PAGE_ID, NEXT_PAGE_ID)?, internal_bytes); + assert_eq!(InternalCellPrefix::decode(&internal_bytes, PAGE_ID, NEXT_PAGE_ID)?, internal); + + let leaf = LeafCellPrefix { + key_length: 3, + value_kind: ValueKind::Inline, + compression: Compression::Lz4, + logical_length: 128, + stored_length: 16, + overflow_head: 0, + stored_crc32: 0x1122_3344, + }; + let leaf_bytes = [ + 3, 0, 0, 1, 128, 0, 0, 0, 16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x44, 0x33, 0x22, 0x11, + ]; + assert_eq!(leaf.encode(PAGE_ID, NEXT_PAGE_ID)?, leaf_bytes); + assert_eq!(LeafCellPrefix::decode(&leaf_bytes, PAGE_ID, NEXT_PAGE_ID)?, leaf); + assert_eq!(stored_value_checksum(b"abc"), 0x3524_41c2); + Ok(()) + } + + #[test] + fn cell_prefixes_reject_reserved_unknown_modes_and_invalid_lengths() -> io::Result<()> { + let mut internal = [3, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0]; + internal[2] = 1; + invalid_data(InternalCellPrefix::decode(&internal, PAGE_ID, NEXT_PAGE_ID))?; + + let mut leaf = [ + 3, 0, 0, 0, 8, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, + ]; + leaf[2] = 9; + invalid_data(LeafCellPrefix::decode(&leaf, PAGE_ID, NEXT_PAGE_ID))?; + leaf[2] = 0; + leaf[3] = 9; + invalid_data(LeafCellPrefix::decode(&leaf, PAGE_ID, NEXT_PAGE_ID))?; + leaf[3] = 0; + leaf[8..12].copy_from_slice(&257u32.to_le_bytes()); + invalid_data(LeafCellPrefix::decode(&leaf, PAGE_ID, NEXT_PAGE_ID))?; + Ok(()) + } + + #[test] + fn leaf_key_and_complete_cell_footprint_boundaries() -> io::Result<()> { + let overflow_at_key_limit = LeafCellPrefix { + key_length: 2004, + value_kind: ValueKind::Overflow, + compression: Compression::None, + logical_length: 4096, + stored_length: 4096, + overflow_head: 9, + stored_crc32: 0x1234_5678, + }; + overflow_at_key_limit.encode(PAGE_ID, NEXT_PAGE_ID)?; + + let mut overflow_above_key_limit = overflow_at_key_limit; + overflow_above_key_limit.key_length = 2005; + invalid_input(overflow_above_key_limit.encode(PAGE_ID, NEXT_PAGE_ID))?; + + let inline_2032 = LeafCellPrefix { + key_length: 1748, + value_kind: ValueKind::Inline, + compression: Compression::None, + logical_length: 256, + stored_length: 256, + overflow_head: 0, + stored_crc32: 0x1234_5678, + }; + inline_2032.encode(PAGE_ID, NEXT_PAGE_ID)?; + + let mut inline_2033 = inline_2032; + inline_2033.key_length = 1749; + invalid_input(inline_2033.encode(PAGE_ID, NEXT_PAGE_ID))?; + Ok(()) + } + + #[test] + fn known_size_prepended_lz4_block_decodes_to_golden_raw_bytes() -> io::Result<()> { + let block = [6, 0, 0, 0, 0x60, b'g', b'o', b'l', b'd', b'e', b'n']; + let mut scratch = Vec::new(); + assert_eq!(decompress_value_into(&block, 6, 6, &mut scratch)?, b"golden"); + Ok(()) + } + + #[test] + fn decompression_reuses_caller_scratch() -> io::Result<()> { + let block = [6, 0, 0, 0, 0x60, b'g', b'o', b'l', b'd', b'e', b'n']; + let mut scratch = Vec::with_capacity(64); + let initial_capacity = scratch.capacity(); + let initial_pointer = scratch.as_ptr(); + + assert_eq!(decompress_value_into(&block, 6, 64, &mut scratch)?, b"golden"); + assert_eq!(scratch.capacity(), initial_capacity); + assert_eq!(scratch.as_ptr(), initial_pointer); + + assert_eq!(decompress_value_into(&block, 6, 64, &mut scratch)?, b"golden"); + assert_eq!(scratch.capacity(), initial_capacity); + assert_eq!(scratch.as_ptr(), initial_pointer); + Ok(()) + } + + #[test] + fn compression_policy_obeys_size_and_ratio_boundaries() -> io::Result<()> { + let mut scratch = Vec::new(); + let short = [0u8; 63]; + let stored = encode_value(&short, &mut scratch)?; + assert_eq!(stored.compression(), Compression::None); + assert_eq!(stored.as_slice(), short); + + let incompressible = (0u8..64).collect::>(); + { + let stored = encode_value(&incompressible, &mut scratch)?; + assert_eq!(stored.compression(), Compression::None); + assert_eq!(stored.as_slice(), incompressible); + } + assert!(scratch.len() >= incompressible.len() * 85 / 100); + + let compressible = [0u8; 128]; + let expected = lz4_flex::compress_prepend_size(&compressible); + let stored = encode_value(&compressible, &mut scratch)?; + assert_eq!(stored.compression(), Compression::Lz4); + assert!(stored.as_slice().len() < compressible.len() * 85 / 100); + assert_eq!(stored.as_slice(), expected); + let mut decompression_scratch = Vec::new(); + assert_eq!( + decompress_value_into(stored.as_slice(), 128, 128, &mut decompression_scratch)?, + compressible + ); + Ok(()) + } + + #[test] + fn lz4_ratio_requires_strictly_less_than_85_percent() -> io::Result<()> { + assert!(compression_is_beneficial(100, 84)?); + assert!(!compression_is_beneficial(100, 85)?); + assert!(!compression_is_beneficial(100, 86)?); + Ok(()) + } + + #[test] + fn decompression_rejects_corrupt_sizes_before_allocation() -> io::Result<()> { + let mut scratch = Vec::new(); + let corrupt = [0xff, 0xff, 0xff, 0xff, 0x00]; + invalid_data(decompress_value_into(&corrupt, u32::MAX, 1024, &mut scratch))?; + assert_eq!(scratch.capacity(), 0); + + let block = [6, 0, 0, 0, 0x60, b'g', b'o', b'l', b'd', b'e', b'n']; + invalid_data(decompress_value_into(&block, 7, 7, &mut scratch))?; + assert_eq!(scratch.capacity(), 0); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree/v3/migration.rs b/backend/src/repository/bplustree/v3/migration.rs new file mode 100644 index 000000000..65075966c --- /dev/null +++ b/backend/src/repository/bplustree/v3/migration.rs @@ -0,0 +1,409 @@ +use super::{BPlusTree, BPlusTreeMetadata}; +use crate::repository::bplustree::v2; +use serde::{Deserialize, Serialize}; +use std::{ + fs::{File, OpenOptions}, + io::{self, Read, Seek, SeekFrom, Write}, + ops::Bound, + path::{Path, PathBuf}, +}; + +const LEGACY_V1: u32 = 1; +const LEGACY_V2: u32 = 2; +const METADATA_LEN_OFFSET: u64 = 16; +const METADATA_MAX_SIZE: u32 = 4000; +const HEADER_FLAG_HAS_METADATA_FLAGS: u32 = 1 << 31; +const HEADER_FLAG_HAS_TOMBSTONES: u32 = 1 << 30; +const HEADER_METADATA_LEN_MASK: u32 = !(HEADER_FLAG_HAS_METADATA_FLAGS | HEADER_FLAG_HAS_TOMBSTONES); + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct MigrationValidation { + pub(crate) entries: usize, + pub(crate) database_id: [u8; 16], + pub(crate) generation: u64, +} + +pub(crate) fn storage_version(path: &Path) -> io::Result> { + let mut file = File::open(path)?; + if file.metadata()?.len() < 8 { + return Ok(None); + } + let mut header = [0; 8]; + file.read_exact(&mut header)?; + if &header[0..4] != b"BTRE" { + return Ok(None); + } + Ok(Some(u32::from_le_bytes( + header[4..8].try_into().map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?, + ))) +} + +fn migration_source(path: &Path, version: u32) -> io::Result> { + if version == LEGACY_V2 { + return Ok(None); + } + if version != LEGACY_V1 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "unsupported legacy B+Tree version")); + } + let name = path.file_name().ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "database has no name"))?; + let temporary = path.with_file_name(format!("{}.{}.v3.tmp", name.to_string_lossy(), uuid::Uuid::new_v4())); + std::fs::copy(path, &temporary)?; + let normalized = normalize_v1_copy(&temporary); + if let Err(error) = normalized { + let _ = std::fs::remove_file(&temporary); + return Err(error); + } + Ok(Some(temporary)) +} + +fn normalize_v1_copy(path: &Path) -> io::Result<()> { + let mut file = OpenOptions::new().read(true).write(true).open(path)?; + file.seek(SeekFrom::Start(METADATA_LEN_OFFSET))?; + let mut encoded = [0; 4]; + file.read_exact(&mut encoded)?; + let raw = u32::from_le_bytes(encoded); + let metadata_len = raw & HEADER_METADATA_LEN_MASK; + if metadata_len > METADATA_MAX_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy metadata exceeds header capacity")); + } + let normalized = (metadata_len | HEADER_FLAG_HAS_METADATA_FLAGS) & !HEADER_FLAG_HAS_TOMBSTONES; + file.seek(SeekFrom::Start(METADATA_LEN_OFFSET))?; + file.write_all(&normalized.to_le_bytes())?; + file.seek(SeekFrom::Start(4))?; + file.write_all(&LEGACY_V2.to_le_bytes())?; + file.sync_all() +} + +fn legacy_metadata(path: &Path) -> io::Result { + let mut file = File::open(path)?; + file.seek(SeekFrom::Start(METADATA_LEN_OFFSET))?; + let mut encoded_len = [0; 4]; + file.read_exact(&mut encoded_len)?; + let metadata_len = u32::from_le_bytes(encoded_len) & HEADER_METADATA_LEN_MASK; + if metadata_len > METADATA_MAX_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy metadata exceeds header capacity")); + } + let metadata_len = usize::try_from(metadata_len).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + let mut encoded = vec![0; metadata_len]; + file.read_exact(&mut encoded)?; + match encoded.as_slice() { + [] => Ok(BPlusTreeMetadata::Empty), + [1, value0, value1, value2, value3] => Ok(BPlusTreeMetadata::TargetIdMapping(u32::from_le_bytes([ + *value0, *value1, *value2, *value3, + ]))), + _ => Err(io::Error::new(io::ErrorKind::InvalidData, "unsupported legacy B+Tree metadata")), + } +} + +pub(crate) fn migrate_v2_typed(source: &Path) -> io::Result +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, +{ + migrate_v2_typed_inner::( + source, + std::convert::identity, + BPlusTree::store_verified, + |_, _| Ok(()), + false, + ) +} + +pub(crate) fn migrate_v2_typed_with_index( + source: &Path, + sort_key: F, +) -> io::Result +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, + SortKey: Ord + Serialize + for<'de> Deserialize<'de>, + F: Fn(&V) -> SortKey, +{ + migrate_v2_typed_inner::( + source, + std::convert::identity, + |tree, destination| tree.store_with_index_verified(destination, sort_key), + |destination, entries| { + let query = super::BPlusTreeQuery::::try_new(destination)?; + let index = crate::repository::storage::get_file_path_for_db_index(destination); + let mut iterator = crate::repository::bplustree::sorted_index::v4::OwnedIterator::::open(query, &index)?; + let indexed_entries = iterator.try_fold(0usize, |count, entry| { + let _ = entry?; + count + .checked_add(1) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "index validation count overflow")) + })?; + if indexed_entries != entries { + return Err(io::Error::new(io::ErrorKind::InvalidData, "migrated sorted-index entry count mismatch")); + } + Ok(()) + }, + true, + ) +} + +pub(crate) fn migrate_v2_typed_map( + source: &Path, + map: Map, +) -> io::Result +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + SourceV: Serialize + for<'de> Deserialize<'de> + Clone, + DestinationV: Serialize + for<'de> Deserialize<'de> + Clone, + Map: FnMut(SourceV) -> DestinationV, +{ + migrate_v2_typed_inner::( + source, + map, + BPlusTree::store_verified, + |_, _| Ok(()), + false, + ) +} + +fn migrate_v2_typed_inner( + source: &Path, + mut map: Map, + store: Store, + validate: Validate, + indexed: bool, +) -> io::Result +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + SourceV: Serialize + for<'de> Deserialize<'de> + Clone, + DestinationV: Serialize + for<'de> Deserialize<'de> + Clone, + Map: FnMut(SourceV) -> DestinationV, + Store: FnOnce(&mut BPlusTree, &Path) -> io::Result, + Validate: FnOnce(&Path, usize) -> io::Result<()>, +{ + let version = storage_version(source)?.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "not a B+Tree"))?; + let normalized_source = migration_source(source, version)?; + let read_path = normalized_source.as_deref().unwrap_or(source); + let destination = migration_destination(source)?; + let destination_index = crate::repository::storage::get_file_path_for_db_index(&destination); + let converted = (|| { + let mut legacy = v2::BPlusTreeQuery::::try_new(read_path)?; + let metadata = legacy_metadata(read_path)?; + let expected_metadata = metadata.clone(); + let mut tree = BPlusTree::new(); + tree.set_metadata(metadata); + let mut entries = 0usize; + for entry in legacy.range_iter(Bound::Unbounded, Bound::Unbounded) { + let (key, value) = entry.map_err(v2::BPlusTreeError::to_io)?; + tree.insert(key, map(value)); + entries = entries + .checked_add(1) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "migration entry count overflow"))?; + } + drop(legacy); + if tree.len() != entries { + return Err(io::Error::new(io::ErrorKind::InvalidData, "legacy migration produced duplicate keys")); + } + let verification = store(&mut tree, &destination)?; + if verification.live_entries + != u64::try_from(entries) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "migration entry count exceeds u64"))? + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "migrated database entry count differs from full verification", + )); + } + let query = super::BPlusTreeQuery::::try_new(&destination)?; + let (database_id, generation) = query.snapshot_identity(); + if database_id == [0; 16] { + return Err(io::Error::new(io::ErrorKind::InvalidData, "migrated database identity is zero")); + } + if generation != 1 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "migrated database generation is not one")); + } + if query.snapshot_metadata() != &expected_metadata { + return Err(io::Error::new(io::ErrorKind::InvalidData, "migrated database metadata mismatch")); + } + drop(query); + validate(&destination, entries)?; + if indexed { + publish(&destination_index, &crate::repository::storage::get_file_path_for_db_index(source))?; + } + publish(&destination, source)?; + Ok(MigrationValidation { entries, database_id, generation }) + })(); + if let Some(temporary) = normalized_source { + let _ = std::fs::remove_file(temporary); + } + cleanup_destination(&destination); + converted +} + +fn migration_destination(path: &Path) -> io::Result { + let name = path.file_name().ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "database has no name"))?; + Ok(path.with_file_name(format!("{}.{}.v3.tmp", name.to_string_lossy(), uuid::Uuid::new_v4()))) +} + +fn publish(temporary: &Path, destination: &Path) -> io::Result<()> { + let temporary = tempfile::TempPath::try_from_path(temporary)?; + temporary.persist(destination).map_err(io::Error::from)?; + super::wal::sync_parent_directory(destination) +} + +fn cleanup_destination(destination: &Path) { + let _ = std::fs::remove_file(destination); + let _ = std::fs::remove_file(crate::repository::storage::get_file_path_for_db_index(destination)); + let _ = std::fs::remove_file(crate::repository::bplustree::common::sidecar_lock_path(destination)); +} + +#[cfg(test)] +mod tests { + use super::*; + use super::super::page::{page_open_count, reset_page_open_count}; + + #[test] + fn typed_migration_validates_the_v3_destination_once() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("single-pass-v2.db"); + let baseline_path = dir.path().join("single-pass-v3-baseline.db"); + let mut legacy = v2::BPlusTree::new(); + legacy.insert(1u32, String::from("one")); + legacy.store(&path)?; + + let mut baseline = BPlusTree::new(); + baseline.insert(1u32, String::from("one")); + reset_page_open_count(); + baseline.store(&baseline_path)?; + let single_verified_store_opens = page_open_count(); + + reset_page_open_count(); + let validation = migrate_v2_typed::(&path)?; + assert_eq!(validation.entries, 1); + assert_eq!(page_open_count(), single_verified_store_opens); + Ok(()) + } + + #[test] + fn typed_v2_and_v1_sources_become_verified_v3_databases() -> io::Result<()> { + for version in [LEGACY_V2, LEGACY_V1] { + let dir = tempfile::tempdir()?; + let path = dir.path().join(format!("source-v{version}.db")); + let mut legacy = v2::BPlusTree::new(); + legacy.insert(1u32, String::from("one")); + legacy.insert(2u32, String::from("two")); + legacy.store(&path)?; + if version == LEGACY_V1 { + let mut file = OpenOptions::new().write(true).open(&path)?; + file.seek(SeekFrom::Start(4))?; + file.write_all(&LEGACY_V1.to_le_bytes())?; + file.sync_all()?; + } + + let validation = migrate_v2_typed::(&path)?; + assert_eq!(validation.entries, 2); + assert_eq!(validation.generation, 1); + assert_eq!(storage_version(&path)?, Some(3)); + let mut query = super::super::BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.iter().collect::>>()?, vec![(1, String::from("one")), (2, String::from("two"))]); + } + Ok(()) + } + + #[test] + fn typed_migration_accepts_verified_historical_fence_key_nodes() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("historical-fence-v2.db"); + let mut legacy = v2::BPlusTree::new_with_orders(2, 2); + for key in 0..80u32 { + legacy.insert(key, format!("value-{key}")); + } + assert!(legacy.add_historical_fence_key()); + legacy.store(&path)?; + + let validation = migrate_v2_typed::(&path)?; + assert_eq!(validation.entries, 80); + let mut query = super::super::BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.iter().collect::>>()?.len(), 80); + Ok(()) + } + + #[test] + fn typed_migration_accepts_unaligned_cow_root() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("unaligned-root-v2.db"); + let mut legacy = v2::BPlusTree::new_with_orders(4, 4); + for key in 0..40u32 { + legacy.insert(key, format!("value-{key}")); + } + legacy.store(&path)?; + + let mut bytes = std::fs::read(&path)?; + let root_offset = u64::from_le_bytes( + bytes[8..16].try_into().map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?, + ); + let root_start = usize::try_from(root_offset).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + let root_end = root_start + .checked_add(v2::PAGE_SIZE_USIZE) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "fixture root end overflow"))?; + let root = bytes + .get(root_start..root_end) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "fixture root is truncated"))? + .to_vec(); + bytes.push(0); + let relocated_root = u64::try_from(bytes.len()).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + bytes.extend_from_slice(&root); + bytes[8..16].copy_from_slice(&relocated_root.to_le_bytes()); + std::fs::write(&path, bytes)?; + + let validation = migrate_v2_typed::(&path)?; + assert_eq!(validation.entries, 40); + Ok(()) + } + + #[test] + fn typed_migration_rejects_missing_root_child_as_corruption() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("missing-root-child-v2.db"); + let mut legacy = v2::BPlusTree::new_with_orders(2, 2); + for key in 0..20u32 { + legacy.insert(key, format!("value-{key}")); + } + assert!(legacy.remove_last_root_child()); + legacy.store(&path)?; + + assert!(migrate_v2_typed::(&path).is_err()); + assert_eq!(storage_version(&path)?, Some(LEGACY_V2)); + Ok(()) + } + + #[test] + fn typed_migration_read_failure_keeps_the_legacy_source() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("corrupt-v2.db"); + let mut legacy = v2::BPlusTree::new(); + legacy.insert(1u32, String::from("one")); + legacy.store(&path)?; + let mut corrupted = std::fs::read(&path)?; + corrupted.truncate(v2::PAGE_SIZE_USIZE + 8); + std::fs::write(&path, &corrupted)?; + + assert!(migrate_v2_typed::(&path).is_err()); + assert_eq!(std::fs::read(&path)?, corrupted); + Ok(()) + } + + #[test] + fn typed_migration_rebuilds_and_validates_the_sorted_index() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("indexed-v2.db"); + let mut legacy = v2::BPlusTree::new(); + legacy.insert(1u32, String::from("bbb")); + legacy.insert(2u32, String::from("a")); + legacy.store(&path)?; + + let validation = migrate_v2_typed_with_index::(&path, String::len)?; + assert_eq!(validation.entries, 2); + let query = super::super::BPlusTreeQuery::::try_new(&path)?; + let index = crate::repository::storage::get_file_path_for_db_index(&path); + let values = crate::repository::bplustree::sorted_index::v4::OwnedIterator::::open(query, &index) + ?.collect::>>()?; + assert_eq!(values, vec![(2, String::from("a")), (1, String::from("bbb"))]); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree/v3/mod.rs b/backend/src/repository/bplustree/v3/mod.rs new file mode 100644 index 000000000..d7d772b5e --- /dev/null +++ b/backend/src/repository/bplustree/v3/mod.rs @@ -0,0 +1,140 @@ +mod format; +pub(crate) mod migration; +mod page; +mod tree; +mod wal; + +#[allow(unused_imports)] +pub use tree::{ + BPlusTree, BPlusTreeDiskIterator, BPlusTreeDiskIteratorOwned, BPlusTreeQuery, BPlusTreeRangeIterator, + BPlusTreeSerialWriter, BPlusTreeUpdate, FlushPolicy, +}; +pub(crate) use format::{Locator, MAGIC, STORAGE_VERSION_V3 as STORAGE_VERSION}; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum BPlusTreeMetadata { + Empty, + TargetIdMapping(u32), +} + +#[cfg(test)] +mod contract_tests { + use super::{tree::verify_full, BPlusTree, BPlusTreeMetadata, BPlusTreeQuery}; + use crate::repository::bplustree::v2; + use std::{io, ops::Bound}; + + fn incompressible_value() -> Vec { + let mut state = 0x1234_5678_9abc_def0u64; + (0..12_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.to_le_bytes()[0] + }) + .collect() + } + + macro_rules! storage_contract { + ($name:ident, $tree:ty, $query:ty, $metadata:expr) => { + #[test] + fn $name() -> io::Result<()> { + type Tree = $tree; + type Query = $query; + + let dir = tempfile::tempdir()?; + let empty_path = dir.path().join("empty.db"); + let mut empty = Tree::new(); + empty.store(&empty_path)?; + let mut empty_query = Query::try_new(&empty_path)?; + assert!(empty_query.is_empty().map_err(|error| error.to_io())?); + assert!(empty_query.iter().collect::>>()?.is_empty()); + + let path = dir.path().join("contract.db"); + let expected = vec![(10, vec![10]), (20, vec![20, 20]), (30, incompressible_value())]; + let mut tree = Tree::new(); + tree.set_metadata($metadata); + for (key, value) in &expected { + tree.insert(*key, value.clone()); + } + assert_eq!(tree.query(&20), Some(&vec![20, 20])); + assert_eq!(tree.find_le(&25).map(|(key, _)| *key), Some(20)); + tree.store(&path)?; + + let loaded = Tree::load(&path)?; + assert_eq!(loaded.get_metadata(), &$metadata); + assert_eq!(loaded.iter().map(|(key, value)| (*key, value.clone())).collect::>(), expected); + + let mut query = Query::try_new(&path)?; + assert_eq!(query.query_zero_copy(&20).map_err(|error| error.to_io())?, Some(vec![20, 20])); + assert_eq!(query.query_le(&25).map_err(|error| error.to_io())?, Some(vec![20, 20])); + assert_eq!(query.iter().collect::>>()?, expected); + + let mut query = Query::try_new(&path)?; + assert_eq!( + query + .range_iter(Bound::Included(&20), Bound::Included(&30)) + .collect::>>()?, + expected[1..].to_vec() + ); + Ok(()) + } + }; + } + + storage_contract!( + v3_storage_contract, + BPlusTree>, + BPlusTreeQuery>, + BPlusTreeMetadata::TargetIdMapping(42) + ); + + #[test] + fn v2_migration_reader_contract() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("v2.db"); + let expected = vec![(10, vec![10]), (20, vec![20, 20]), (30, incompressible_value())]; + let mut tree = v2::BPlusTree::new(); + tree.set_metadata(v2::BPlusTreeMetadata::TargetIdMapping(42)); + for (key, value) in &expected { + tree.insert(*key, value.clone()); + } + tree.store(&path)?; + + let mut query = v2::BPlusTreeQuery::>::try_new(&path)?; + let actual = query + .range_iter(Bound::Unbounded, Bound::Unbounded) + .collect::, _>>() + .map_err(v2::BPlusTreeError::to_io)?; + assert_eq!(actual, expected); + Ok(()) + } + + #[test] + fn failed_v3_store_preserves_published_database() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("preserved.db"); + let mut tree = BPlusTree::::new(); + tree.insert("published".into(), 7); + tree.store(&path)?; + tree.insert("x".repeat(2_100), 8); + assert!(tree.store(&path).is_err()); + + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&"published".into()).map_err(v2::BPlusTreeError::to_io)?, Some(7)); + assert_eq!(query.query(&"x".repeat(2_100)).map_err(v2::BPlusTreeError::to_io)?, None); + Ok(()) + } + + #[test] + fn v3_full_verification_contract() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("verified.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, incompressible_value()); + tree.store(&path)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(verify_full(&mut query)?.live_entries, 1); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree/v3/page.rs b/backend/src/repository/bplustree/v3/page.rs new file mode 100644 index 000000000..f6e1cd3d6 --- /dev/null +++ b/backend/src/repository/bplustree/v3/page.rs @@ -0,0 +1,689 @@ +use super::format::{ + encode_free_body, validate_free_body, write_page_checksum, InternalPreamble, OverflowHeader, PageHeader, PageType, + Slot, INTERNAL_PREAMBLE_LEN, OVERFLOW_PAYLOAD_LEN, PAGE_HEADER_LEN, PAGE_SIZE, SLOT_LEN, +}; +use std::{io, iter::FusedIterator}; + +#[cfg(test)] +thread_local! { + static PAGE_OPEN_COUNT: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +#[cfg(test)] +pub(crate) fn reset_page_open_count() { PAGE_OPEN_COUNT.set(0); } + +#[cfg(test)] +pub(crate) fn page_open_count() -> usize { PAGE_OPEN_COUNT.get() } + +fn invalid_data(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidData, message) } + +fn invalid_input(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message) } + +fn checked_end(offset: usize, length: usize, kind: io::ErrorKind) -> io::Result { + offset.checked_add(length).ok_or_else(|| io::Error::new(kind, "page offset overflow")) +} + +fn slot_base(page_type: PageType) -> io::Result { + match page_type { + PageType::Leaf => Ok(PAGE_HEADER_LEN), + PageType::Internal => PAGE_HEADER_LEN + .checked_add(INTERNAL_PREAMBLE_LEN) + .ok_or_else(|| invalid_data("slot base overflow")), + PageType::Overflow | PageType::Free => Err(invalid_input("chain pages do not contain slots")), + } +} + +pub(crate) struct SlottedPage { + bytes: B, + header: PageHeader, + page_id: u64, + next_page_id: u64, + overflow_payload_length: Option, +} + +#[derive(Clone, Copy)] +pub(crate) struct PageValidation { + header: PageHeader, + page_id: u64, + next_page_id: u64, + overflow_payload_length: Option, +} + +impl> SlottedPage { + pub(crate) fn open(bytes: B, page_id: u64, next_page_id: u64) -> io::Result { + #[cfg(test)] + PAGE_OPEN_COUNT.with(|count| count.set(count.get().saturating_add(1))); + let page = bytes.as_ref(); + if page.len() != PAGE_SIZE { + return Err(invalid_data("page must be exactly 4096 bytes")); + } + let header = PageHeader::decode(page, page_id, next_page_id)?; + let overflow_payload_length = match header.page_type { + PageType::Internal => { + InternalPreamble::decode(page, page_id, next_page_id)?; + None + } + PageType::Overflow => Some(OverflowHeader::decode(page)?.payload_length), + PageType::Free => { + validate_free_body(page)?; + None + } + PageType::Leaf => None, + }; + let slotted = Self { bytes, header, page_id, next_page_id, overflow_payload_length }; + slotted.validate_slots()?; + Ok(slotted) + } + + pub(crate) fn from_immutable_snapshot(bytes: B, validation: PageValidation) -> io::Result { + if bytes.as_ref().len() != PAGE_SIZE { + return Err(invalid_data("page must be exactly 4096 bytes")); + } + Ok(Self { + bytes, + header: validation.header, + page_id: validation.page_id, + next_page_id: validation.next_page_id, + overflow_payload_length: validation.overflow_payload_length, + }) + } + + pub(crate) const fn validation(&self) -> PageValidation { + PageValidation { + header: self.header, + page_id: self.page_id, + next_page_id: self.next_page_id, + overflow_payload_length: self.overflow_payload_length, + } + } + + pub(crate) const fn header(&self) -> PageHeader { self.header } + + pub(crate) const fn page_id(&self) -> u64 { self.page_id } + + pub(crate) const fn next_page_id(&self) -> u64 { self.next_page_id } + + pub(crate) fn as_bytes(&self) -> &[u8] { self.bytes.as_ref() } + + fn slot(&self, index: usize) -> io::Result { + if index >= usize::from(self.header.cell_count) { + return Err(invalid_input("cell index is outside page")); + } + let base = slot_base(self.header.page_type)?; + let offset = index + .checked_mul(SLOT_LEN) + .and_then(|size| base.checked_add(size)) + .ok_or_else(|| invalid_data("slot offset overflow"))?; + let end = checked_end(offset, SLOT_LEN, io::ErrorKind::InvalidData)?; + Slot::decode( + self.bytes + .as_ref() + .get(offset..end) + .ok_or_else(|| invalid_data("slot is outside page"))?, + ) + } + + pub(crate) fn cell(&self, index: usize) -> io::Result<&[u8]> { + let range = self.cell_range(index)?; + self.bytes + .as_ref() + .get(range) + .ok_or_else(|| invalid_data("cell is outside page")) + } + + pub(crate) fn cell_range(&self, index: usize) -> io::Result> { + let slot = self.slot(index)?; + let offset = usize::from(slot.offset); + let end = checked_end(offset, usize::from(slot.length), io::ErrorKind::InvalidData)?; + Ok(offset..end) + } + + pub(crate) fn cells(&self) -> Cells<'_, B> { Cells { page: self, next: 0 } } + + fn validate_slots(&self) -> io::Result<()> { + if matches!(self.header.page_type, PageType::Overflow | PageType::Free) { + return Ok(()); + } + if self.header.cell_count == 0 { + return Ok(()); + } + + let free_start = usize::from(self.header.free_start); + let mut previous_offset = PAGE_SIZE; + for index in 0..usize::from(self.header.cell_count) { + let slot = self.slot(index)?; + let offset = usize::from(slot.offset); + let length = usize::from(slot.length); + let end = checked_end(offset, length, io::ErrorKind::InvalidData)?; + if length == 0 + || offset < free_start + || end > PAGE_SIZE + || offset >= previous_offset + || end > previous_offset + { + return Err(invalid_data("invalid or overlapping cell range")); + } + previous_offset = offset; + } + if previous_offset != usize::from(self.header.free_end) { + return Err(invalid_data("free_end does not match the lowest cell offset")); + } + Ok(()) + } +} + +pub(crate) fn encode_overflow_page( + page_id: u64, + next_page_id: u64, + next: u64, + payload: &[u8], +) -> io::Result<[u8; PAGE_SIZE]> { + if payload.len() > OVERFLOW_PAYLOAD_LEN { + return Err(invalid_input("overflow payload is too large")); + } + let mut page = [0; PAGE_SIZE]; + let payload_end = checked_end(40, payload.len(), io::ErrorKind::InvalidInput)?; + page.get_mut(40..payload_end) + .ok_or_else(|| invalid_input("overflow payload exceeds page"))? + .copy_from_slice(payload); + OverflowHeader { + payload_length: u16::try_from(payload.len()).map_err(|_| invalid_input("overflow payload exceeds u16"))?, + } + .encode_into(&mut page)?; + PageHeader { + page_type: PageType::Overflow, + cell_count: 0, + free_start: 0, + free_end: 0, + left: 0, + right: next, + } + .encode_into(&mut page, page_id, next_page_id)?; + Ok(page) +} + +pub(crate) fn encode_free_page(page_id: u64, next_page_id: u64, next: u64) -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0; PAGE_SIZE]; + encode_free_body(&mut page)?; + PageHeader { + page_type: PageType::Free, + cell_count: 0, + free_start: 0, + free_end: 0, + left: 0, + right: next, + } + .encode_into(&mut page, page_id, next_page_id)?; + Ok(page) +} + +pub(crate) fn overflow_payload>(page: &SlottedPage) -> io::Result<&[u8]> { + let payload_length = page.overflow_payload_length.ok_or_else(|| invalid_data("expected overflow page"))?; + let end = checked_end(40, usize::from(payload_length), io::ErrorKind::InvalidData)?; + page.bytes.as_ref().get(40..end).ok_or_else(|| invalid_data("truncated overflow payload")) +} + +pub(crate) struct Cells<'a, B> { + page: &'a SlottedPage, + next: usize, +} + +impl<'a, B: AsRef<[u8]>> Iterator for Cells<'a, B> { + type Item = io::Result<&'a [u8]>; + + fn next(&mut self) -> Option { + if self.next >= usize::from(self.page.header.cell_count) { + return None; + } + let index = self.next; + self.next += 1; + Some(self.page.cell(index)) + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.len(); + (remaining, Some(remaining)) + } +} + +impl> ExactSizeIterator for Cells<'_, B> { + fn len(&self) -> usize { usize::from(self.page.header.cell_count).saturating_sub(self.next) } +} + +impl> FusedIterator for Cells<'_, B> {} + +impl + AsMut<[u8]>> SlottedPage { + pub(crate) fn rebuild_ordered<'a, I>(&mut self, cells: I) -> io::Result<()> + where + I: IntoIterator, + { + let base = slot_base(self.header.page_type)?; + let mut rebuilt = [0; PAGE_SIZE]; + if self.header.page_type == PageType::Internal { + let preamble_end = checked_end(PAGE_HEADER_LEN, INTERNAL_PREAMBLE_LEN, io::ErrorKind::InvalidData)?; + let preamble = self + .bytes + .as_ref() + .get(PAGE_HEADER_LEN..preamble_end) + .ok_or_else(|| invalid_data("internal preamble is outside page"))?; + rebuilt + .get_mut(PAGE_HEADER_LEN..preamble_end) + .ok_or_else(|| invalid_data("internal preamble destination is outside page"))? + .copy_from_slice(preamble); + } + + let mut count = 0u16; + let mut cell_start = PAGE_SIZE; + for cell in cells { + if cell.is_empty() { + return Err(invalid_input("cells must not be empty")); + } + let length = u16::try_from(cell.len()).map_err(|_| invalid_input("cell length exceeds u16"))?; + let next_count = count.checked_add(1).ok_or_else(|| invalid_input("cell count exceeds u16"))?; + let slot_offset = usize::from(count) + .checked_mul(SLOT_LEN) + .and_then(|size| base.checked_add(size)) + .ok_or_else(|| invalid_input("slot offset overflow"))?; + let slot_end = checked_end(slot_offset, SLOT_LEN, io::ErrorKind::InvalidInput)?; + let next_cell_start = cell_start + .checked_sub(cell.len()) + .ok_or_else(|| invalid_input("cells exceed page capacity"))?; + if next_cell_start < slot_end { + return Err(invalid_input("cells exceed page capacity")); + } + rebuilt + .get_mut(next_cell_start..cell_start) + .ok_or_else(|| invalid_input("cell destination is outside page"))? + .copy_from_slice(cell); + let slot = Slot { + offset: u16::try_from(next_cell_start).map_err(|_| invalid_input("cell offset exceeds u16"))?, + length, + }; + rebuilt + .get_mut(slot_offset..slot_end) + .ok_or_else(|| invalid_input("slot destination is outside page"))? + .copy_from_slice(&slot.encode()); + count = next_count; + cell_start = next_cell_start; + } + + if self.header.page_type == PageType::Internal && count == 0 { + return Err(invalid_input("internal pages must not be empty")); + } + let slot_bytes = usize::from(count) + .checked_mul(SLOT_LEN) + .ok_or_else(|| invalid_input("slot directory size overflow"))?; + let free_start = base.checked_add(slot_bytes).ok_or_else(|| invalid_input("slot directory overflow"))?; + let header = PageHeader { + page_type: self.header.page_type, + cell_count: count, + free_start: u16::try_from(free_start).map_err(|_| invalid_input("slot directory exceeds u16"))?, + free_end: u16::try_from(cell_start).map_err(|_| invalid_input("cell offset exceeds u16"))?, + left: self.header.left, + right: self.header.right, + }; + header.encode_into(&mut rebuilt, self.page_id, self.next_page_id)?; + self.bytes + .as_mut() + .get_mut(..PAGE_SIZE) + .ok_or_else(|| invalid_data("page destination is truncated"))? + .copy_from_slice(&rebuilt); + self.header = header; + Ok(()) + } + + pub(crate) fn replace_same_len(&mut self, index: usize, replacement: &[u8]) -> io::Result<()> { + let slot = self.slot(index)?; + if replacement.len() != usize::from(slot.length) { + return Err(invalid_input("replacement length differs from cell length")); + } + let offset = usize::from(slot.offset); + let end = checked_end(offset, replacement.len(), io::ErrorKind::InvalidData)?; + let page = self.bytes.as_mut(); + if page.len() != PAGE_SIZE { + return Err(invalid_data("page must be exactly 4096 bytes")); + } + page + .get_mut(offset..end) + .ok_or_else(|| invalid_data("cell destination is outside page"))? + .copy_from_slice(replacement); + write_page_checksum(page) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::repository::bplustree::v3::format::{ + encode_free_body, page_checksum, write_page_checksum, InternalPreamble, OverflowHeader, PageHeader, PageType, + Slot, PAGE_HEADER_LEN, PAGE_SIZE, SLOT_LEN, + }; + use std::io; + + const PAGE_ID: u64 = 1; + const NEXT_PAGE_ID: u64 = 8; + + struct ChangingMutView { + bytes: [u8; PAGE_SIZE], + calls: usize, + } + + impl AsRef<[u8]> for ChangingMutView { + fn as_ref(&self) -> &[u8] { &self.bytes } + } + + impl AsMut<[u8]> for ChangingMutView { + fn as_mut(&mut self) -> &mut [u8] { + self.calls += 1; + if self.calls == 1 { &mut self.bytes[..200] } else { &mut self.bytes[..32] } + } + } + + fn invalid_data(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidData => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidData, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidData")), + } + } + + fn invalid_input(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidInput => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidInput, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidInput")), + } + } + + fn empty_leaf() -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0; PAGE_SIZE]; + PageHeader { + page_type: PageType::Leaf, + cell_count: 0, + free_start: u16::try_from(PAGE_HEADER_LEN).map_err(io::Error::other)?, + free_end: u16::try_from(PAGE_SIZE).map_err(io::Error::other)?, + left: 0, + right: 0, + } + .encode_into(&mut page, PAGE_ID, NEXT_PAGE_ID)?; + Ok(page) + } + + fn leaf_with_slots(slots: &[Slot], free_end: u16) -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0; PAGE_SIZE]; + for (index, slot) in slots.iter().copied().enumerate() { + let offset = PAGE_HEADER_LEN + .checked_add(index.checked_mul(SLOT_LEN).ok_or_else(|| io::Error::other("test slot overflow"))?) + .ok_or_else(|| io::Error::other("test slot overflow"))?; + let end = offset.checked_add(SLOT_LEN).ok_or_else(|| io::Error::other("test slot overflow"))?; + page.get_mut(offset..end) + .ok_or_else(|| io::Error::other("test slot outside page"))? + .copy_from_slice(&slot.encode()); + } + let cell_count = u16::try_from(slots.len()).map_err(io::Error::other)?; + let free_start = PAGE_HEADER_LEN + .checked_add(slots.len().checked_mul(SLOT_LEN).ok_or_else(|| io::Error::other("test slot overflow"))?) + .and_then(|value| u16::try_from(value).ok()) + .ok_or_else(|| io::Error::other("test slot overflow"))?; + PageHeader { + page_type: PageType::Leaf, + cell_count, + free_start, + free_end, + left: 0, + right: 0, + } + .encode_into(&mut page, PAGE_ID, NEXT_PAGE_ID)?; + Ok(page) + } + + fn chain_page(page_type: PageType) -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0; PAGE_SIZE]; + match page_type { + PageType::Overflow => OverflowHeader { payload_length: 0 }.encode_into(&mut page)?, + PageType::Free => encode_free_body(&mut page)?, + PageType::Leaf | PageType::Internal => return Err(io::Error::other("test requires a chain page")), + } + PageHeader { + page_type, + cell_count: 0, + free_start: 0, + free_end: 0, + left: 0, + right: 0, + } + .encode_into(&mut page, PAGE_ID, NEXT_PAGE_ID)?; + Ok(page) + } + + fn set_right_reference(page: &mut [u8], reference: u64) -> io::Result<()> { + page.get_mut(16..24) + .ok_or_else(|| io::Error::other("missing right reference"))? + .copy_from_slice(&reference.to_le_bytes()); + write_page_checksum(page) + } + + #[test] + fn opens_empty_leaf_from_immutable_and_mutable_buffers() -> io::Result<()> { + let page = empty_leaf()?; + let immutable = SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(immutable.header().page_type, PageType::Leaf); + assert_eq!(immutable.cells().len(), 0); + invalid_input(immutable.cell(0))?; + + let mut page = page; + let mut mutable = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + mutable.rebuild_ordered([b"a".as_slice(), b"bc".as_slice()])?; + assert_eq!(mutable.cells().collect::>>()?, [b"a".as_slice(), b"bc".as_slice()]); + Ok(()) + } + + #[test] + fn rebuild_packs_slots_in_descending_offset_order_and_exactly_fits() -> io::Result<()> { + let mut page = empty_leaf()?; + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + let first = [1; 2028]; + let second = [2; 2028]; + slotted.rebuild_ordered([first.as_slice(), second.as_slice()])?; + + assert_eq!(slotted.header().free_start, 40); + assert_eq!(slotted.header().free_end, 40); + assert_eq!(slotted.cell(0)?, first); + assert_eq!(slotted.cell(1)?, second); + assert_eq!(slotted.slot(0)?.offset, 2068); + assert_eq!(slotted.slot(1)?.offset, 40); + Ok(()) + } + + #[test] + fn rebuild_rejects_insufficient_space_without_changing_page() -> io::Result<()> { + let original = empty_leaf()?; + let mut page = original; + { + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + let cell = [1; 2029]; + invalid_input(slotted.rebuild_ordered([cell.as_slice(), cell.as_slice()]))?; + } + assert_eq!(page, original); + Ok(()) + } + + #[test] + fn equal_size_replacement_refreshes_page_checksum() -> io::Result<()> { + let mut page = empty_leaf()?; + { + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + slotted.rebuild_ordered([b"old".as_slice()])?; + } + let old_checksum = page_checksum(&page)?; + { + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + slotted.replace_same_len(0, b"new")?; + assert_eq!(slotted.cell(0)?, b"new"); + invalid_input(slotted.replace_same_len(0, b"longer"))?; + } + let new_checksum = page_checksum(&page)?; + assert_ne!(old_checksum, new_checksum); + SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)?; + Ok(()) + } + + #[test] + fn replacement_rejects_a_short_mutable_view_without_partial_mutation() -> io::Result<()> { + let slots = [Slot { offset: 100, length: 3 }]; + let mut bytes = leaf_with_slots(&slots, 100)?; + bytes[100..103].copy_from_slice(b"old"); + write_page_checksum(&mut bytes)?; + let original = bytes; + let mut changing = ChangingMutView { bytes, calls: 0 }; + + { + let mut slotted = SlottedPage::open(&mut changing, PAGE_ID, NEXT_PAGE_ID)?; + invalid_data(slotted.replace_same_len(0, b"new"))?; + } + + assert_eq!(changing.calls, 1); + assert_eq!(changing.bytes, original); + Ok(()) + } + + #[test] + fn rebuild_compacts_gapped_cells() -> io::Result<()> { + let slots = [Slot { offset: 4080, length: 4 }, Slot { offset: 4000, length: 4 }]; + let mut page = leaf_with_slots(&slots, 4000)?; + page[4080..4084].copy_from_slice(b"left"); + page[4000..4004].copy_from_slice(b"rght"); + write_page_checksum(&mut page)?; + + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + slotted.rebuild_ordered([b"left".as_slice(), b"rght".as_slice()])?; + assert_eq!(slotted.header().free_end, 4088); + assert_eq!(slotted.slot(0)?.offset, 4092); + assert_eq!(slotted.slot(1)?.offset, 4088); + Ok(()) + } + + #[test] + fn opens_a_large_densely_packed_page() -> io::Result<()> { + let mut page = empty_leaf()?; + let cell = [7]; + { + let mut slotted = SlottedPage::open(page.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + slotted.rebuild_ordered(std::iter::repeat_n(cell.as_slice(), 800))?; + } + let slotted = SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(slotted.cells().len(), 800); + assert_eq!(slotted.cells().count(), 800); + Ok(()) + } + + #[test] + fn open_rejects_malformed_slot_ranges() -> io::Result<()> { + let cases = [ + (vec![Slot { offset: 4090, length: 0 }], 4090), + (vec![Slot { offset: 35, length: 1 }], 36), + (vec![Slot { offset: 4090, length: 10 }], 4090), + ( + vec![Slot { offset: 4080, length: 8 }, Slot { offset: 4080, length: 8 }], + 4080, + ), + ( + vec![Slot { offset: 4080, length: 8 }, Slot { offset: 4090, length: 6 }], + 4090, + ), + ( + vec![Slot { offset: 4080, length: 16 }, Slot { offset: 4070, length: 16 }], + 4070, + ), + ( + vec![Slot { offset: 4080, length: 16 }, Slot { offset: 4064, length: 16 }], + 4060, + ), + ]; + + for (slots, free_end) in cases { + invalid_data(SlottedPage::open( + leaf_with_slots(&slots, free_end)?.as_slice(), + PAGE_ID, + NEXT_PAGE_ID, + ))?; + } + Ok(()) + } + + #[test] + fn open_rejects_leaf_self_reference() -> io::Result<()> { + let mut page = empty_leaf()?; + page[8..16].copy_from_slice(&PAGE_ID.to_le_bytes()); + write_page_checksum(&mut page)?; + invalid_data(SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID))?; + + let mut page = empty_leaf()?; + set_right_reference(&mut page, PAGE_ID)?; + invalid_data(SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)) + } + + #[test] + fn open_rejects_overflow_self_reference() -> io::Result<()> { + let mut page = chain_page(PageType::Overflow)?; + set_right_reference(&mut page, PAGE_ID)?; + invalid_data(SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)) + } + + #[test] + fn open_rejects_free_self_reference() -> io::Result<()> { + let mut page = chain_page(PageType::Free)?; + set_right_reference(&mut page, PAGE_ID)?; + invalid_data(SlottedPage::open(page.as_slice(), PAGE_ID, NEXT_PAGE_ID)) + } + + #[test] + fn open_validates_internal_preamble() -> io::Result<()> { + let mut page = [0; PAGE_SIZE]; + InternalPreamble { leftmost_child: PAGE_ID }.encode_into(&mut page, 2, NEXT_PAGE_ID)?; + PageHeader { + page_type: PageType::Internal, + cell_count: 1, + free_start: 44, + free_end: 4095, + left: 0, + right: 0, + } + .encode_into(&mut page, 2, NEXT_PAGE_ID)?; + page[40..44].copy_from_slice(&Slot { offset: 4095, length: 1 }.encode()); + write_page_checksum(&mut page)?; + SlottedPage::open(page.as_slice(), 2, NEXT_PAGE_ID)?; + + page[32..40].copy_from_slice(&2u64.to_le_bytes()); + write_page_checksum(&mut page)?; + invalid_data(SlottedPage::open(page.as_slice(), 2, NEXT_PAGE_ID)) + } + + #[test] + fn internal_rebuild_preserves_preamble_cells_and_checksum() -> io::Result<()> { + let mut page = [0; PAGE_SIZE]; + InternalPreamble { leftmost_child: 3 }.encode_into(&mut page, 2, NEXT_PAGE_ID)?; + page[40..44].copy_from_slice(&Slot { offset: 4095, length: 1 }.encode()); + page[4095] = b'x'; + PageHeader { + page_type: PageType::Internal, + cell_count: 1, + free_start: 44, + free_end: 4095, + left: 0, + right: 0, + } + .encode_into(&mut page, 2, NEXT_PAGE_ID)?; + + { + let mut slotted = SlottedPage::open(page.as_mut_slice(), 2, NEXT_PAGE_ID)?; + slotted.rebuild_ordered([b"a".as_slice(), b"bc".as_slice()])?; + assert_eq!(slotted.cells().collect::>>()?, [b"a".as_slice(), b"bc".as_slice()]); + } + + let reopened = SlottedPage::open(page.as_slice(), 2, NEXT_PAGE_ID)?; + assert_eq!(InternalPreamble::decode(&page, 2, NEXT_PAGE_ID)?.leftmost_child, 3); + assert_eq!(reopened.cells().collect::>>()?, [b"a".as_slice(), b"bc".as_slice()]); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree/v3/tree.rs b/backend/src/repository/bplustree/v3/tree.rs new file mode 100644 index 000000000..04f5ee547 --- /dev/null +++ b/backend/src/repository/bplustree/v3/tree.rs @@ -0,0 +1,5720 @@ +use super::{ + format::{ + decompress_value_in_place, decompress_value_into, encode_inline_leaf_cell, encode_internal_cell, + encode_overflow_leaf_cell, encode_tombstone_leaf_cell, encode_value, stored_value_checksum, Compression, + DatabaseHeader, InternalCellRef, InternalPreamble, LeafCellRef, LeafValueRef, Locator, PageHeader, PageType, Slot, + MAX_CELL_FOOTPRINT, MAX_INLINE_STORED_VALUE, OVERFLOW_PAYLOAD_LEN, PAGE_HEADER_LEN, PAGE_SIZE, SLOT_LEN, + }, + page::{encode_free_page, encode_overflow_page, overflow_payload, PageValidation, SlottedPage}, + wal::{ + commit_ordered_page_refs_under_existing_lock, invalidate_sorted_index, recover_pending, + recover_pending_under_existing_lock, recovery_required, sync_parent_directory, wal_path, wal_temporary_path, + with_exclusive_sidecar, ExclusiveSidecarGuard, SharedSidecarGuard, WalOperationError, WalOutcome, + }, + BPlusTreeMetadata, +}; +use crate::{ + repository::bplustree::common::{mmap_with_advice, read_exact_at_offset, BPlusTreeError}, + utils::{binary_deserialize, binary_serialize, binary_serialize_into}, +}; +use memmap2::{Advice, Mmap}; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + fs::{File, OpenOptions}, + io::{self, BufWriter, Read, Write}, + marker::PhantomData, + ops::{Bound, Range}, + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, OnceLock, + }, + thread::JoinHandle, + time::Duration, +}; + +fn invalid_data(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidData, message) } + +fn invalid_input(message: &'static str) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message) } + +fn checked_page_usage<'a, I>(base: usize, cells: I, maximum_cell_footprint: usize) -> io::Result +where + I: IntoIterator, +{ + cells.into_iter().try_fold(base, |used, cell| { + let footprint = SLOT_LEN + .checked_add(cell.len()) + .ok_or_else(|| invalid_input("cell footprint overflow"))?; + if footprint > maximum_cell_footprint { + return Err(invalid_input("cell footprint exceeds format limit")); + } + used.checked_add(footprint).ok_or_else(|| invalid_input("page usage overflow")) + }) +} + +pub(crate) fn used_leaf_bytes>(cells: &[T]) -> io::Result { + checked_page_usage(PAGE_HEADER_LEN, cells.iter().map(AsRef::as_ref), MAX_CELL_FOOTPRINT) +} + +pub(crate) fn choose_leaf_split>(cells: &[T]) -> io::Result { + if cells.len() < 2 { + return Err(invalid_input("leaf split requires two non-empty outputs")); + } + let mut total = 0usize; + for cell in cells { + let footprint = SLOT_LEN + .checked_add(cell.as_ref().len()) + .ok_or_else(|| invalid_input("leaf cell footprint overflow"))?; + if footprint > MAX_CELL_FOOTPRINT { + return Err(invalid_input("leaf cell footprint exceeds format limit")); + } + total = total.checked_add(footprint).ok_or_else(|| invalid_input("leaf split usage overflow"))?; + } + + let mut left_payload = 0usize; + let mut best = None; + for boundary in 1..cells.len() { + let cell = cells + .get(boundary - 1) + .ok_or_else(|| invalid_input("leaf split boundary is outside cells"))?; + let footprint = SLOT_LEN + .checked_add(cell.as_ref().len()) + .ok_or_else(|| invalid_input("leaf cell footprint overflow"))?; + left_payload = left_payload + .checked_add(footprint) + .ok_or_else(|| invalid_input("leaf split usage overflow"))?; + let right_payload = total + .checked_sub(left_payload) + .ok_or_else(|| invalid_input("leaf split usage underflow"))?; + let left_used = PAGE_HEADER_LEN + .checked_add(left_payload) + .ok_or_else(|| invalid_input("leaf split usage overflow"))?; + let right_used = PAGE_HEADER_LEN + .checked_add(right_payload) + .ok_or_else(|| invalid_input("leaf split usage overflow"))?; + if left_used <= PAGE_SIZE && right_used <= PAGE_SIZE { + let imbalance = left_used.abs_diff(right_used); + if best.is_none_or(|(_, best_imbalance)| imbalance < best_imbalance) { + best = Some((boundary, imbalance)); + } + } + } + best.map(|(boundary, _)| boundary) + .ok_or_else(|| invalid_input("leaf split has no valid boundary")) +} + +fn used_internal_bytes<'a, I>(cells: I) -> io::Result +where + I: IntoIterator, +{ + checked_page_usage(PAGE_HEADER_LEN + 8, cells, SLOT_LEN + 12 + 2004) +} + +#[derive(Debug)] +pub(crate) struct InternalSplit<'a, T> { + #[cfg(test)] + pub(crate) promoted_index: usize, + pub(crate) promoted: InternalCellRef<'a>, + pub(crate) right_leftmost_child: u64, + pub(crate) left_cells: &'a [T], + pub(crate) right_cells: &'a [T], +} + +pub(crate) fn choose_internal_split>( + cells: &[T], + page_id: u64, + next_page_id: u64, +) -> io::Result> { + if cells.len() < 3 { + return Err(invalid_input("internal split requires two non-empty outputs and a promoted separator")); + } + for cell in cells { + InternalCellRef::decode(cell.as_ref(), page_id, next_page_id)?; + } + + let mut best = None; + for promoted_index in 1..cells.len() - 1 { + let left_used = used_internal_bytes(cells[..promoted_index].iter().map(AsRef::as_ref))?; + let right_used = used_internal_bytes(cells[promoted_index + 1..].iter().map(AsRef::as_ref))?; + if left_used <= PAGE_SIZE && right_used <= PAGE_SIZE { + let imbalance = left_used.abs_diff(right_used); + if best.is_none_or(|(_, best_imbalance)| imbalance < best_imbalance) { + best = Some((promoted_index, imbalance)); + } + } + } + let promoted_index = best + .map(|(index, _)| index) + .ok_or_else(|| invalid_input("internal split has no valid boundary"))?; + let promoted = InternalCellRef::decode( + cells + .get(promoted_index) + .ok_or_else(|| invalid_input("promoted separator is outside cells"))? + .as_ref(), + page_id, + next_page_id, + )?; + Ok(InternalSplit { + #[cfg(test)] + promoted_index, + right_leftmost_child: promoted.right_child, + promoted, + left_cells: cells + .get(..promoted_index) + .ok_or_else(|| invalid_input("left split range is outside cells"))?, + right_cells: cells + .get(promoted_index + 1..) + .ok_or_else(|| invalid_input("right split range is outside cells"))?, + }) +} + +fn decode_key(encoded: &[u8]) -> io::Result +where + K: for<'de> Deserialize<'de>, +{ + rmp_serde::from_slice(encoded) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("invalid serialized key: {err}"))) +} + +#[cfg(test)] +thread_local! { + static INTERNAL_KEY_DECODE_COUNT: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +#[cfg(test)] +fn reset_internal_key_decode_count() { INTERNAL_KEY_DECODE_COUNT.set(0); } + +#[cfg(test)] +fn internal_key_decode_count() -> usize { INTERNAL_KEY_DECODE_COUNT.get() } + +fn decode_internal_key(encoded: &[u8]) -> io::Result +where + K: for<'de> Deserialize<'de>, +{ + #[cfg(test)] + INTERNAL_KEY_DECODE_COUNT.with(|count| count.set(count.get().saturating_add(1))); + decode_key(encoded) +} + +pub(crate) fn search_leaf(page: &SlottedPage, target: &K) -> io::Result> +where + K: Ord + for<'de> Deserialize<'de>, + B: AsRef<[u8]>, +{ + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("typed leaf search requires a leaf page")); + } + let mut left = 0usize; + let mut right = usize::from(page.header().cell_count); + while left < right { + let middle = left + (right - left) / 2; + let cell = LeafCellRef::decode(page.cell(middle)?, page.page_id(), page.next_page_id())?; + match decode_key::(cell.key_bytes)?.cmp(target) { + std::cmp::Ordering::Less => left = middle + 1, + std::cmp::Ordering::Greater => right = middle, + std::cmp::Ordering::Equal => return Ok(Ok(middle)), + } + } + Ok(Err(left)) +} + +#[cfg(test)] +pub(crate) fn validate_locator>( + page: &SlottedPage, + locator: Locator, + serialized_primary_key: &[u8], +) -> io::Result<()> { + if page.header().page_type != PageType::Leaf || page.page_id() != locator.leaf_page_id { + return Err(invalid_data("locator does not reference this leaf page")); + } + let cell = LeafCellRef::decode( + page.cell(usize::from(locator.slot_index))?, + page.page_id(), + page.next_page_id(), + )?; + let cell_crc = crc32fast::hash(cell.key_bytes); + if cell_crc != locator.serialized_key_crc32 + || crc32fast::hash(serialized_primary_key) != locator.serialized_key_crc32 + || cell.key_bytes != serialized_primary_key + { + return Err(invalid_data("locator serialized key mismatch")); + } + Ok(()) +} + +#[cfg(test)] +fn database_page(database: &[u8], page_id: u64, next_page_id: u64) -> io::Result<&[u8]> { + if page_id == 0 || page_id >= next_page_id { + return Err(invalid_data("overflow page id is outside database")); + } + let page_id = usize::try_from(page_id).map_err(|_| invalid_data("overflow page id exceeds usize"))?; + let offset = page_id + .checked_mul(PAGE_SIZE) + .ok_or_else(|| invalid_data("overflow page offset overflow"))?; + let end = offset + .checked_add(PAGE_SIZE) + .ok_or_else(|| invalid_data("overflow page end overflow"))?; + database.get(offset..end).ok_or_else(|| invalid_data("truncated overflow page")) +} + +#[cfg(test)] +pub(crate) fn read_leaf_value<'a>( + database: &'a [u8], + value: &LeafValueRef<'a>, + next_page_id: u64, + maximum_length: usize, + scratch: &'a mut Vec, +) -> io::Result> { + match *value { + LeafValueRef::Tombstone => Ok(None), + LeafValueRef::Inline { compression, logical_len, stored, .. } => { + let logical_len = usize::try_from(logical_len).map_err(|_| invalid_data("logical length exceeds usize"))?; + if logical_len > maximum_length { + return Err(invalid_data("logical length exceeds allocation limit")); + } + match compression { + Compression::None => Ok(Some(stored)), + Compression::Lz4 => decompress_value_into( + stored, + u32::try_from(logical_len).map_err(|_| invalid_data("logical length exceeds u32"))?, + maximum_length, + scratch, + ) + .map(Some), + } + } + LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } => { + let logical_length = usize::try_from(logical_len).map_err(|_| invalid_data("logical length exceeds usize"))?; + let stored_length = usize::try_from(stored_len).map_err(|_| invalid_data("stored length exceeds usize"))?; + if logical_length > maximum_length || stored_length > maximum_length { + return Err(invalid_data("overflow value exceeds allocation limit")); + } + scratch.clear(); + scratch + .try_reserve(stored_length) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + let mut page_id = head; + let mut visited = HashSet::new(); + while page_id != 0 { + if u64::try_from(visited.len()).map_err(|_| invalid_data("overflow chain length exceeds u64"))? + >= next_page_id + { + return Err(invalid_data("overflow chain contains a cycle")); + } + let page = SlottedPage::open(database_page(database, page_id, next_page_id)?, page_id, next_page_id)?; + if page.header().page_type != PageType::Overflow { + return Err(invalid_data("overflow chain references a non-overflow page")); + } + let payload = overflow_payload(&page)?; + if payload.is_empty() { + return Err(invalid_data("overflow chain contains an empty payload")); + } + visited + .try_reserve(1) + .map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + if !visited.insert(page_id) { + return Err(invalid_data("overflow chain contains a cycle")); + } + let new_length = scratch + .len() + .checked_add(payload.len()) + .ok_or_else(|| invalid_data("overflow chain length overflow"))?; + if new_length > stored_length { + return Err(invalid_data("overflow chain exceeds stored length")); + } + scratch.extend_from_slice(payload); + page_id = page.header().right; + } + if scratch.len() != stored_length { + return Err(invalid_data("overflow chain stored length mismatch")); + } + if crc32fast::hash(scratch) != crc32 { + return Err(invalid_data("stored value checksum mismatch")); + } + if compression == Compression::Lz4 { + decompress_value_in_place(scratch, logical_len, maximum_length)?; + } else if scratch.len() != logical_length { + return Err(invalid_data("uncompressed overflow length mismatch")); + } + Ok(Some(scratch.as_slice())) + } + } +} + +#[derive(Clone, Debug)] +pub struct BPlusTree { + entries: BTreeMap, + metadata: BPlusTreeMetadata, + dirty: bool, +} + +impl Default for BPlusTree { + fn default() -> Self { Self::new() } +} + +impl BPlusTree { + pub const fn new() -> Self { + Self { entries: BTreeMap::new(), metadata: BPlusTreeMetadata::Empty, dirty: true } + } + + pub fn get_metadata(&self) -> &BPlusTreeMetadata { &self.metadata } + + pub fn set_metadata(&mut self, metadata: BPlusTreeMetadata) { + self.metadata = metadata; + self.dirty = true; + } + + pub fn is_empty(&self) -> bool { self.entries.is_empty() } + + pub fn len(&self) -> usize { self.entries.len() } + + pub fn insert(&mut self, key: K, value: V) { + let _ = self.entries.insert(key, value); + self.dirty = true; + } + + pub fn query(&self, key: &K) -> Option<&V> { self.entries.get(key) } + + pub fn find_le(&self, key: &K) -> Option<(&K, &V)> { self.entries.range(..=key).next_back() } + + pub fn iter(&self) -> std::collections::btree_map::Iter<'_, K, V> { self.entries.iter() } +} + +impl<'a, K: Ord, V> IntoIterator for &'a BPlusTree { + type Item = (&'a K, &'a V); + type IntoIter = std::collections::btree_map::Iter<'a, K, V>; + + fn into_iter(self) -> Self::IntoIter { self.entries.iter() } +} + +#[derive(Clone)] +struct NodeInfo { + page_id: u64, + minimum_key: Vec, +} + +fn allocate_page(pages: &mut Vec<[u8; PAGE_SIZE]>) -> io::Result { + let page_id = u64::try_from(pages.len()).map_err(|_| invalid_input("page count exceeds u64"))?; + pages.try_reserve(1).map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + pages.push([0; PAGE_SIZE]); + Ok(page_id) +} + +fn leaf_page>(page_id: u64, left: u64, right: u64, cells: &[T]) -> io::Result<[u8; PAGE_SIZE]> { + let mut bytes = [0; PAGE_SIZE]; + PageHeader { + page_type: PageType::Leaf, + cell_count: 0, + free_start: u16::try_from(PAGE_HEADER_LEN).map_err(|_| invalid_input("leaf header exceeds u16"))?, + free_end: u16::try_from(PAGE_SIZE).map_err(|_| invalid_input("page size exceeds u16"))?, + left, + right, + } + .encode_into(&mut bytes, page_id, u64::MAX)?; + SlottedPage::open(bytes.as_mut_slice(), page_id, u64::MAX)? + .rebuild_ordered(cells.iter().map(AsRef::as_ref))?; + Ok(bytes) +} + +fn internal_page>( + page_id: u64, + leftmost_child: u64, + cells: &[T], +) -> io::Result<[u8; PAGE_SIZE]> { + let first = cells + .first() + .ok_or_else(|| invalid_input("internal page must contain a separator"))? + .as_ref(); + let first_offset = PAGE_SIZE + .checked_sub(first.len()) + .ok_or_else(|| invalid_input("internal cell exceeds page"))?; + let mut bytes = [0; PAGE_SIZE]; + InternalPreamble { leftmost_child }.encode_into(&mut bytes, page_id, u64::MAX)?; + bytes + .get_mut(40..44) + .ok_or_else(|| invalid_input("internal slot is outside page"))? + .copy_from_slice( + &Slot { + offset: u16::try_from(first_offset).map_err(|_| invalid_input("internal cell offset exceeds u16"))?, + length: u16::try_from(first.len()).map_err(|_| invalid_input("internal cell length exceeds u16"))?, + } + .encode(), + ); + bytes + .get_mut(first_offset..) + .ok_or_else(|| invalid_input("internal cell is outside page"))? + .copy_from_slice(first); + PageHeader { + page_type: PageType::Internal, + cell_count: 1, + free_start: 44, + free_end: u16::try_from(first_offset).map_err(|_| invalid_input("internal free_end exceeds u16"))?, + left: 0, + right: 0, + } + .encode_into(&mut bytes, page_id, u64::MAX)?; + SlottedPage::open(bytes.as_mut_slice(), page_id, u64::MAX)? + .rebuild_ordered(cells.iter().map(AsRef::as_ref))?; + Ok(bytes) +} + +fn allocate_overflow_chain(pages: &mut Vec<[u8; PAGE_SIZE]>, stored: &[u8]) -> io::Result { + let count = stored.len().div_ceil(OVERFLOW_PAYLOAD_LEN); + if count == 0 { + return Err(invalid_input("overflow value must not be empty")); + } + let head = u64::try_from(pages.len()).map_err(|_| invalid_input("page count exceeds u64"))?; + pages.try_reserve(count).map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + for _ in 0..count { + pages.push([0; PAGE_SIZE]); + } + for (index, payload) in stored.chunks(OVERFLOW_PAYLOAD_LEN).enumerate() { + let page_id = head + .checked_add(u64::try_from(index).map_err(|_| invalid_input("overflow index exceeds u64"))?) + .ok_or_else(|| invalid_input("overflow page id overflow"))?; + let next = if index + 1 == count { + 0 + } else { + page_id.checked_add(1).ok_or_else(|| invalid_input("overflow page id overflow"))? + }; + let page = encode_overflow_page(page_id, u64::MAX, next, payload)?; + let slot = usize::try_from(page_id).map_err(|_| invalid_input("overflow page id exceeds usize"))?; + *pages.get_mut(slot).ok_or_else(|| invalid_input("overflow page allocation missing"))? = page; + } + Ok(head) +} + +fn finish_leaf( + pages: &mut [[u8; PAGE_SIZE]], + page_id: u64, + left: u64, + right: u64, + cells: &[Vec], +) -> io::Result<()> { + let page = leaf_page(page_id, left, right, cells)?; + let index = usize::try_from(page_id).map_err(|_| invalid_input("leaf page id exceeds usize"))?; + *pages.get_mut(index).ok_or_else(|| invalid_input("leaf page allocation missing"))? = page; + Ok(()) +} + +fn build_leaf_level(entries: &BTreeMap, pages: &mut Vec<[u8; PAGE_SIZE]>) -> io::Result> +where + K: Ord + Serialize, + V: Serialize, +{ + let mut leaf_id = allocate_page(pages)?; + let mut left_leaf = 0; + let mut cells = Vec::>::new(); + let mut leaves = Vec::::new(); + let mut leaf_minimum = None; + let mut compression_scratch = Vec::new(); + let mut cell = Vec::new(); + + for (key, value) in entries { + let key_bytes = binary_serialize(key)?; + let raw_value = binary_serialize(value)?; + let logical_len = u32::try_from(raw_value.len()).map_err(|_| invalid_input("serialized value exceeds u32"))?; + let stored = encode_value(&raw_value, &mut compression_scratch)?; + let stored_bytes = stored.as_slice(); + let inline_footprint = SLOT_LEN + .checked_add(24) + .and_then(|size| size.checked_add(key_bytes.len())) + .and_then(|size| size.checked_add(stored_bytes.len())) + .ok_or_else(|| invalid_input("leaf cell footprint overflow"))?; + if stored_bytes.len() <= MAX_INLINE_STORED_VALUE && inline_footprint <= MAX_CELL_FOOTPRINT { + encode_inline_leaf_cell(&key_bytes, logical_len, stored.compression(), stored_bytes, &mut cell)?; + } else { + let head = allocate_overflow_chain(pages, stored_bytes)?; + encode_overflow_leaf_cell( + &key_bytes, + logical_len, + stored.compression(), + u32::try_from(stored_bytes.len()).map_err(|_| invalid_input("stored value exceeds u32"))?, + head, + stored_value_checksum(stored_bytes), + leaf_id, + u64::MAX, + &mut cell, + )?; + } + + let next_usage = checked_page_usage( + PAGE_HEADER_LEN, + cells.iter().map(Vec::as_slice).chain(std::iter::once(cell.as_slice())), + MAX_CELL_FOOTPRINT, + )?; + if next_usage > PAGE_SIZE { + let next_leaf = allocate_page(pages)?; + finish_leaf(pages, leaf_id, left_leaf, next_leaf, &cells)?; + leaves.push(NodeInfo { + page_id: leaf_id, + minimum_key: leaf_minimum.take().ok_or_else(|| invalid_input("non-empty leaf has no minimum key"))?, + }); + left_leaf = leaf_id; + leaf_id = next_leaf; + cells.clear(); + } + if cells.is_empty() { + leaf_minimum = Some(key_bytes); + } + cells.try_reserve(1).map_err(|err| io::Error::new(io::ErrorKind::OutOfMemory, err))?; + cells.push(std::mem::take(&mut cell)); + } + + finish_leaf(pages, leaf_id, left_leaf, 0, &cells)?; + leaves.push(NodeInfo { page_id: leaf_id, minimum_key: leaf_minimum.unwrap_or_default() }); + Ok(leaves) +} + +fn group_internal_children(level: &[NodeInfo]) -> io::Result> { + let mut groups = Vec::<(usize, usize)>::new(); + let mut start = 0usize; + while start < level.len() { + let mut end = start + 1; + let mut used = PAGE_HEADER_LEN + 8; + while end < level.len() { + let child = level.get(end).ok_or_else(|| invalid_input("internal child index is invalid"))?; + let cell_len = 12usize + .checked_add(child.minimum_key.len()) + .ok_or_else(|| invalid_input("internal cell length overflow"))?; + let next = used + .checked_add(SLOT_LEN) + .and_then(|size| size.checked_add(cell_len)) + .ok_or_else(|| invalid_input("internal page usage overflow"))?; + if next > PAGE_SIZE { + break; + } + used = next; + end += 1; + } + groups.push((start, end)); + start = end; + } + if groups.len() > 1 && groups.last().is_some_and(|(from, to)| to - from == 1) { + let last = groups.len() - 1; + let previous = groups + .get(last - 1) + .copied() + .ok_or_else(|| invalid_input("internal grouping is missing its previous group"))?; + let moved = previous.1.checked_sub(1).ok_or_else(|| invalid_input("internal grouping underflow"))?; + if moved <= previous.0 { + return Err(invalid_input("internal page cannot retain two children")); + } + groups + .get_mut(last - 1) + .ok_or_else(|| invalid_input("internal grouping is missing its previous group"))? + .1 = moved; + groups + .get_mut(last) + .ok_or_else(|| invalid_input("internal grouping is missing its last group"))? + .0 = moved; + } + Ok(groups) +} + +fn build_parent_level(level: &[NodeInfo], pages: &mut Vec<[u8; PAGE_SIZE]>) -> io::Result> { + let groups = group_internal_children(level)?; + let mut parents = Vec::with_capacity(groups.len()); + for (start, end) in groups { + let children = level.get(start..end).ok_or_else(|| invalid_input("internal child range is invalid"))?; + if children.len() < 2 { + return Err(invalid_input("internal page requires two children")); + } + let page_id = allocate_page(pages)?; + let (first_child, remaining_children) = + children.split_first().ok_or_else(|| invalid_input("internal page has no children"))?; + let mut encoded_cells = Vec::with_capacity(remaining_children.len()); + for child in remaining_children { + let mut encoded = Vec::new(); + encode_internal_cell(&child.minimum_key, child.page_id, page_id, u64::MAX, &mut encoded)?; + encoded_cells.push(encoded); + } + let encoded = internal_page(page_id, first_child.page_id, &encoded_cells)?; + let page_index = usize::try_from(page_id).map_err(|_| invalid_input("internal page id exceeds usize"))?; + *pages.get_mut(page_index).ok_or_else(|| invalid_input("internal page allocation missing"))? = encoded; + parents.push(NodeInfo { page_id, minimum_key: first_child.minimum_key.clone() }); + } + Ok(parents) +} + +fn build_pages(entries: &BTreeMap) -> io::Result<(Vec<[u8; PAGE_SIZE]>, u64)> +where + K: Ord + Serialize, + V: Serialize, +{ + let mut pages = vec![[0; PAGE_SIZE]]; + let mut level = build_leaf_level(entries, &mut pages)?; + if entries.is_empty() { + let root = level.first().ok_or_else(|| invalid_input("tree has no root page"))?.page_id; + return Ok((pages, root)); + } + while level.len() > 1 { + level = build_parent_level(&level, &mut pages)?; + } + let root_page_id = level.first().ok_or_else(|| invalid_input("tree has no root page"))?.page_id; + Ok((pages, root_page_id)) +} + +fn temporary_path(filepath: &Path) -> io::Result { + let name = filepath + .file_name() + .ok_or_else(|| invalid_input("database path has no file name"))? + .to_string_lossy(); + Ok(filepath.with_file_name(format!("{name}.{}.v3.tmp", uuid::Uuid::new_v4()))) +} + +fn publish_database( + temporary: &Path, + destination: &Path, + sync_directory: impl FnOnce(&Path) -> io::Result<()>, +) -> io::Result<()> { + let temporary = match tempfile::TempPath::try_from_path(temporary) { + Ok(path) => path, + Err(error) => { + let _ = std::fs::remove_file(temporary); + return Err(error); + } + }; + temporary.persist(destination).map_err(io::Error::from)?; + sync_directory(destination).map_err(|error| { + io::Error::new( + error.kind(), + format!("database published but directory sync failed; durability unknown: {error}"), + ) + }) +} + +impl BPlusTree +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de>, +{ + pub fn store(&mut self, filepath: &Path) -> io::Result { + with_exclusive_sidecar(filepath, || { + recover_pending_under_existing_lock(filepath)?; + if self.dirty { + self.store_exclusive(filepath).map(|stored| stored.root_page_id) + } else { + Ok(0) + } + }) + } + + pub(crate) fn store_verified(&mut self, filepath: &Path) -> io::Result { + with_exclusive_sidecar(filepath, || { + recover_pending_under_existing_lock(filepath)?; + if !self.dirty { + return Err(invalid_input("verified store requires a dirty tree")); + } + self.store_exclusive(filepath).map(|stored| stored.verification) + }) + } + + fn store_exclusive(&mut self, filepath: &Path) -> io::Result { + let (mut pages, root_page_id) = build_pages(&self.entries)?; + let next_page_id = u64::try_from(pages.len()).map_err(|_| invalid_input("page count exceeds u64"))?; + *pages.first_mut().ok_or_else(|| invalid_input("database header page is missing"))? = DatabaseHeader { + root_page_id, + next_page_id, + free_page_head: 0, + generation: 1, + database_id: *uuid::Uuid::new_v4().as_bytes(), + metadata: self.metadata.clone(), + } + .encode()?; + + let temp_path = temporary_path(filepath)?; + let prepared = (|| { + let mut file = OpenOptions::new().write(true).create_new(true).open(&temp_path)?; + { + let mut writer = BufWriter::with_capacity(1024 * 1024, &mut file); + for page in &pages { + writer.write_all(page)?; + } + writer.flush()?; + } + file.sync_all()?; + drop(file); + let mut query = BPlusTreeQuery::::from_file_unlocked(File::open(&temp_path)?)?; + let verification = verify_full(&mut query)?; + drop(query); + Ok(verification) + })(); + let verification = match prepared { + Ok(verification) => verification, + Err(error) => { + let _ = std::fs::remove_file(&temp_path); + return Err(error); + } + }; + publish_database(&temp_path, filepath, sync_parent_directory)?; + invalidate_sorted_index(filepath)?; + self.dirty = false; + Ok(StoredDatabase { root_page_id, verification }) + } + + pub fn store_with_index(&mut self, filepath: &Path, sort_key_extractor: F) -> io::Result + where + SortKey: Ord + Serialize, + F: Fn(&V) -> SortKey, + { + self.store_with_index_result(filepath, sort_key_extractor) + .map(|stored| stored.map_or(0, |stored| stored.root_page_id)) + } + + pub(crate) fn store_with_index_verified( + &mut self, + filepath: &Path, + sort_key_extractor: F, + ) -> io::Result + where + SortKey: Ord + Serialize, + F: Fn(&V) -> SortKey, + { + self.store_with_index_result(filepath, sort_key_extractor)? + .map(|stored| stored.verification) + .ok_or_else(|| invalid_input("verified indexed store requires a dirty tree")) + } + + fn store_with_index_result( + &mut self, + filepath: &Path, + sort_key_extractor: F, + ) -> io::Result> + where + SortKey: Ord + Serialize, + F: Fn(&V) -> SortKey, + { + with_exclusive_sidecar(filepath, || { + recover_pending_under_existing_lock(filepath)?; + if !self.dirty { + return Ok(None); + } + let stored = self.store_exclusive(filepath)?; + Self::store_index_exclusive(filepath, sort_key_extractor, stored.verification.live_entries)?; + Ok(Some(stored)) + }) + } + + fn store_index_exclusive( + filepath: &Path, + sort_key_extractor: F, + expected_entries: u64, + ) -> io::Result<()> + where + SortKey: Ord + Serialize, + F: Fn(&V) -> SortKey, + { + let mut query = BPlusTreeQuery::::from_file_unlocked(File::open(filepath)?)?; + let (database_id, generation) = query.snapshot_identity(); + let mut entries = query + .collect_with_locators()? + .into_iter() + .map(|(key, value, locator)| (sort_key_extractor(&value), key, locator)) + .collect::>(); + if u64::try_from(entries.len()).map_err(|_| invalid_data("sorted-index entry count exceeds u64"))? + != expected_entries + { + return Err(invalid_data("sorted-index source entry count differs from verified database")); + } + entries.sort_unstable_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))); + drop(query); + + let index_path = crate::repository::storage::get_file_path_for_db_index(filepath); + let temporary = temporary_path(&index_path)?; + let prepared = (|| { + let mut writer = crate::repository::bplustree::sorted_index::v4::Writer::new(&temporary, database_id, generation)?; + for (sort_key, primary_key, locator) in &entries { + writer.push(sort_key, primary_key, *locator)?; + } + let _ = writer.finish()?; + Ok(()) + })(); + if let Err(error) = prepared { + let _ = std::fs::remove_file(&temporary); + return Err(error); + } + publish_database(&temporary, &index_path, sync_parent_directory) + } + +} + +struct StoredDatabase { + root_page_id: u64, + verification: VerificationReport, +} + +impl BPlusTree +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + pub fn load(filepath: &Path) -> io::Result { + let mut query = BPlusTreeQuery::::try_new(filepath)?; + let metadata = query.header.metadata.clone(); + let mut entries = BTreeMap::new(); + for entry in query.iter() { + let (key, value) = entry?; + let _ = entries.insert(key, value); + } + Ok(Self { entries, metadata, dirty: false }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum FlushPolicy { + Immediate, + Batch, +} + +struct DatabaseImage { + mmap: Option, + fallback: Vec, +} + +impl DatabaseImage { + fn open(path: &Path) -> io::Result<(Self, DatabaseHeader)> { + let mut file = File::open(path)?; + let file_len = usize::try_from(file.metadata()?.len()) + .map_err(|_| invalid_data("database length exceeds usize"))?; + let mmap = mmap_with_advice(&file, Advice::Normal, "v3 B+Tree update"); + let mut fallback = Vec::new(); + if mmap.is_none() { + fallback + .try_reserve_exact(file_len) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + file.read_to_end(&mut fallback)?; + } + let image = Self { mmap, fallback }; + let bytes = image.as_slice(); + let header = DatabaseHeader::decode( + bytes.get(..PAGE_SIZE).ok_or_else(|| invalid_data("database header is truncated"))?, + )?; + let expected = usize::try_from(header.next_page_id) + .map_err(|_| invalid_data("next page id exceeds usize"))? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| invalid_data("database length overflow"))?; + if bytes.len() != expected { + return Err(invalid_data("database file length does not match header")); + } + Ok((image, header)) + } + + fn as_slice(&self) -> &[u8] { self.mmap.as_deref().unwrap_or(&self.fallback) } +} + +struct WriteTransaction { + original_header: DatabaseHeader, + next_header: DatabaseHeader, + original_file_len: u64, + dirty_pages: BTreeMap>, + allocated_pages: HashSet, + freed_pages: HashSet, +} + +impl WriteTransaction { + fn new(header: DatabaseHeader, original_file_len: usize) -> io::Result { + Ok(Self { + original_header: header.clone(), + next_header: header, + original_file_len: u64::try_from(original_file_len) + .map_err(|_| invalid_data("database length exceeds u64"))?, + dirty_pages: BTreeMap::new(), + allocated_pages: HashSet::new(), + freed_pages: HashSet::new(), + }) + } + + fn page<'a>(&'a self, base: &'a [u8], page_id: u64) -> io::Result<&'a [u8]> { + if page_id == 0 || page_id >= self.next_header.next_page_id { + return Err(invalid_data("transaction page id is outside database")); + } + if let Some(page) = self.dirty_pages.get(&page_id) { + return Ok(page.as_slice()); + } + if page_id >= self.original_header.next_page_id { + return Err(invalid_data("transaction appended page is missing")); + } + let range = page_byte_range(page_id, base.len())?; + base.get(range).ok_or_else(|| invalid_data("transaction base page is truncated")) + } + + fn page_copy(&self, base: &[u8], page_id: u64) -> io::Result<[u8; PAGE_SIZE]> { + let mut copied = [0; PAGE_SIZE]; + copied.copy_from_slice(self.page(base, page_id)?); + Ok(copied) + } + + fn page_mut<'a>(&'a mut self, base: &[u8], page_id: u64) -> io::Result<&'a mut [u8; PAGE_SIZE]> { + if page_id == 0 || page_id >= self.next_header.next_page_id { + return Err(invalid_data("transaction page id is outside database")); + } + if self.dirty_pages.contains_key(&page_id) { + return self + .dirty_pages + .get_mut(&page_id) + .map(Box::as_mut) + .ok_or_else(|| invalid_data("transaction page copy is missing")); + } + if page_id >= self.original_header.next_page_id { + return Err(invalid_data("appended transaction page is missing")); + } + let range = page_byte_range(page_id, base.len())?; + let source = base.get(range).ok_or_else(|| invalid_data("transaction base page is truncated"))?; + let mut copied = Box::new([0; PAGE_SIZE]); + copied.copy_from_slice(source); + let _ = self.dirty_pages.insert(page_id, copied); + self.dirty_pages + .get_mut(&page_id) + .map(Box::as_mut) + .ok_or_else(|| invalid_data("transaction page copy is missing")) + } + + #[allow(clippy::large_types_passed_by_value)] + fn write_page(&mut self, page_id: u64, page: [u8; PAGE_SIZE]) -> io::Result<()> { + if page_id == 0 || page_id >= self.next_header.next_page_id { + return Err(invalid_input("written page id is outside transaction bounds")); + } + let _ = self.dirty_pages.insert(page_id, Box::new(page)); + Ok(()) + } + + fn allocate_page(&mut self, base: &[u8]) -> io::Result { + if self.next_header.free_page_head != 0 { + let page_id = self.next_header.free_page_head; + let page = SlottedPage::open(self.page(base, page_id)?, page_id, self.next_header.next_page_id)?; + if page.header().page_type != PageType::Free { + return Err(invalid_data("free list references a non-free page")); + } + self.next_header.free_page_head = page.header().right; + if !self.allocated_pages.insert(page_id) { + return Err(invalid_data("free page was allocated twice")); + } + let _ = self.freed_pages.remove(&page_id); + return Ok(page_id); + } + let page_id = self.next_header.next_page_id; + self.next_header.next_page_id = page_id + .checked_add(1) + .ok_or_else(|| invalid_input("next page id overflow"))?; + if !self.allocated_pages.insert(page_id) { + return Err(invalid_data("appended page was allocated twice")); + } + Ok(page_id) + } + + fn free_page(&mut self, page_id: u64) -> io::Result<()> { + if page_id == 0 || page_id == self.next_header.root_page_id || page_id >= self.next_header.next_page_id { + return Err(invalid_input("page cannot be added to the free list")); + } + if !self.freed_pages.insert(page_id) { + return Err(invalid_data("page was freed twice in one transaction")); + } + let _ = self.allocated_pages.remove(&page_id); + let page = encode_free_page(page_id, self.next_header.next_page_id, self.next_header.free_page_head)?; + self.next_header.free_page_head = page_id; + self.write_page(page_id, page) + } + + fn has_changes(&self) -> bool { + !self.dirty_pages.is_empty() || self.next_header.metadata != self.original_header.metadata + } + + fn prepared_pages(&mut self) -> io::Result> { + if !self.has_changes() { + return Ok(Vec::new()); + } + self.next_header.generation = self + .original_header + .generation + .checked_add(1) + .ok_or_else(|| invalid_input("database generation overflow"))?; + let _ = self.dirty_pages.insert(0, Box::new(self.next_header.encode()?)); + let mut prepared = Vec::new(); + prepared + .try_reserve_exact(self.dirty_pages.len()) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for (page_id, page) in &self.dirty_pages { + if *page_id == 0 { + DatabaseHeader::decode(page.as_slice())?; + } else { + SlottedPage::open(page.as_slice(), *page_id, self.next_header.next_page_id)?; + } + prepared.push((*page_id, page.as_ref())); + } + Ok(prepared) + } +} + +fn page_byte_range(page_id: u64, database_len: usize) -> io::Result> { + let start = usize::try_from(page_id) + .map_err(|_| invalid_data("page id exceeds usize"))? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| invalid_data("page offset overflow"))?; + let end = start.checked_add(PAGE_SIZE).ok_or_else(|| invalid_data("page end overflow"))?; + if end > database_len { + return Err(invalid_data("page is outside database")); + } + Ok(start..end) +} + +struct ActiveBatch { + base: DatabaseImage, + transaction: WriteTransaction, + _guard: ExclusiveSidecarGuard, +} + +#[derive(Default)] +struct WriteScratch { + key: Vec, + value: Vec, + compression: Vec, + cell: Vec, + read_value: Vec, +} + +pub struct BPlusTreeUpdate { + filepath: PathBuf, + database_id: [u8; 16], + verified_generation: u64, + verified_next_page_id: u64, + flush_policy: FlushPolicy, + active: Option, + scratch: WriteScratch, + _types: PhantomData<(K, V)>, +} + +impl BPlusTreeUpdate +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de>, +{ + pub fn try_new(filepath: &Path) -> io::Result { + let mut query = BPlusTreeQuery::::try_new(filepath)?; + let _ = verify_full(&mut query)?; + let database_id = query.header.database_id; + let verified_generation = query.header.generation; + let verified_next_page_id = query.header.next_page_id; + drop(query); + Ok(Self { + filepath: filepath.to_path_buf(), + database_id, + verified_generation, + verified_next_page_id, + flush_policy: FlushPolicy::Immediate, + active: None, + scratch: WriteScratch::default(), + _types: PhantomData, + }) + } + + pub fn try_new_with_backoff(filepath: &Path) -> io::Result { Self::try_new(filepath) } + + pub fn try_new_with_backoff_stats(filepath: &Path) -> io::Result<(Self, u64)> { + Self::try_new(filepath).map(|updater| (updater, 0)) + } + + pub fn set_flush_policy(&mut self, policy: FlushPolicy) { self.flush_policy = policy; } + + fn ensure_transaction(&mut self) -> io::Result<()> { + if self.active.is_some() { + return Ok(()); + } + let guard = ExclusiveSidecarGuard::acquire(&self.filepath)?; + recover_pending_under_existing_lock(&self.filepath)?; + let (base, header) = DatabaseImage::open(&self.filepath)?; + if header.database_id != self.database_id + || header.generation != self.verified_generation + || header.next_page_id != self.verified_next_page_id + { + let mut query = BPlusTreeQuery::::from_file_unlocked(File::open(&self.filepath)?)?; + let _ = verify_full(&mut query)?; + self.database_id = query.header.database_id; + self.verified_generation = query.header.generation; + self.verified_next_page_id = query.header.next_page_id; + } + let transaction = WriteTransaction::new(header, base.as_slice().len())?; + self.active = Some(ActiveBatch { base, transaction, _guard: guard }); + Ok(()) + } + + pub fn update(&mut self, key: &K, value: V) -> Result { + self.upsert(key, &value).map_err(Into::into) + } + + pub fn update_batch(&mut self, items: &[(&K, &V)]) -> Result { + self.upsert_batch(items).map_err(Into::into) + } + + pub fn prepare_upsert_batch(items: &[(&K, &V)]) -> io::Result)>> { + let mut prepared = Vec::new(); + prepared + .try_reserve_exact(items.len()) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for (key, value) in items { + prepared.push(((*key).clone(), binary_serialize(value)?)); + } + prepared.sort_by(|left, right| left.0.cmp(&right.0)); + 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) + } + + pub fn upsert_batch(&mut self, items: &[(&K, &V)]) -> io::Result { + let prepared = match Self::prepare_upsert_batch(items) { + Ok(prepared) => prepared, + Err(error) => { + self.active = None; + return Err(error); + } + }; + self.upsert_batch_encoded(prepared) + } + + pub fn upsert_batch_encoded(&mut self, items: Vec<(K, Vec)>) -> io::Result { + if items.is_empty() { + return self.upsert_batch(&[]); + } + let policy = self.flush_policy; + self.flush_policy = FlushPolicy::Batch; + let mut root = 0; + for (key, value) in items { + match self.upsert_serialized(&key, &value) { + Ok(next_root) => root = next_root, + Err(error) => { + self.active = None; + self.flush_policy = policy; + return Err(error); + } + } + } + self.flush_policy = policy; + if policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(root) + } + + pub fn upsert(&mut self, key: &K, value: &V) -> io::Result { + let result = self.upsert_inner(key, value); + if result.is_err() { + self.active = None; + } + result + } + + fn upsert_inner(&mut self, key: &K, value: &V) -> io::Result { + self.scratch.key.clear(); + binary_serialize_into(&mut self.scratch.key, key)?; + self.scratch.value.clear(); + binary_serialize_into(&mut self.scratch.value, value)?; + self.ensure_transaction()?; + let stored = match encode_value(&self.scratch.value, &mut self.scratch.compression) { + Ok(stored) => stored, + Err(error) => { + self.active = None; + return Err(error); + } + }; + let active = self.active.as_mut().ok_or_else(|| invalid_data("write transaction is missing"))?; + let staged = stage_upsert::( + &mut active.transaction, + active.base.as_slice(), + key, + &self.scratch.key, + u32::try_from(self.scratch.value.len()).map_err(|_| invalid_input("serialized value exceeds u32"))?, + stored.compression(), + stored.as_slice(), + &mut self.scratch.cell, + ); + if let Err(error) = staged { + self.active = None; + return Err(error); + } + let root = self + .active + .as_ref() + .ok_or_else(|| invalid_data("write transaction is missing"))? + .transaction + .next_header + .root_page_id; + if self.flush_policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(root) + } + + fn upsert_serialized(&mut self, key: &K, raw_value: &[u8]) -> io::Result { + let result = self.upsert_serialized_inner(key, raw_value); + if result.is_err() { + self.active = None; + } + result + } + + fn upsert_serialized_inner(&mut self, key: &K, raw_value: &[u8]) -> io::Result { + self.scratch.key.clear(); + binary_serialize_into(&mut self.scratch.key, key)?; + self.ensure_transaction()?; + let stored = match encode_value(raw_value, &mut self.scratch.compression) { + Ok(stored) => stored, + Err(error) => { + self.active = None; + return Err(error); + } + }; + let active = self.active.as_mut().ok_or_else(|| invalid_data("write transaction is missing"))?; + let staged = stage_upsert::( + &mut active.transaction, + active.base.as_slice(), + key, + &self.scratch.key, + u32::try_from(raw_value.len()).map_err(|_| invalid_input("serialized value exceeds u32"))?, + stored.compression(), + stored.as_slice(), + &mut self.scratch.cell, + ); + if let Err(error) = staged { + self.active = None; + return Err(error); + } + let root = self + .active + .as_ref() + .ok_or_else(|| invalid_data("write transaction is missing"))? + .transaction + .next_header + .root_page_id; + if self.flush_policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(root) + } + + pub fn delete(&mut self, key: &K) -> io::Result { + let result = self.delete_inner(key); + if result.is_err() { + self.active = None; + } + result + } + + fn delete_inner(&mut self, key: &K) -> io::Result { + let started = self.active.is_none(); + self.scratch.key.clear(); + binary_serialize_into(&mut self.scratch.key, key)?; + self.ensure_transaction()?; + let active = self.active.as_mut().ok_or_else(|| invalid_data("write transaction is missing"))?; + let deleted = stage_delete::( + &mut active.transaction, + active.base.as_slice(), + key, + &self.scratch.key, + &mut self.scratch.cell, + ); + let deleted = match deleted { + Ok(deleted) => deleted, + Err(error) => { + self.active = None; + return Err(error); + } + }; + if !deleted && started { + self.active = None; + } else if deleted && self.flush_policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(deleted) + } + + pub fn delete_batch(&mut self, keys: &[&K]) -> io::Result { + if keys.is_empty() { + return Ok(0); + } + let policy = self.flush_policy; + self.flush_policy = FlushPolicy::Batch; + let mut deleted = 0usize; + let mut ordered = Vec::new(); + ordered + .try_reserve_exact(keys.len()) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + ordered.extend_from_slice(keys); + ordered.sort(); + for key in ordered { + match self.delete(key) { + Ok(true) => deleted = deleted.checked_add(1).ok_or_else(|| invalid_input("delete count overflow"))?, + Ok(false) => {} + Err(error) => { + self.active = None; + self.flush_policy = policy; + return Err(error); + } + } + } + self.flush_policy = policy; + if policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(deleted) + } + + pub fn get_metadata(&self) -> io::Result { + if let Some(active) = &self.active { + return Ok(active.transaction.next_header.metadata.clone()); + } + BPlusTreeQuery::::try_new(&self.filepath).map(|query| query.header.metadata) + } + + pub fn set_metadata(&mut self, metadata: &BPlusTreeMetadata) -> io::Result<()> { + let started = self.active.is_none(); + self.ensure_transaction()?; + let active = self.active.as_mut().ok_or_else(|| invalid_data("write transaction is missing"))?; + if active.transaction.next_header.metadata == *metadata { + if started { + self.active = None; + } + return Ok(()); + } + active.transaction.next_header.metadata = metadata.clone(); + if self.flush_policy == FlushPolicy::Immediate { + self.commit()?; + } + Ok(()) + } + + pub fn query(&mut self, key: &K) -> Result, BPlusTreeError> { + if let Some(active) = &mut self.active { + return query_transaction::( + &active.transaction, + active.base.as_slice(), + key, + &mut self.scratch.read_value, + ) + .map_err(Into::into); + } + BPlusTreeQuery::::try_new(&self.filepath) + .and_then(|mut query| query.query_io(key)) + .map_err(Into::into) + } + + pub fn commit(&mut self) -> io::Result<()> { + let Some(mut active) = self.active.take() else { return Ok(()) }; + validate_transaction_links::(&active.transaction, active.base.as_slice())?; + let prepared = active.transaction.prepared_pages()?; + if prepared.is_empty() { + return Ok(()); + } + let committed_header = DatabaseHeader::decode(prepared[0].1)?; + let committed_generation = committed_header.generation; + let committed_next_page_id = committed_header.next_page_id; + let result = commit_ordered_page_refs_under_existing_lock(&self.filepath, &prepared); + match result { + Ok(()) => { + self.verified_generation = committed_generation; + self.verified_next_page_id = committed_next_page_id; + Ok(()) + } + Err(error) => { + let outcome = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .map(WalOperationError::outcome); + if outcome == Some(WalOutcome::CommittedCleanupPending) { + self.verified_generation = committed_generation; + self.verified_next_page_id = committed_next_page_id; + } + if let Err(recovery_error) = recover_pending_under_existing_lock(&self.filepath) { + log::error!( + "B+Tree commit failed and recovery remains pending for {}: {recovery_error}", + self.filepath.display() + ); + } + Err(error) + } + } + } + + pub fn compact(&mut self) -> io::Result<()> { + self.commit()?; + let filepath = self.filepath.clone(); + let header = with_exclusive_sidecar(&filepath, || { + recover_pending_under_existing_lock(&filepath)?; + let mut query = BPlusTreeQuery::::from_file_unlocked(File::open(&filepath)?)?; + let _ = verify_full(&mut query)?; + let metadata = query.header.metadata.clone(); + let mut entries = BTreeMap::new(); + for entry in query.iter() { + let (key, value) = entry?; + let _ = entries.insert(key, value); + } + drop(query); + + let mut replacement = BPlusTree { entries, metadata, dirty: true }; + let _ = replacement.store_exclusive(&filepath)?; + let file = File::open(&filepath)?; + let mut page = [0; PAGE_SIZE]; + read_exact_at_offset(&file, &mut page, 0)?; + DatabaseHeader::decode(&page) + })?; + self.database_id = header.database_id; + self.verified_generation = header.generation; + self.verified_next_page_id = header.next_page_id; + Ok(()) + } +} + +pub struct BPlusTreeSerialWriter { + updater: Arc>>, + flush_policy: FlushPolicy, + dirty: Arc, + shutdown: Arc, + background_error: Arc>>, + background_handle: Mutex>>, +} + +impl BPlusTreeSerialWriter +where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, + V: Serialize + for<'de> Deserialize<'de> + 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: Arc::new(Mutex::new(updater)), + flush_policy, + dirty: Arc::new(AtomicBool::new(false)), + shutdown: Arc::new(AtomicBool::new(false)), + background_error: Arc::new(Mutex::new(None)), + background_handle: Mutex::new(None), + }) + } + + pub fn upsert_prepared(&self, items: Vec<(K, Vec)>) -> io::Result { + let result = self.updater.lock().upsert_batch_encoded(items); + if result.is_ok() { + self.dirty.store(self.flush_policy == FlushPolicy::Batch, Ordering::Release); + } + result + } + + pub fn upsert(&self, items: &[(&K, &V)]) -> io::Result { + self.upsert_prepared(BPlusTreeUpdate::::prepare_upsert_batch(items)?) + } + + pub fn start_background_commit(&self, interval: Duration) -> io::Result<()> { + if self.flush_policy != FlushPolicy::Batch { + return Err(invalid_input("background commit requires batch flush policy")); + } + if interval.is_zero() { + return Err(invalid_input("background commit interval must be greater than zero")); + } + let mut slot = self.background_handle.lock(); + if slot.is_some() { + return Ok(()); + } + self.shutdown.store(false, Ordering::Release); + let updater = Arc::clone(&self.updater); + let dirty = Arc::clone(&self.dirty); + let shutdown = Arc::clone(&self.shutdown); + let background_error = Arc::clone(&self.background_error); + *slot = Some( + std::thread::Builder::new() + .name(String::from("bplustree-commit")) + .spawn(move || { + while !shutdown.load(Ordering::Acquire) { + std::thread::park_timeout(interval); + if shutdown.load(Ordering::Acquire) { + break; + } + commit_if_dirty(&updater, &dirty, &background_error); + } + commit_if_dirty(&updater, &dirty, &background_error); + }) + .map_err(io::Error::other)?, + ); + Ok(()) + } + + pub fn stop_background_commit(&self) -> io::Result<()> { + self.shutdown.store(true, Ordering::Release); + if let Some(handle) = self.background_handle.lock().take() { + handle.thread().unpark(); + handle.join().map_err(|_| io::Error::other("background B+Tree commit thread panicked"))?; + } + if let Some(error) = self.background_error.lock().take() { + return Err(error); + } + Ok(()) + } + + pub fn flush_now(&self) -> io::Result<()> { self.commit() } + + pub fn commit(&self) -> io::Result<()> { + self.updater.lock().commit()?; + self.dirty.store(false, Ordering::Release); + Ok(()) + } + + pub fn shutdown(&self) -> io::Result<()> { + self.stop_background_commit()?; + self.commit() + } +} + +fn commit_if_dirty( + updater: &Mutex>, + dirty: &AtomicBool, + background_error: &Mutex>, +) where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de>, +{ + if !dirty.swap(false, Ordering::AcqRel) { + return; + } + if let Err(error) = updater.lock().commit() { + dirty.store(true, Ordering::Release); + log::error!("Background B+Tree commit failed: {error}"); + *background_error.lock() = Some(error); + } +} + +impl Drop for BPlusTreeSerialWriter { + fn drop(&mut self) { + self.shutdown.store(true, Ordering::Release); + if self.dirty.load(Ordering::Acquire) { + log::warn!("Dropping dirty B+Tree serial writer without an explicit shutdown"); + } + if let Some(handle) = self.background_handle.lock().take() { + handle.thread().unpark(); + drop(handle); + } + } +} + +fn minimum_tree_key( + transaction: &WriteTransaction, + base: &[u8], + mut page_id: u64, +) -> io::Result>> { + let mut visited = HashSet::new(); + loop { + if !visited.insert(page_id) { + return Err(invalid_data("minimum-key descent contains a cycle")); + } + let page = SlottedPage::open(transaction.page(base, page_id)?, page_id, transaction.next_header.next_page_id)?; + match page.header().page_type { + PageType::Leaf => { + let Some(cell) = page.cells().next() else { return Ok(None) }; + return LeafCellRef::decode(cell?, page_id, transaction.next_header.next_page_id) + .map(|cell| Some(cell.key_bytes.to_vec())); + } + PageType::Internal => { + page_id = InternalPreamble::decode(page.as_bytes(), page_id, transaction.next_header.next_page_id)? + .leftmost_child; + } + PageType::Overflow | PageType::Free => return Err(invalid_data("tree references a non-tree page")), + } + } +} + +fn validate_leaf_backlink( + transaction: &WriteTransaction, + base: &[u8], + page_id: u64, + sibling_id: u64, + sibling_points_left: bool, +) -> io::Result<()> { + if sibling_id == 0 { + return Ok(()); + } + let sibling = SlottedPage::open( + transaction.page(base, sibling_id)?, + sibling_id, + transaction.next_header.next_page_id, + )?; + let backlink = if sibling_points_left { sibling.header().left } else { sibling.header().right }; + if sibling.header().page_type != PageType::Leaf || backlink != page_id { + return Err(invalid_data("asymmetric leaf sibling link")); + } + Ok(()) +} + +fn validate_transaction_free_list(transaction: &WriteTransaction, base: &[u8]) -> io::Result<()> { + let mut free = HashSet::new(); + let mut page_id = transaction.next_header.free_page_head; + while page_id != 0 { + if !free.insert(page_id) { + return Err(invalid_data("free list contains a cycle")); + } + if transaction.allocated_pages.contains(&page_id) { + return Err(invalid_data("allocated page remains linked from free list")); + } + let page = SlottedPage::open(transaction.page(base, page_id)?, page_id, transaction.next_header.next_page_id)?; + if page.header().page_type != PageType::Free { + return Err(invalid_data("free list references a non-free page")); + } + page_id = page.header().right; + } + if !transaction.freed_pages.iter().all(|page| free.contains(page)) { + return Err(invalid_data("freed page is missing from free list")); + } + for page_id in &transaction.allocated_pages { + let page = transaction.page(base, *page_id)?; + if SlottedPage::open(page, *page_id, transaction.next_header.next_page_id)?.header().page_type == PageType::Free { + return Err(invalid_data("allocated page still has free-page type")); + } + } + Ok(()) +} + +fn validate_transaction_links(transaction: &WriteTransaction, base: &[u8]) -> io::Result<()> +where + K: Ord + for<'de> Deserialize<'de>, +{ + let original_length = transaction + .original_header + .next_page_id + .checked_mul(u64::try_from(PAGE_SIZE).map_err(|_| invalid_data("page size exceeds u64"))?) + .ok_or_else(|| invalid_data("original database length overflow"))?; + if original_length != transaction.original_file_len { + return Err(invalid_data("transaction original length does not match header")); + } + + let mut owned_overflow_pages = HashSet::new(); + for (page_id, bytes) in &transaction.dirty_pages { + if *page_id == 0 { + continue; + } + let page = SlottedPage::open(bytes.as_slice(), *page_id, transaction.next_header.next_page_id)?; + match page.header().page_type { + PageType::Leaf => { + validate_leaf_backlink(transaction, base, *page_id, page.header().left, false)?; + validate_leaf_backlink(transaction, base, *page_id, page.header().right, true)?; + let mut previous_key = None; + for cell in page.cells() { + let cell = LeafCellRef::decode(cell?, *page_id, transaction.next_header.next_page_id)?; + let key = decode_key::(cell.key_bytes)?; + if previous_key.as_ref().is_some_and(|previous| previous >= &key) { + return Err(invalid_data("leaf keys are not strictly increasing")); + } + previous_key = Some(key); + if let LeafValueRef::Overflow { stored_len, head, crc32, .. } = cell.value { + for overflow_page in + validated_overflow_chain_pages(transaction, base, head, stored_len, crc32)? + { + if !owned_overflow_pages.insert(overflow_page) { + return Err(invalid_data("overflow page is owned by multiple values")); + } + } + } + } + } + PageType::Internal => { + let preamble = InternalPreamble::decode(bytes.as_slice(), *page_id, transaction.next_header.next_page_id)?; + let child = SlottedPage::open( + transaction.page(base, preamble.leftmost_child)?, + preamble.leftmost_child, + transaction.next_header.next_page_id, + )?; + if !matches!(child.header().page_type, PageType::Leaf | PageType::Internal) { + return Err(invalid_data("internal page references a non-tree child")); + } + let mut previous_key = None; + for cell in page.cells() { + let cell = InternalCellRef::decode(cell?, *page_id, transaction.next_header.next_page_id)?; + let key = decode_key::(cell.key_bytes)?; + if previous_key.as_ref().is_some_and(|previous| previous >= &key) { + return Err(invalid_data("internal separator keys are not strictly increasing")); + } + previous_key = Some(key); + let minimum = minimum_tree_key(transaction, base, cell.right_child)? + .ok_or_else(|| invalid_data("internal separator references an empty subtree"))?; + if minimum != cell.key_bytes { + return Err(invalid_data("internal separator differs from right subtree minimum")); + } + } + } + PageType::Overflow => { + if page.header().right != 0 { + let next = SlottedPage::open( + transaction.page(base, page.header().right)?, + page.header().right, + transaction.next_header.next_page_id, + )?; + if next.header().page_type != PageType::Overflow { + return Err(invalid_data("overflow page references a non-overflow page")); + } + } + } + PageType::Free => {} + } + } + + validate_transaction_free_list(transaction, base) +} + +#[derive(Debug)] +struct Promotion { + key: Vec, + right_child: u64, +} + +fn internal_child_position Deserialize<'de>, B: AsRef<[u8]>>( + page: &SlottedPage, + key: &K, +) -> io::Result<(usize, u64)> { + let mut left = 0usize; + let mut right = usize::from(page.header().cell_count); + while left < right { + let middle = left + (right - left) / 2; + let cell = InternalCellRef::decode(page.cell(middle)?, page.page_id(), page.next_page_id())?; + if decode_key::(cell.key_bytes)? <= *key { + left = middle + 1; + } else { + right = middle; + } + } + let child = if left == 0 { + InternalPreamble::decode(page.as_bytes(), page.page_id(), page.next_page_id())?.leftmost_child + } else { + InternalCellRef::decode(page.cell(left - 1)?, page.page_id(), page.next_page_id())?.right_child + }; + Ok((left, child)) +} + +fn locate_transaction_leaf Deserialize<'de>>( + transaction: &WriteTransaction, + base: &[u8], + key: &K, +) -> io::Result<(u64, Vec<(u64, usize)>)> { + let mut page_id = transaction.next_header.root_page_id; + let mut path = Vec::new(); + let mut visited = HashSet::new(); + loop { + if !visited.insert(page_id) { + return Err(invalid_data("tree descent contains a cycle")); + } + let page = SlottedPage::open(transaction.page(base, page_id)?, page_id, transaction.next_header.next_page_id)?; + match page.header().page_type { + PageType::Leaf => return Ok((page_id, path)), + PageType::Internal => { + let (position, child) = internal_child_position(&page, key)?; + path.try_reserve(1) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + path.push((page_id, position)); + page_id = child; + } + PageType::Overflow | PageType::Free => return Err(invalid_data("tree references a non-tree page")), + } + } +} + +fn overflow_chain_pages( + transaction: &WriteTransaction, + base: &[u8], + mut page_id: u64, +) -> io::Result> { + let mut pages = Vec::new(); + let mut visited = HashSet::new(); + while page_id != 0 { + if !visited.insert(page_id) { + return Err(invalid_data("overflow chain contains a cycle")); + } + let page = SlottedPage::open(transaction.page(base, page_id)?, page_id, transaction.next_header.next_page_id)?; + if page.header().page_type != PageType::Overflow || overflow_payload(&page)?.is_empty() { + return Err(invalid_data("invalid overflow chain page")); + } + pages + .try_reserve(1) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + pages.push(page_id); + page_id = page.header().right; + } + if pages.is_empty() { + return Err(invalid_data("overflow chain is empty")); + } + Ok(pages) +} + +fn validated_overflow_chain_pages( + transaction: &WriteTransaction, + base: &[u8], + head: u64, + stored_len: u32, + crc32: u32, +) -> io::Result> { + let pages = overflow_chain_pages(transaction, base, head)?; + let expected = usize::try_from(stored_len).map_err(|_| invalid_data("stored length exceeds usize"))?; + let mut actual = 0usize; + let mut hasher = crc32fast::Hasher::new(); + for page_id in &pages { + let page = SlottedPage::open( + transaction.page(base, *page_id)?, + *page_id, + transaction.next_header.next_page_id, + )?; + let payload = overflow_payload(&page)?; + actual = actual + .checked_add(payload.len()) + .ok_or_else(|| invalid_data("overflow chain length overflow"))?; + if actual > expected { + return Err(invalid_data("overflow chain exceeds declared length")); + } + hasher.update(payload); + } + if actual != expected || hasher.finalize() != crc32 { + return Err(invalid_data("overflow value checksum or length mismatch")); + } + Ok(pages) +} + +fn free_pages(transaction: &mut WriteTransaction, pages: &[u64]) -> io::Result<()> { + for page_id in pages.iter().rev() { + transaction.free_page(*page_id)?; + } + Ok(()) +} + +fn write_overflow_chain( + transaction: &mut WriteTransaction, + base: &[u8], + stored: &[u8], + old_pages: &[u64], +) -> io::Result { + let required = stored.len().div_ceil(OVERFLOW_PAYLOAD_LEN); + if required == 0 { + return Err(invalid_input("overflow value must not be empty")); + } + let mut pages = Vec::new(); + pages + .try_reserve_exact(required) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + if old_pages.len() >= required { + pages.extend_from_slice( + old_pages + .get(..required) + .ok_or_else(|| invalid_data("overflow reuse range is invalid"))?, + ); + free_pages( + transaction, + old_pages + .get(required..) + .ok_or_else(|| invalid_data("overflow tail range is invalid"))?, + )?; + } else { + for _ in 0..required { + pages.push(transaction.allocate_page(base)?); + } + free_pages(transaction, old_pages)?; + } + for (index, payload) in stored.chunks(OVERFLOW_PAYLOAD_LEN).enumerate() { + let page_id = *pages.get(index).ok_or_else(|| invalid_data("overflow page allocation is missing"))?; + let next = pages.get(index + 1).copied().unwrap_or(0); + transaction.write_page( + page_id, + encode_overflow_page(page_id, transaction.next_header.next_page_id, next, payload)?, + )?; + } + pages.first().copied().ok_or_else(|| invalid_data("overflow head is missing")) +} + +fn repair_right_leaf_backlink( + transaction: &mut WriteTransaction, + base: &[u8], + right_page_id: u64, + left_page_id: u64, +) -> io::Result<()> { + if right_page_id == 0 { + return Ok(()); + } + let next_page_id = transaction.next_header.next_page_id; + let page = transaction.page_mut(base, right_page_id)?; + let mut header = SlottedPage::open(page.as_slice(), right_page_id, next_page_id)?.header(); + if header.page_type != PageType::Leaf { + return Err(invalid_data("leaf sibling references a non-leaf page")); + } + header.left = left_page_id; + header.encode_into(page, right_page_id, next_page_id) +} + +fn mutate_leaf( + transaction: &mut WriteTransaction, + base: &[u8], + leaf_id: u64, + index: usize, + replace: bool, + cell: &[u8], +) -> io::Result> { + let snapshot = transaction.page_copy(base, leaf_id)?; + let page = SlottedPage::open(snapshot.as_slice(), leaf_id, transaction.next_header.next_page_id)?; + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("mutation target is not a leaf page")); + } + let count = usize::from(page.header().cell_count); + if (replace && index >= count) || (!replace && index > count) { + return Err(invalid_input("leaf mutation index is outside page")); + } + if replace && page.cell(index)?.len() == cell.len() { + let next_page_id = transaction.next_header.next_page_id; + let dirty = transaction.page_mut(base, leaf_id)?; + return SlottedPage::open(dirty.as_mut_slice(), leaf_id, next_page_id)? + .replace_same_len(index, cell) + .map(|()| None); + } + let mut cells = Vec::<&[u8]>::new(); + cells + .try_reserve_exact(count + usize::from(!replace)) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for current in 0..count { + if current == index { + cells.push(cell); + if replace { + continue; + } + } + cells.push(page.cell(current)?); + } + if index == count { + cells.push(cell); + } + if used_leaf_bytes(&cells)? <= PAGE_SIZE { + let next_page_id = transaction.next_header.next_page_id; + let dirty = transaction.page_mut(base, leaf_id)?; + SlottedPage::open(dirty.as_mut_slice(), leaf_id, next_page_id)? + .rebuild_ordered(cells.iter().copied())?; + return Ok(None); + } + + let new_leaf_id = transaction.allocate_page(base)?; + let boundary = choose_leaf_split(&cells)?; + let (left_cells, right_cells) = cells.split_at(boundary); + let right = page.header().right; + let left_page = leaf_page(leaf_id, page.header().left, new_leaf_id, left_cells)?; + let right_page = leaf_page(new_leaf_id, leaf_id, right, right_cells)?; + let separator = LeafCellRef::decode( + right_cells.first().ok_or_else(|| invalid_data("split right leaf is empty"))?, + new_leaf_id, + transaction.next_header.next_page_id, + )? + .key_bytes + .to_vec(); + transaction.write_page(leaf_id, left_page)?; + transaction.write_page(new_leaf_id, right_page)?; + repair_right_leaf_backlink(transaction, base, right, new_leaf_id)?; + Ok(Some(Promotion { key: separator, right_child: new_leaf_id })) +} + +fn insert_internal_promotion( + transaction: &mut WriteTransaction, + base: &[u8], + page_id: u64, + position: usize, + promotion: &Promotion, + cell_scratch: &mut Vec, +) -> io::Result> { + let snapshot = transaction.page_copy(base, page_id)?; + let page = SlottedPage::open(snapshot.as_slice(), page_id, transaction.next_header.next_page_id)?; + if page.header().page_type != PageType::Internal { + return Err(invalid_data("promotion target is not an internal page")); + } + let count = usize::from(page.header().cell_count); + if position > count { + return Err(invalid_input("internal insertion position is outside page")); + } + encode_internal_cell( + &promotion.key, + promotion.right_child, + page_id, + transaction.next_header.next_page_id, + cell_scratch, + )?; + let mut cells = Vec::<&[u8]>::new(); + cells + .try_reserve_exact(count + 1) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for current in 0..count { + if current == position { + cells.push(cell_scratch); + } + cells.push(page.cell(current)?); + } + if position == count { + cells.push(cell_scratch); + } + if used_internal_bytes(cells.iter().copied())? <= PAGE_SIZE { + let next_page_id = transaction.next_header.next_page_id; + let dirty = transaction.page_mut(base, page_id)?; + SlottedPage::open(dirty.as_mut_slice(), page_id, next_page_id)? + .rebuild_ordered(cells.iter().copied())?; + return Ok(None); + } + + let right_page_id = transaction.allocate_page(base)?; + let split = choose_internal_split(&cells, page_id, transaction.next_header.next_page_id)?; + let leftmost = InternalPreamble::decode(snapshot.as_slice(), page_id, transaction.next_header.next_page_id)? + .leftmost_child; + let promoted = Promotion { key: split.promoted.key_bytes.to_vec(), right_child: right_page_id }; + let left_page = internal_page(page_id, leftmost, split.left_cells)?; + let right_page = internal_page(right_page_id, split.right_leftmost_child, split.right_cells)?; + transaction.write_page(page_id, left_page)?; + transaction.write_page(right_page_id, right_page)?; + Ok(Some(promoted)) +} + +fn propagate_promotion( + transaction: &mut WriteTransaction, + base: &[u8], + mut path: Vec<(u64, usize)>, + mut promotion: Promotion, + cell_scratch: &mut Vec, +) -> io::Result<()> { + while let Some((parent, position)) = path.pop() { + let Some(next) = insert_internal_promotion(transaction, base, parent, position, &promotion, cell_scratch)? else { + return Ok(()); + }; + promotion = next; + } + let old_root = transaction.next_header.root_page_id; + let new_root = transaction.allocate_page(base)?; + encode_internal_cell( + &promotion.key, + promotion.right_child, + new_root, + transaction.next_header.next_page_id, + cell_scratch, + )?; + transaction.write_page(new_root, internal_page(new_root, old_root, &[cell_scratch.as_slice()])?)?; + transaction.next_header.root_page_id = new_root; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn stage_upsert Deserialize<'de>>( + transaction: &mut WriteTransaction, + base: &[u8], + key: &K, + encoded_key: &[u8], + logical_len: u32, + compression: Compression, + stored: &[u8], + cell_scratch: &mut Vec, +) -> io::Result<()> { + let (leaf_id, path) = locate_transaction_leaf(transaction, base, key)?; + let page = SlottedPage::open(transaction.page(base, leaf_id)?, leaf_id, transaction.next_header.next_page_id)?; + let search = search_leaf(&page, key)?; + let (replace, index, old_overflow) = match search { + Ok(index) => { + let old = LeafCellRef::decode(page.cell(index)?, leaf_id, transaction.next_header.next_page_id)?; + let overflow = match old.value { + LeafValueRef::Overflow { stored_len, head, crc32, .. } => { + validated_overflow_chain_pages(transaction, base, head, stored_len, crc32)? + } + LeafValueRef::Inline { .. } | LeafValueRef::Tombstone => Vec::new(), + }; + (true, index, overflow) + } + Err(index) => (false, index, Vec::new()), + }; + let inline_footprint = SLOT_LEN + .checked_add(24) + .and_then(|size| size.checked_add(encoded_key.len())) + .and_then(|size| size.checked_add(stored.len())) + .ok_or_else(|| invalid_input("leaf cell footprint overflow"))?; + if stored.len() <= MAX_INLINE_STORED_VALUE && inline_footprint <= MAX_CELL_FOOTPRINT { + free_pages(transaction, &old_overflow)?; + encode_inline_leaf_cell(encoded_key, logical_len, compression, stored, cell_scratch)?; + } else { + let head = write_overflow_chain(transaction, base, stored, &old_overflow)?; + encode_overflow_leaf_cell( + encoded_key, + logical_len, + compression, + u32::try_from(stored.len()).map_err(|_| invalid_input("stored value exceeds u32"))?, + head, + stored_value_checksum(stored), + leaf_id, + transaction.next_header.next_page_id, + cell_scratch, + )?; + } + if let Some(promotion) = mutate_leaf(transaction, base, leaf_id, index, replace, cell_scratch)? { + propagate_promotion(transaction, base, path, promotion, cell_scratch)?; + } + Ok(()) +} + +fn stage_delete Deserialize<'de>>( + transaction: &mut WriteTransaction, + base: &[u8], + key: &K, + encoded_key: &[u8], + cell_scratch: &mut Vec, +) -> io::Result { + let (leaf_id, _) = locate_transaction_leaf(transaction, base, key)?; + let page = SlottedPage::open(transaction.page(base, leaf_id)?, leaf_id, transaction.next_header.next_page_id)?; + let Ok(index) = search_leaf(&page, key)? else { return Ok(false) }; + let old = LeafCellRef::decode(page.cell(index)?, leaf_id, transaction.next_header.next_page_id)?; + match old.value { + LeafValueRef::Tombstone => return Ok(false), + LeafValueRef::Overflow { stored_len, head, crc32, .. } => { + let pages = validated_overflow_chain_pages(transaction, base, head, stored_len, crc32)?; + free_pages(transaction, &pages)?; + } + LeafValueRef::Inline { .. } => {} + } + encode_tombstone_leaf_cell(encoded_key, cell_scratch)?; + let _ = mutate_leaf(transaction, base, leaf_id, index, true, cell_scratch)?; + Ok(true) +} + +fn query_transaction( + transaction: &WriteTransaction, + base: &[u8], + key: &K, + scratch: &mut Vec, +) -> io::Result> +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + let (leaf_id, _) = locate_transaction_leaf(transaction, base, key)?; + let page = SlottedPage::open(transaction.page(base, leaf_id)?, leaf_id, transaction.next_header.next_page_id)?; + let Ok(index) = search_leaf(&page, key)? else { return Ok(None) }; + let cell = LeafCellRef::decode(page.cell(index)?, leaf_id, transaction.next_header.next_page_id)?; + let bytes = match cell.value { + LeafValueRef::Tombstone => return Ok(None), + LeafValueRef::Inline { compression: Compression::None, stored, .. } => stored, + LeafValueRef::Inline { compression: Compression::Lz4, logical_len, stored, .. } => { + decompress_value_into(stored, logical_len, transaction_value_limit(transaction)?, scratch)? + } + LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } => { + read_transaction_overflow( + transaction, + base, + compression, + logical_len, + stored_len, + head, + crc32, + scratch, + )? + } + }; + binary_deserialize(bytes).map(Some) +} + +fn transaction_value_limit(transaction: &WriteTransaction) -> io::Result { + usize::try_from(transaction.next_header.next_page_id) + .map_err(|_| invalid_data("next page id exceeds usize"))? + .checked_mul(PAGE_SIZE) + .and_then(|size| size.checked_mul(256)) + .map(|size| size.min(usize::try_from(u32::MAX).unwrap_or(usize::MAX))) + .ok_or_else(|| invalid_data("value allocation limit overflow")) +} + +#[allow(clippy::too_many_arguments)] +fn read_transaction_overflow<'a>( + transaction: &WriteTransaction, + base: &[u8], + compression: Compression, + logical_len: u32, + stored_len: u32, + head: u64, + crc32: u32, + scratch: &'a mut Vec, +) -> io::Result<&'a [u8]> { + let stored_len = usize::try_from(stored_len).map_err(|_| invalid_data("stored length exceeds usize"))?; + if stored_len > transaction_value_limit(transaction)? { + return Err(invalid_data("overflow value exceeds allocation limit")); + } + scratch.clear(); + scratch + .try_reserve(stored_len) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for page_id in overflow_chain_pages(transaction, base, head)? { + let page = SlottedPage::open(transaction.page(base, page_id)?, page_id, transaction.next_header.next_page_id)?; + let payload = overflow_payload(&page)?; + if scratch.len().saturating_add(payload.len()) > stored_len { + return Err(invalid_data("overflow chain exceeds declared length")); + } + scratch.extend_from_slice(payload); + } + if scratch.len() != stored_len || crc32fast::hash(scratch) != crc32 { + return Err(invalid_data("overflow value checksum or length mismatch")); + } + if compression == Compression::Lz4 { + decompress_value_in_place(scratch, logical_len, transaction_value_limit(transaction)?)?; + } else if scratch.len() != usize::try_from(logical_len).map_err(|_| invalid_data("logical length exceeds usize"))? { + return Err(invalid_data("uncompressed overflow length mismatch")); + } + Ok(scratch) +} + +enum DecodedEntry { + Ready(K, V), + Overflow(K, Compression, u32, u32, u64, u32), + Tombstone, +} + +enum DecodedValue { + Ready(V), + Overflow(Compression, u32, u32, u64, u32), + Tombstone, +} + +struct InternalRoute { + leftmost_child: u64, + separators: Vec<(K, u64)>, +} + +impl InternalRoute { + fn child_for(&self, target: &K) -> u64 { + let index = self.separators.partition_point(|(key, _)| key <= target); + index.checked_sub(1).map_or(self.leftmost_child, |previous| self.separators[previous].1) + } +} + +fn decode_internal_route(page: &SlottedPage) -> io::Result> +where + K: Ord + for<'de> Deserialize<'de>, + B: AsRef<[u8]>, +{ + let preamble = InternalPreamble::decode(page.as_bytes(), page.page_id(), page.next_page_id())?; + let count = usize::from(page.header().cell_count); + let mut separators = Vec::new(); + separators + .try_reserve_exact(count) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for index in 0..count { + let cell = InternalCellRef::decode(page.cell(index)?, page.page_id(), page.next_page_id())?; + separators.push((decode_internal_key(cell.key_bytes)?, cell.right_child)); + } + Ok(InternalRoute { leftmost_child: preamble.leftmost_child, separators }) +} + +enum LocateCell { + Internal(InternalRoute), + Leaf(Option>), +} + +enum LocateLeaf { + Internal(InternalRoute), + Leaf, +} + +struct QuerySnapshot { + file: Option, + mmap: Option, + filepath: PathBuf, + page_validations: Vec>, + internal_routes: Vec>>, + sidecar_guard: Option, +} + +pub struct BPlusTreeQuery { + snapshot: Arc>, + header: DatabaseHeader, + file_len: usize, + page_buffer: Vec, + value_scratch: Vec, + locator_page_id: u64, + locator_cell_ranges: Vec>, + _value: PhantomData, +} + +impl BPlusTreeQuery { + fn from_file_unlocked(file: File) -> io::Result { + let file_len = usize::try_from(file.metadata()?.len()).map_err(|_| invalid_data("database length exceeds usize"))?; + if file_len < PAGE_SIZE { + return Err(invalid_data("database is shorter than its header page")); + } + let mut header_page = [0; PAGE_SIZE]; + read_exact_at_offset(&file, &mut header_page, 0)?; + let header = DatabaseHeader::decode(&header_page)?; + let expected = usize::try_from(header.next_page_id) + .map_err(|_| invalid_data("next page id exceeds usize"))? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| invalid_data("database length overflow"))?; + if file_len != expected { + return Err(invalid_data("database file length does not match header")); + } + let mmap = mmap_with_advice(&file, Advice::Normal, "v3 B+Tree query"); + let page_count = usize::try_from(header.next_page_id).map_err(|_| invalid_data("next page id exceeds usize"))?; + let mut page_validations = Vec::new(); + page_validations + .try_reserve_exact(page_count) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + let mut internal_routes = Vec::new(); + internal_routes + .try_reserve_exact(page_count) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for _ in 0..page_count { + page_validations.push(OnceLock::new()); + internal_routes.push(OnceLock::new()); + } + let file = mmap.is_none().then_some(file); + let mut page_buffer = Vec::new(); + if mmap.is_none() { + page_buffer + .try_reserve_exact(PAGE_SIZE) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + page_buffer.resize(PAGE_SIZE, 0); + } + Ok(Self { + snapshot: Arc::new(QuerySnapshot { + file, + mmap, + filepath: PathBuf::new(), + page_validations, + internal_routes, + sidecar_guard: None, + }), + header, + file_len, + page_buffer, + value_scratch: Vec::new(), + locator_page_id: 0, + locator_cell_ranges: Vec::new(), + _value: PhantomData, + }) + } + + pub fn try_new(filepath: &Path) -> io::Result { + loop { + let sidecar_guard = SharedSidecarGuard::acquire(filepath)?; + let pending = wal_path(filepath).try_exists()? || wal_temporary_path(filepath).try_exists()?; + if !pending { + let mut query = Self::from_file_unlocked(File::open(filepath)?)?; + let snapshot = Arc::get_mut(&mut query.snapshot) + .ok_or_else(|| invalid_data("new query snapshot is unexpectedly shared"))?; + snapshot.filepath = filepath.to_path_buf(); + snapshot.sidecar_guard = Some(sidecar_guard); + return Ok(query); + } + drop(sidecar_guard); + match recover_pending(filepath) { + Ok(()) => {} + Err(error) + if matches!( + error.kind(), + io::ErrorKind::PermissionDenied | io::ErrorKind::ReadOnlyFilesystem + ) => + { + return Err(recovery_required(filepath, error)); + } + Err(error) => return Err(error), + } + } + } + + pub fn try_clone(&self) -> io::Result { + if self.snapshot.filepath.as_os_str().is_empty() { + return Err(invalid_input("mapped query without a path cannot be cloned")); + } + let mut page_buffer = Vec::new(); + if self.snapshot.mmap.is_none() { + page_buffer + .try_reserve_exact(PAGE_SIZE) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + page_buffer.resize(PAGE_SIZE, 0); + } + Ok(Self { + snapshot: Arc::clone(&self.snapshot), + header: self.header.clone(), + file_len: self.file_len, + page_buffer, + value_scratch: Vec::new(), + locator_page_id: 0, + locator_cell_ranges: Vec::new(), + _value: PhantomData, + }) + } + + #[cfg(test)] + pub(crate) fn clone_error_fixture() -> Self { + Self { + snapshot: Arc::new(QuerySnapshot { + file: None, + mmap: None, + filepath: PathBuf::new(), + page_validations: Vec::new(), + internal_routes: Vec::new(), + sidecar_guard: None, + }), + header: DatabaseHeader { + root_page_id: 1, + next_page_id: 2, + free_page_head: 0, + generation: 1, + database_id: [0; 16], + metadata: BPlusTreeMetadata::Empty, + }, + file_len: 0, + page_buffer: Vec::new(), + value_scratch: Vec::new(), + locator_page_id: 0, + locator_cell_ranges: Vec::new(), + _value: PhantomData, + } + } + + pub fn filepath(&self) -> &Path { &self.snapshot.filepath } + + pub(crate) fn snapshot_identity(&self) -> ([u8; 16], u64) { + (self.header.database_id, self.header.generation) + } + + pub(crate) fn snapshot_metadata(&self) -> &BPlusTreeMetadata { &self.header.metadata } + + fn value_allocation_limit(&self) -> usize { + self.file_len + .saturating_mul(256) + .min(usize::try_from(u32::MAX).unwrap_or(usize::MAX)) + } + + fn assemble_overflow_chain( + &mut self, + compression: Compression, + logical_len: u32, + stored_len: u32, + mut page_id: u64, + crc32: u32, + mut owned_pages: Option<&mut HashSet>, + ) -> io::Result<()> { + let stored_len = usize::try_from(stored_len).map_err(|_| invalid_data("stored length exceeds usize"))?; + let logical_len_usize = + usize::try_from(logical_len).map_err(|_| invalid_data("logical length exceeds usize"))?; + if stored_len > self.file_len || logical_len_usize > self.value_allocation_limit() { + return Err(invalid_data("overflow value exceeds allocation limit")); + } + self.value_scratch.clear(); + self.value_scratch + .try_reserve(stored_len) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + let mut chain = HashSet::new(); + while page_id != 0 { + record_page_visit(&mut chain, page_id, "overflow chain contains a cycle")?; + if let Some(pages) = owned_pages.as_deref_mut() { + record_page_visit(pages, page_id, "overflow page is owned by multiple values")?; + } + page_id = self.with_slotted_page(page_id, |page, scratch| { + if page.header().page_type != PageType::Overflow { + return Err(invalid_data("overflow chain references a non-overflow page")); + } + let payload = overflow_payload(page)?; + if payload.is_empty() { + return Err(invalid_data("overflow chain contains an empty payload")); + } + let next_len = scratch + .len() + .checked_add(payload.len()) + .ok_or_else(|| invalid_data("overflow value length overflow"))?; + if next_len > stored_len { + return Err(invalid_data("overflow chain exceeds declared length")); + } + scratch.extend_from_slice(payload); + Ok(page.header().right) + })?; + } + if self.value_scratch.len() != stored_len || crc32fast::hash(&self.value_scratch) != crc32 { + return Err(invalid_data("overflow value checksum or length mismatch")); + } + if compression == Compression::Lz4 { + let allocation_limit = self.value_allocation_limit(); + decompress_value_in_place(&mut self.value_scratch, logical_len, allocation_limit)?; + } else if self.value_scratch.len() != logical_len_usize { + return Err(invalid_data("uncompressed overflow length mismatch")); + } + Ok(()) + } + + fn with_page(&mut self, page_id: u64, read: impl FnOnce(&[u8], &mut Vec) -> io::Result) -> io::Result { + if page_id == 0 || page_id >= self.header.next_page_id { + return Err(invalid_data("page id is outside database")); + } + let offset = usize::try_from(page_id) + .map_err(|_| invalid_data("page id exceeds usize"))? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| invalid_data("page offset overflow"))?; + let end = offset.checked_add(PAGE_SIZE).ok_or_else(|| invalid_data("page end overflow"))?; + if let Some(mmap) = &self.snapshot.mmap { + let page = mmap.get(offset..end).ok_or_else(|| invalid_data("page is truncated"))?; + return read(page, &mut self.value_scratch); + } + let file = self.snapshot.file.as_ref().ok_or_else(|| invalid_data("query has no data source"))?; + if self.page_buffer.len() != PAGE_SIZE { + return Err(invalid_data("query page buffer has invalid length")); + } + read_exact_at_offset( + file, + &mut self.page_buffer, + u64::try_from(offset).map_err(|_| invalid_data("page offset exceeds u64"))?, + )?; + read(&self.page_buffer, &mut self.value_scratch) + } + + fn with_slotted_page( + &mut self, + page_id: u64, + read: impl FnOnce(&SlottedPage<&[u8]>, &mut Vec) -> io::Result, + ) -> io::Result { + let index = usize::try_from(page_id).map_err(|_| invalid_data("page id exceeds usize"))?; + let cached = self + .snapshot + .page_validations + .get(index) + .ok_or_else(|| invalid_data("page id is outside cache"))? + .get() + .copied(); + let next_page_id = self.header.next_page_id; + let (result, validation) = self.with_page(page_id, |bytes, scratch| { + let page = match cached { + Some(validation) => SlottedPage::from_immutable_snapshot(bytes, validation)?, + None => SlottedPage::open(bytes, page_id, next_page_id)?, + }; + let validation = cached.is_none().then(|| page.validation()); + read(&page, scratch).map(|result| (result, validation)) + })?; + if let Some(validation) = validation { + let slot = self + .snapshot + .page_validations + .get(index) + .ok_or_else(|| invalid_data("page id is outside cache"))?; + let _ = slot.set(validation); + } + Ok(result) + } + + fn cached_internal_child(&self, page_id: u64, key: &K) -> io::Result> + where + K: Ord, + { + let index = usize::try_from(page_id).map_err(|_| invalid_data("page id exceeds usize"))?; + Ok(self + .snapshot + .internal_routes + .get(index) + .ok_or_else(|| invalid_data("page id is outside route cache"))? + .get() + .map(|route| route.child_for(key))) + } + + fn cache_internal_route(&self, page_id: u64, route: InternalRoute) -> io::Result<()> { + let index = usize::try_from(page_id).map_err(|_| invalid_data("page id exceeds usize"))?; + let slot = self + .snapshot + .internal_routes + .get(index) + .ok_or_else(|| invalid_data("page id is outside route cache"))?; + let _ = slot.set(route); + Ok(()) + } + + fn locate_leaf(&mut self, key: &K) -> io::Result + where + K: Ord + for<'de> Deserialize<'de>, + { + let mut page_id = self.header.root_page_id; + let mut depth = 0u64; + loop { + if depth >= self.header.next_page_id { + return Err(invalid_data("tree descent contains a cycle")); + } + if let Some(child) = self.cached_internal_child(page_id, key)? { + page_id = child; + depth += 1; + continue; + } + let step = self.with_slotted_page(page_id, |page, _| { + match page.header().page_type { + PageType::Leaf => Ok(LocateLeaf::Leaf), + PageType::Internal => decode_internal_route(page).map(LocateLeaf::Internal), + PageType::Overflow | PageType::Free => Err(invalid_data("tree references a non-tree page")), + } + })?; + match step { + LocateLeaf::Internal(route) => { + let child = route.child_for(key); + self.cache_internal_route(page_id, route)?; + page_id = child; + depth += 1; + } + LocateLeaf::Leaf => return Ok(page_id), + } + } + } + + fn locate_cell(&mut self, key: &K) -> io::Result)>> + where + K: Ord + for<'de> Deserialize<'de>, + { + let mut page_id = self.header.root_page_id; + let mut depth = 0u64; + loop { + if depth >= self.header.next_page_id { + return Err(invalid_data("tree descent contains a cycle")); + } + if let Some(child) = self.cached_internal_child(page_id, key)? { + page_id = child; + depth += 1; + continue; + } + let step = self.with_slotted_page(page_id, |page, _| { + match page.header().page_type { + PageType::Internal => decode_internal_route(page).map(LocateCell::Internal), + PageType::Leaf => search_leaf(page, key)? + .ok() + .map(|index| page.cell_range(index)) + .transpose() + .map(LocateCell::Leaf), + PageType::Overflow | PageType::Free => Err(invalid_data("tree references a non-tree page")), + } + })?; + match step { + LocateCell::Internal(route) => { + let child = route.child_for(key); + self.cache_internal_route(page_id, route)?; + page_id = child; + depth += 1; + } + LocateCell::Leaf(range) => return Ok(range.map(|range| (page_id, range))), + } + } + } + + fn leftmost_leaf(&mut self) -> io::Result + where + K: for<'de> Deserialize<'de>, + { + let mut page_id = self.header.root_page_id; + let mut depth = 0u64; + loop { + if depth >= self.header.next_page_id { + return Err(invalid_data("tree descent contains a cycle")); + } + let result = self.with_slotted_page(page_id, |page, _| { + match page.header().page_type { + PageType::Leaf => Ok(None), + PageType::Internal => InternalPreamble::decode(page.as_bytes(), page_id, page.next_page_id()) + .map(|preamble| Some(preamble.leftmost_child)), + PageType::Overflow | PageType::Free => Err(invalid_data("tree references a non-tree page")), + } + })?; + let Some(child) = result else { return Ok(page_id) }; + page_id = child; + depth += 1; + } + } + + fn decode_entry(&mut self, leaf_page_id: u64, slot_index: usize) -> io::Result> + where + K: for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, + { + let range = self.with_slotted_page(leaf_page_id, |page, _| { + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("iterator expected a leaf page")); + } + page.cell_range(slot_index) + })?; + self.decode_entry_range(leaf_page_id, range) + } + + fn decode_entry_range(&mut self, leaf_page_id: u64, range: Range) -> io::Result> + where + K: for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, + { + let next_page_id = self.header.next_page_id; + let allocation_limit = self.value_allocation_limit(); + let decoded = self.with_page(leaf_page_id, |bytes, scratch| { + let cell_bytes = bytes.get(range).ok_or_else(|| invalid_data("leaf cell is outside page"))?; + let cell = LeafCellRef::decode(cell_bytes, leaf_page_id, next_page_id)?; + let key = binary_deserialize(cell.key_bytes)?; + match cell.value { + LeafValueRef::Inline { compression: Compression::None, stored, .. } => { + binary_deserialize(stored).map(|value| DecodedEntry::Ready(key, value)) + } + LeafValueRef::Inline { compression: Compression::Lz4, logical_len, stored, .. } => { + let decompressed = decompress_value_into(stored, logical_len, allocation_limit, scratch)?; + binary_deserialize(decompressed).map(|value| DecodedEntry::Ready(key, value)) + } + LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } => { + Ok(DecodedEntry::Overflow(key, compression, logical_len, stored_len, head, crc32)) + } + LeafValueRef::Tombstone => Ok(DecodedEntry::Tombstone), + } + })?; + match decoded { + DecodedEntry::Ready(key, value) => Ok(Some((key, value))), + DecodedEntry::Tombstone => Ok(None), + DecodedEntry::Overflow(key, compression, logical_len, stored_len, page_id, crc32) => { + self.assemble_overflow_chain(compression, logical_len, stored_len, page_id, crc32, None)?; + let value = binary_deserialize(&self.value_scratch)?; + Ok(Some((key, value))) + } + } + } +} + +impl BPlusTreeQuery +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + fn query_io(&mut self, key: &K) -> io::Result> { + let Some((leaf_id, range)) = self.locate_cell(key)? else { return Ok(None) }; + self.decode_entry_range(leaf_id, range).map(|entry| entry.map(|(_, value)| value)) + } + + pub fn query(&mut self, key: &K) -> Result, BPlusTreeError> { self.query_io(key).map_err(Into::into) } + + pub fn query_zero_copy(&mut self, key: &K) -> Result, BPlusTreeError> { self.query(key) } + + fn locator_cell_range(&mut self, locator: Locator) -> io::Result> { + if self.locator_page_id != locator.leaf_page_id { + let page_id = locator.leaf_page_id; + let mut ranges = std::mem::take(&mut self.locator_cell_ranges); + self.with_slotted_page(page_id, |page, _| { + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("locator does not reference a leaf page")); + } + let count = usize::from(page.header().cell_count); + ranges.clear(); + ranges + .try_reserve(count) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for index in 0..count { + ranges.push(page.cell_range(index)?); + } + Ok(()) + })?; + self.locator_page_id = page_id; + self.locator_cell_ranges = ranges; + } + self.locator_cell_ranges + .get(usize::from(locator.slot_index)) + .cloned() + .ok_or_else(|| invalid_data("locator slot index is outside leaf page")) + } + + pub(crate) fn read_locator_value(&mut self, locator: Locator, primary_key: &[u8]) -> io::Result { + let page_id = locator.leaf_page_id; + let next_page_id = self.header.next_page_id; + let range = self.locator_cell_range(locator)?; + let allocation_limit = self.value_allocation_limit(); + let decoded = self.with_page(page_id, |bytes, scratch| { + let cell = LeafCellRef::decode( + bytes.get(range).ok_or_else(|| invalid_data("locator cell is outside page"))?, + page_id, + next_page_id, + )?; + if crc32fast::hash(cell.key_bytes) != locator.serialized_key_crc32 + || crc32fast::hash(primary_key) != locator.serialized_key_crc32 + || cell.key_bytes != primary_key + { + return Err(invalid_data("locator serialized key mismatch")); + } + match cell.value { + LeafValueRef::Inline { compression: Compression::None, stored, .. } => { + binary_deserialize(stored).map(DecodedValue::Ready) + } + LeafValueRef::Inline { compression: Compression::Lz4, logical_len, stored, .. } => { + let decompressed = decompress_value_into(stored, logical_len, allocation_limit, scratch)?; + binary_deserialize(decompressed).map(DecodedValue::Ready) + } + LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } => { + Ok(DecodedValue::Overflow(compression, logical_len, stored_len, head, crc32)) + } + LeafValueRef::Tombstone => Ok(DecodedValue::Tombstone), + } + })?; + match decoded { + DecodedValue::Ready(value) => Ok(value), + DecodedValue::Overflow(compression, logical_len, stored_len, head, crc32) => { + self.assemble_overflow_chain(compression, logical_len, stored_len, head, crc32, None)?; + binary_deserialize(&self.value_scratch) + } + DecodedValue::Tombstone => Err(invalid_data("locator references a tombstone")), + } + } + + pub(crate) fn collect_with_locators(&mut self) -> io::Result> { + let mut result = Vec::new(); + let mut page_id = self.header.root_page_id; + let mut visited = HashSet::new(); + let mut descending = true; + let mut descent_depth = 0u64; + loop { + let next_page_id = self.header.next_page_id; + let (child, right, locators) = self.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + if descending && page.header().page_type == PageType::Internal { + let leftmost = InternalPreamble::decode(bytes, page_id, next_page_id)?.leftmost_child; + return Ok((Some(leftmost), 0, Vec::new())); + } + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("locator scan expected a tree page")); + } + let mut locators = Vec::new(); + for index in 0..usize::from(page.header().cell_count) { + let cell = LeafCellRef::decode(page.cell(index)?, page_id, next_page_id)?; + if !matches!(cell.value, LeafValueRef::Tombstone) { + locators.push(( + Locator::for_key( + page_id, + u16::try_from(index).map_err(|_| invalid_data("slot index exceeds u16"))?, + cell.key_bytes, + )?, + page.cell_range(index)?, + )); + } + } + Ok((None, page.header().right, locators)) + })?; + if let Some(child) = child { + if descent_depth >= self.header.next_page_id { + return Err(invalid_data("locator descent contains a cycle")); + } + page_id = child; + descent_depth += 1; + continue; + } + descending = false; + record_page_visit(&mut visited, page_id, "right sibling chain contains a cycle")?; + result + .try_reserve(locators.len()) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for (locator, range) in locators { + let entry = self + .decode_entry_range(page_id, range)? + .ok_or_else(|| invalid_data("live locator became a tombstone"))?; + result.push((entry.0, entry.1, locator)); + } + if right == 0 { + return Ok(result); + } + page_id = right; + } + } + + pub fn contains_live_key(&mut self, key: &K) -> Result { + let result: io::Result = (|| { + let leaf_id = self.locate_leaf(key)?; + let next_page_id = self.header.next_page_id; + let index = self.with_slotted_page(leaf_id, |page, _| { + search_leaf(page, key) + })?; + let Ok(index) = index else { return Ok(false) }; + self.with_slotted_page(leaf_id, |page, _| { + let cell = LeafCellRef::decode(page.cell(index)?, leaf_id, next_page_id)?; + Ok(!matches!(cell.value, LeafValueRef::Tombstone)) + }) + })(); + result.map_err(Into::into) + } + + pub fn query_le(&mut self, key: &K) -> Result, BPlusTreeError> { + self.query_le_io(key).map_err(Into::into) + } + + fn query_le_io(&mut self, key: &K) -> io::Result> { + let mut leaf_id = self.locate_leaf(key)?; + let mut first = true; + let mut visited = HashSet::new(); + record_page_visit(&mut visited, leaf_id, "left sibling chain contains a cycle")?; + loop { + let (left, count, start) = self.with_slotted_page(leaf_id, |page, _| { + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("left sibling is not a leaf")); + } + let count = usize::from(page.header().cell_count); + let start = if first { + match search_leaf(page, key)? { + Ok(index) => Some(index), + Err(0) => None, + Err(index) => index.checked_sub(1), + } + } else { + count.checked_sub(1) + }; + Ok((page.header().left, count, start)) + })?; + if let Some(start) = start { + for index in (0..=start.min(count.saturating_sub(1))).rev() { + if let Some((_, value)) = self.decode_entry(leaf_id, index)? { + return Ok(Some(value)); + } + } + } + if left == 0 { + return Ok(None); + } + let current = leaf_id; + record_page_visit(&mut visited, left, "left sibling chain contains a cycle")?; + leaf_id = left; + self.with_slotted_page(leaf_id, |page, _| { + if page.header().page_type != PageType::Leaf || page.header().right != current { + return Err(invalid_data("asymmetric leaf sibling link")); + } + Ok(()) + })?; + first = false; + } + } + + pub fn len(&mut self) -> Result { + let result: io::Result = (|| { + let mut page_id = self.leftmost_leaf()?; + let mut visited = HashSet::new(); + record_page_visit(&mut visited, page_id, "right sibling chain contains a cycle")?; + let mut total = 0usize; + loop { + let next_page_id = self.header.next_page_id; + let (right, live) = self.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("length scan expected a leaf page")); + } + let mut live = 0usize; + for index in 0..usize::from(page.header().cell_count) { + let cell = LeafCellRef::decode(page.cell(index)?, page_id, next_page_id)?; + if !matches!(cell.value, LeafValueRef::Tombstone) { + live = live.checked_add(1).ok_or_else(|| invalid_data("entry count overflow"))?; + } + } + Ok((page.header().right, live)) + })?; + total = total.checked_add(live).ok_or_else(|| invalid_data("entry count overflow"))?; + if right == 0 { + return Ok(total); + } + record_page_visit(&mut visited, right, "right sibling chain contains a cycle")?; + self.with_page(right, |bytes, _| { + let page = SlottedPage::open(bytes, right, next_page_id)?; + if page.header().page_type != PageType::Leaf || page.header().left != page_id { + return Err(invalid_data("asymmetric leaf sibling link")); + } + Ok(()) + })?; + page_id = right; + } + })(); + result.map_err(Into::into) + } + + pub fn is_empty(&mut self) -> Result { + let mut iterator = self.iter(); + match iterator.next() { + None => Ok(true), + Some(Ok(_)) => Ok(false), + Some(Err(err)) => Err(BPlusTreeError::Io(err)), + } + } + + pub fn iter(&mut self) -> BPlusTreeDiskIterator<'_, K, V> { BPlusTreeDiskIterator::new(self) } + + pub fn disk_iter(self) -> BPlusTreeDiskIteratorOwned { BPlusTreeDiskIteratorOwned::new(self) } + + pub fn range_iter( + &mut self, + start: Bound<&K>, + end: Bound<&K>, + ) -> BPlusTreeRangeIterator<'_, K, V> + where + K: Clone, + { + let start = start.map(Clone::clone); + let end = end.map(Clone::clone); + BPlusTreeRangeIterator { iterator: BPlusTreeDiskIterator::from_bound(self, start.clone()), start, end } + } + + pub fn range_page( + &mut self, + start: Bound<&K>, + end: Bound<&K>, + offset: usize, + limit: usize, + ) -> Result<(Vec<(K, V)>, bool), BPlusTreeError> + where + K: Clone, + { + let mut iterator = self.range_iter(start, end); + for _ in 0..offset { + if let Some(entry) = iterator.next() { + let _ = entry.map_err(BPlusTreeError::Io)?; + } else { + return Ok((Vec::new(), false)); + } + } + let mut result = Vec::new(); + while result.len() < limit { + match iterator.next() { + Some(Ok(entry)) => { + result.try_reserve(1).map_err(|error| { + BPlusTreeError::Io(io::Error::new(io::ErrorKind::OutOfMemory, error)) + })?; + result.push(entry); + } + Some(Err(err)) => return Err(BPlusTreeError::Io(err)), + None => return Ok((result, false)), + } + } + let has_more = match iterator.next() { + Some(Ok(_)) => true, + Some(Err(err)) => return Err(BPlusTreeError::Io(err)), + None => false, + }; + Ok((result, has_more)) + } +} + +struct CursorState { + leaf_page_id: u64, + slot_index: usize, + cell_ranges: Vec>, + right_sibling: u64, + expected_left_sibling: Option, + page_loaded: bool, + start: Bound, + initialized: bool, + finished: bool, + visited_leaves: HashSet, + pending: Option<(K, V)>, +} + +impl CursorState { + fn new(start: Bound) -> Self { + Self { + leaf_page_id: 0, + slot_index: 0, + cell_ranges: Vec::new(), + right_sibling: 0, + expected_left_sibling: None, + page_loaded: false, + start, + initialized: false, + finished: false, + visited_leaves: HashSet::new(), + pending: None, + } + } +} + +fn load_cursor_page(query: &mut BPlusTreeQuery, state: &mut CursorState) -> io::Result<()> { + let page_id = state.leaf_page_id; + let next_page_id = query.header.next_page_id; + let expected_left = state.expected_left_sibling; + let mut cell_ranges = std::mem::take(&mut state.cell_ranges); + let right_sibling = query.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + if page.header().page_type != PageType::Leaf { + return Err(invalid_data("iterator sibling is not a leaf")); + } + if expected_left.is_some_and(|left| page.header().left != left) { + return Err(invalid_data("asymmetric leaf sibling link")); + } + let count = usize::from(page.header().cell_count); + cell_ranges.clear(); + cell_ranges + .try_reserve(count) + .map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + for index in 0..count { + cell_ranges.push(page.cell_range(index)?); + } + Ok(page.header().right) + })?; + state.cell_ranges = cell_ranges; + state.right_sibling = right_sibling; + state.page_loaded = true; + Ok(()) +} + +fn record_page_visit(visited: &mut HashSet, page_id: u64, cycle_error: &'static str) -> io::Result<()> { + visited.try_reserve(1).map_err(|error| io::Error::new(io::ErrorKind::OutOfMemory, error))?; + if !visited.insert(page_id) { + return Err(invalid_data(cycle_error)); + } + Ok(()) +} + +fn cursor_next(query: &mut BPlusTreeQuery, state: &mut CursorState) -> Option> +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + if let Some(entry) = state.pending.take() { + return Some(Ok(entry)); + } + if state.finished { + return None; + } + let mut entry_error = false; + let result = (|| { + if !state.initialized { + state.leaf_page_id = match &state.start { + Bound::Included(key) | Bound::Excluded(key) => query.locate_leaf(key)?, + Bound::Unbounded => query.leftmost_leaf()?, + }; + record_page_visit( + &mut state.visited_leaves, + state.leaf_page_id, + "right sibling chain contains a cycle", + )?; + if let Bound::Included(key) | Bound::Excluded(key) = &state.start { + let page_id = state.leaf_page_id; + let next_page_id = query.header.next_page_id; + state.slot_index = query.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + let index = match search_leaf(&page, key)? { + Ok(index) if matches!(state.start, Bound::Excluded(_)) => index + 1, + Ok(index) | Err(index) => index, + }; + Ok(index) + })?; + } + state.initialized = true; + } + loop { + if !state.page_loaded { + load_cursor_page(query, state)?; + } + let page_id = state.leaf_page_id; + while state.slot_index < state.cell_ranges.len() { + let index = state.slot_index; + state.slot_index += 1; + match query.decode_entry_range(page_id, state.cell_ranges[index].clone()) { + Ok(Some(entry)) => return Ok(Some(entry)), + Ok(None) => {} + Err(error) => { + entry_error = true; + return Err(error); + } + } + } + let right = state.right_sibling; + if right == 0 { + return Ok(None); + } + record_page_visit(&mut state.visited_leaves, right, "right sibling chain contains a cycle")?; + state.expected_left_sibling = Some(page_id); + state.leaf_page_id = right; + state.slot_index = 0; + state.page_loaded = false; + state.cell_ranges.clear(); + } + })(); + match result { + Ok(Some(entry)) => Some(Ok(entry)), + Ok(None) => { + state.finished = true; + None + } + Err(err) => { + state.finished = !entry_error; + Some(Err(err)) + } + } +} + +pub struct BPlusTreeDiskIterator<'a, K, V> { + query: &'a mut BPlusTreeQuery, + state: CursorState, +} + +impl<'a, K, V> BPlusTreeDiskIterator<'a, K, V> { + fn new(query: &'a mut BPlusTreeQuery) -> Self { + Self { query, state: CursorState::new(Bound::Unbounded) } + } + + fn from_bound(query: &'a mut BPlusTreeQuery, start: Bound) -> Self { + Self { query, state: CursorState::new(start) } + } +} + +impl BPlusTreeDiskIterator<'_, K, V> +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + pub fn try_is_empty(&mut self) -> io::Result { + match self.next() { + None => Ok(true), + Some(Ok(entry)) => { + self.state.pending = Some(entry); + Ok(false) + } + Some(Err(err)) => Err(err), + } + } + +} + +impl Iterator for BPlusTreeDiskIterator<'_, K, V> +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + type Item = io::Result<(K, V)>; + + fn next(&mut self) -> Option { cursor_next(self.query, &mut self.state) } +} + +pub struct BPlusTreeDiskIteratorOwned { + query: BPlusTreeQuery, + state: CursorState, +} + +impl BPlusTreeDiskIteratorOwned { + fn new(query: BPlusTreeQuery) -> Self { Self { query, state: CursorState::new(Bound::Unbounded) } } +} + +impl BPlusTreeDiskIteratorOwned +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + pub fn try_is_empty(&mut self) -> io::Result { + match self.next() { + None => Ok(true), + Some(Ok(entry)) => { + self.state.pending = Some(entry); + Ok(false) + } + Some(Err(err)) => Err(err), + } + } + +} + +impl Iterator for BPlusTreeDiskIteratorOwned +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + type Item = io::Result<(K, V)>; + + fn next(&mut self) -> Option { cursor_next(&mut self.query, &mut self.state) } +} + +pub struct BPlusTreeRangeIterator<'a, K, V> { + iterator: BPlusTreeDiskIterator<'a, K, V>, + start: Bound, + end: Bound, +} + +fn within_start(key: &K, bound: &Bound) -> bool { + match bound { + Bound::Included(start) => key >= start, + Bound::Excluded(start) => key > start, + Bound::Unbounded => true, + } +} + +fn past_end(key: &K, bound: &Bound) -> bool { + match bound { + Bound::Included(end) => key > end, + Bound::Excluded(end) => key >= end, + Bound::Unbounded => false, + } +} + +impl Iterator for BPlusTreeRangeIterator<'_, K, V> +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, +{ + type Item = io::Result<(K, V)>; + + fn next(&mut self) -> Option { + loop { + let entry = self.iterator.next()?; + match entry { + Ok((key, _value)) if past_end(&key, &self.end) => { + self.iterator.state.finished = true; + return None; + } + Ok((key, value)) if within_start(&key, &self.start) => return Some(Ok((key, value))), + Ok(_) => {} + Err(err) => return Some(Err(err)), + } + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct VerificationReport { + pub(crate) live_entries: u64, + pub(crate) tree_pages: u64, + pub(crate) overflow_pages: u64, + pub(crate) free_pages: u64, +} + +struct VerifyLeaf { + page_id: u64, + left: u64, + right: u64, + minimum: Option, + maximum: Option, +} + +enum Visit { + Enter(u64, Option, Option), + Exit(u64), +} + +enum VerifyValue { + Inline, + Overflow(Compression, u32, u32, u64, u32), + Tombstone, +} + +fn finish_verified_page( + page_id: u64, + active: &mut HashSet, + internal_children: &mut HashMap)>>, + page_minimum: &mut HashMap>, +) -> io::Result<()> { + if !active.remove(&page_id) { + return Err(invalid_data("tree verifier active set is inconsistent")); + } + let Some(children) = internal_children.remove(&page_id) else { return Ok(()) }; + for (child, separator) in children.iter().skip(1) { + let actual = page_minimum + .get(child) + .and_then(Option::as_ref) + .ok_or_else(|| invalid_data("internal child has no minimum key"))?; + if separator.as_ref() != Some(actual) { + return Err(invalid_data("internal separator is not the right child minimum")); + } + } + let minimum = children + .first() + .and_then(|(child, _)| page_minimum.get(child)) + .cloned() + .ok_or_else(|| invalid_data("internal leftmost child has no verified minimum"))?; + page_minimum.insert(page_id, minimum); + Ok(()) +} + +fn verify_leaf_page( + query: &mut BPlusTreeQuery, + page_id: u64, + lower: Option<&K>, + upper: Option<&K>, + overflow_pages: &mut HashSet, +) -> io::Result<(VerifyLeaf, u64)> +where + K: Ord + for<'de> Deserialize<'de> + Clone, +{ + let next_page_id = query.header.next_page_id; + let (left, right, count) = query.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + Ok((page.header().left, page.header().right, usize::from(page.header().cell_count))) + })?; + let mut minimum = None; + let mut maximum = None; + let mut live_entries = 0u64; + for index in 0..count { + let (key, value) = query.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + let cell = LeafCellRef::decode(page.cell(index)?, page_id, next_page_id)?; + let value = match cell.value { + LeafValueRef::Inline { .. } => VerifyValue::Inline, + LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } => { + VerifyValue::Overflow(compression, logical_len, stored_len, head, crc32) + } + LeafValueRef::Tombstone => VerifyValue::Tombstone, + }; + Ok((binary_deserialize::(cell.key_bytes)?, value)) + })?; + if maximum.as_ref().is_some_and(|previous| previous >= &key) { + return Err(invalid_data("leaf keys are not strictly ordered")); + } + if lower.is_some_and(|bound| &key < bound) || upper.is_some_and(|bound| &key >= bound) { + return Err(invalid_data("leaf key is outside parent separator range")); + } + minimum.get_or_insert_with(|| key.clone()); + maximum = Some(key); + match value { + VerifyValue::Tombstone => continue, + VerifyValue::Inline => {} + VerifyValue::Overflow(compression, logical_len, stored_len, head, crc32) => query + .assemble_overflow_chain( + compression, + logical_len, + stored_len, + head, + crc32, + Some(overflow_pages), + )?, + } + live_entries = live_entries + .checked_add(1) + .ok_or_else(|| invalid_data("live entry count overflow"))?; + } + Ok((VerifyLeaf { page_id, left, right, minimum, maximum }, live_entries)) +} + +fn verify_internal_page(query: &mut BPlusTreeQuery, page_id: u64) -> io::Result)>> +where + K: Ord + for<'de> Deserialize<'de> + Clone, +{ + let next_page_id = query.header.next_page_id; + query.with_page(page_id, |bytes, _| { + let page = SlottedPage::open(bytes, page_id, next_page_id)?; + let mut children = Vec::with_capacity(usize::from(page.header().cell_count) + 1); + children.push((InternalPreamble::decode(bytes, page_id, next_page_id)?.leftmost_child, None)); + let mut previous = None; + for cell in page.cells() { + let cell = InternalCellRef::decode(cell?, page_id, next_page_id)?; + let key = binary_deserialize::(cell.key_bytes)?; + if previous.as_ref().is_some_and(|prior| prior >= &key) { + return Err(invalid_data("internal separators are not strictly ordered")); + } + previous = Some(key.clone()); + children.push((cell.right_child, Some(key))); + } + Ok(children) + }) +} + +fn push_child_visits( + stack: &mut Vec>, + children: &[(u64, Option)], + lower: Option<&K>, + upper: Option<&K>, +) -> io::Result<()> { + for index in (0..children.len()).rev() { + let child_lower = if index == 0 { lower.cloned() } else { children.get(index).and_then(|(_, key)| key.clone()) }; + let child_upper = children.get(index + 1).and_then(|(_, key)| key.clone()).or_else(|| upper.cloned()); + let child = children + .get(index) + .map(|(child, _)| *child) + .ok_or_else(|| invalid_data("internal child index is invalid"))?; + stack.push(Visit::Enter(child, child_lower, child_upper)); + } + Ok(()) +} + +fn verify_leaf_links(leaves: &[VerifyLeaf]) -> io::Result<()> { + for (index, leaf) in leaves.iter().enumerate() { + let expected_left = index.checked_sub(1).map_or(0, |previous| leaves[previous].page_id); + let expected_right = leaves.get(index + 1).map_or(0, |next| next.page_id); + if leaf.left != expected_left || leaf.right != expected_right { + return Err(invalid_data("leaf sibling links do not match tree order")); + } + if index > 0 + && leaves[index - 1] + .maximum + .as_ref() + .zip(leaf.minimum.as_ref()) + .is_some_and(|(previous, current)| previous >= current) + { + return Err(invalid_data("keys are inverted across leaf siblings")); + } + } + Ok(()) +} + +fn verify_free_pages( + query: &mut BPlusTreeQuery, + tree_pages: &HashSet, + overflow_pages: &HashSet, +) -> io::Result> { + let mut free_pages = HashSet::new(); + let mut free = query.header.free_page_head; + while free != 0 { + record_page_visit(&mut free_pages, free, "free list contains a duplicate or cycle")?; + if tree_pages.contains(&free) || overflow_pages.contains(&free) { + return Err(invalid_data("page is reachable from both live data and free list")); + } + let next_page_id = query.header.next_page_id; + free = query.with_page(free, |bytes, _| { + let page = SlottedPage::open(bytes, free, next_page_id)?; + if page.header().page_type != PageType::Free { + return Err(invalid_data("free list references a non-free page")); + } + Ok(page.header().right) + })?; + } + Ok(free_pages) +} + +fn verify_page_ownership( + next_page_id: u64, + tree_pages: &HashSet, + overflow_pages: &HashSet, + free_pages: &HashSet, +) -> io::Result<()> { + for page_id in 1..next_page_id { + let memberships = u8::from(tree_pages.contains(&page_id)) + + u8::from(overflow_pages.contains(&page_id)) + + u8::from(free_pages.contains(&page_id)); + if memberships != 1 { + return Err(invalid_data("database contains an orphan or multiply-owned page")); + } + } + Ok(()) +} + +pub(crate) fn verify_full(query: &mut BPlusTreeQuery) -> io::Result +where + K: Ord + for<'de> Deserialize<'de> + Clone, +{ + let mut tree_pages = HashSet::new(); + let mut active = HashSet::new(); + let mut overflow_pages = HashSet::new(); + let mut leaves = Vec::>::new(); + let mut page_minimum = HashMap::>::new(); + let mut internal_children = HashMap::)>>::new(); + let mut live_entries = 0u64; + let mut stack = vec![Visit::Enter(query.header.root_page_id, None, None)]; + + while let Some(visit) = stack.pop() { + let (page_id, lower, upper) = match visit { + Visit::Exit(page_id) => { + finish_verified_page(page_id, &mut active, &mut internal_children, &mut page_minimum)?; + continue; + } + Visit::Enter(page_id, lower, upper) => (page_id, lower, upper), + }; + if active.contains(&page_id) { + return Err(invalid_data("tree child graph contains a cycle")); + } + if !tree_pages.insert(page_id) { + return Err(invalid_data("tree page has multiple parents")); + } + active.insert(page_id); + stack.push(Visit::Exit(page_id)); + let next_page_id = query.header.next_page_id; + let page_type = query.with_page(page_id, |bytes, _| { + SlottedPage::open(bytes, page_id, next_page_id).map(|page| page.header().page_type) + })?; + match page_type { + PageType::Leaf => { + let (leaf, page_live_entries) = + verify_leaf_page(query, page_id, lower.as_ref(), upper.as_ref(), &mut overflow_pages)?; + live_entries = live_entries + .checked_add(page_live_entries) + .ok_or_else(|| invalid_data("live entry count overflow"))?; + page_minimum.insert(page_id, leaf.minimum.clone()); + leaves.push(leaf); + } + PageType::Internal => { + let children = verify_internal_page(query, page_id)?; + push_child_visits(&mut stack, &children, lower.as_ref(), upper.as_ref())?; + internal_children.insert(page_id, children); + } + PageType::Overflow | PageType::Free => return Err(invalid_data("tree child has the wrong page type")), + } + } + + verify_leaf_links(&leaves)?; + let free_pages = verify_free_pages(query, &tree_pages, &overflow_pages)?; + verify_page_ownership(query.header.next_page_id, &tree_pages, &overflow_pages, &free_pages)?; + Ok(VerificationReport { + live_entries, + tree_pages: u64::try_from(tree_pages.len()).map_err(|_| invalid_data("tree page count exceeds u64"))?, + overflow_pages: u64::try_from(overflow_pages.len()).map_err(|_| invalid_data("overflow page count exceeds u64"))?, + free_pages: u64::try_from(free_pages.len()).map_err(|_| invalid_data("free page count exceeds u64"))?, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + repository::bplustree::v3::{ + format::{ + encode_inline_leaf_cell, encode_internal_cell, encode_overflow_leaf_cell, + encode_tombstone_leaf_cell, encode_value, Compression, PageHeader, PageType, + OVERFLOW_PAYLOAD_LEN, PAGE_HEADER_LEN, PAGE_SIZE, + }, + page::{ + encode_free_page, encode_overflow_page, page_open_count, reset_page_open_count, SlottedPage, + }, + }, + utils::binary_serialize, + }; + use std::{ + fs, io, + path::{Path, PathBuf}, + process::Command, + sync::mpsc::{self, Receiver, RecvTimeoutError}, + thread::{self, JoinHandle}, + time::Duration, + }; + use fs2::FileExt as _; + + const PAGE_ID: u64 = 7; + const NEXT_PAGE_ID: u64 = 20; + + fn invalid_data(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidData => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidData, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidData")), + } + } + + fn invalid_input(result: io::Result) -> io::Result<()> { + match result { + Err(err) if err.kind() == io::ErrorKind::InvalidInput => Ok(()), + Err(err) => Err(io::Error::other(format!("expected InvalidInput, got {err}"))), + Ok(_) => Err(io::Error::other("expected InvalidInput")), + } + } + + fn database_header(path: &Path) -> io::Result { + let bytes = fs::read(path)?; + DatabaseHeader::decode( + bytes + .get(..PAGE_SIZE) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "database header is truncated"))?, + ) + } + + #[test] + fn point_query_validates_a_single_leaf_once() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("point-query-page-opens.db"); + let mut tree = BPlusTree::new(); + tree.insert(7u32, String::from("value")); + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + reset_page_open_count(); + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(String::from("value"))); + assert_eq!(page_open_count(), 1); + Ok(()) + } + + #[test] + fn repeated_point_query_reuses_snapshot_page_validation_and_internal_keys() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("repeated-point-query.db"); + let mut tree = BPlusTree::new(); + for key in 0..2_000u32 { + tree.insert(key, format!("value-{key:04}")); + } + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + reset_page_open_count(); + reset_internal_key_decode_count(); + assert_eq!(query.query(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + let first_page_opens = page_open_count(); + let first_internal_decodes = internal_key_decode_count(); + if first_page_opens < 2 || first_internal_decodes == 0 { + return Err(io::Error::other("test fixture must contain internal pages")); + } + + assert_eq!(query.query(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + assert_eq!(page_open_count(), first_page_opens); + assert_eq!(internal_key_decode_count(), first_internal_decodes); + Ok(()) + } + + #[test] + fn cloned_query_reuses_snapshot_page_validation_and_internal_keys() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("cloned-point-query.db"); + let mut tree = BPlusTree::new(); + for key in 0..2_000u32 { + tree.insert(key, format!("value-{key:04}")); + } + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + reset_page_open_count(); + reset_internal_key_decode_count(); + assert_eq!(query.query(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + let first_page_opens = page_open_count(); + let first_internal_decodes = internal_key_decode_count(); + if first_page_opens < 2 || first_internal_decodes == 0 { + return Err(io::Error::other("test fixture must contain internal pages")); + } + + let mut cloned = query.try_clone()?; + assert_eq!(cloned.query(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + assert_eq!(page_open_count(), first_page_opens); + assert_eq!(internal_key_decode_count(), first_internal_decodes); + Ok(()) + } + + #[test] + fn repeated_query_le_reuses_snapshot_page_validation_and_internal_keys() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("repeated-query-le.db"); + let mut tree = BPlusTree::new(); + for key in 0..2_000u32 { + tree.insert(key, format!("value-{key:04}")); + } + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + reset_page_open_count(); + reset_internal_key_decode_count(); + assert_eq!(query.query_le(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + let first_page_opens = page_open_count(); + let first_internal_decodes = internal_key_decode_count(); + if first_page_opens < 2 || first_internal_decodes == 0 { + return Err(io::Error::other("test fixture must contain internal pages")); + } + + assert_eq!(query.query_le(&1_337).map_err(BPlusTreeError::to_io)?, Some(String::from("value-1337"))); + assert_eq!(page_open_count(), first_page_opens); + assert_eq!(internal_key_decode_count(), first_internal_decodes); + Ok(()) + } + + #[test] + fn locator_collection_validates_each_leaf_once() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("locator-page-opens.db"); + let mut tree = BPlusTree::new(); + for key in 0..3u32 { + tree.insert(key, format!("value-{key}")); + } + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + reset_page_open_count(); + assert_eq!(query.collect_with_locators()?.len(), 3); + assert_eq!(page_open_count(), 1); + Ok(()) + } + + #[test] + fn prepared_upsert_batch_is_key_sorted_and_stable() -> io::Result<()> { + let keys = [3u32, 1, 2, 1]; + let values = ["three", "first", "two", "last"].map(String::from); + let items = keys.iter().zip(&values).collect::>(); + + let prepared = BPlusTreeUpdate::::prepare_upsert_batch(&items)?; + let decoded = prepared + .into_iter() + .map(|(key, value)| binary_deserialize::(&value).map(|value| (key, value))) + .collect::>>()?; + assert_eq!( + decoded, + vec![ + (1, String::from("first")), + (1, String::from("last")), + (2, String::from("two")), + (3, String::from("three")), + ] + ); + Ok(()) + } + + #[test] + fn equal_size_inline_update_rewrites_without_appending_pages() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("equal-inline.db"); + let mut tree = BPlusTree::new(); + tree.insert(7u32, String::from("old")); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.update(&7, String::from("new")).map_err(BPlusTreeError::to_io)?; + drop(updater); + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(String::from("new"))); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn smaller_inline_update_compacts_without_appending_pages() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("smaller-inline.db"); + let mut tree = BPlusTree::new(); + tree.insert(7u32, String::from("a much longer inline value")); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.update(&7, String::from("x")).map_err(BPlusTreeError::to_io)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(String::from("x"))); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn growing_inline_value_uses_page_local_compaction_before_split() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("growing-inline.db"); + let mut tree = BPlusTree::new(); + for key in 0..8u32 { + tree.insert(key, vec![u8::try_from(key).map_err(io::Error::other)?; 8]); + } + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + let grown = vec![0x5a; 200]; + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.update(&4, grown.clone()).map_err(BPlusTreeError::to_io)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(query.query(&4).map_err(BPlusTreeError::to_io)?, Some(grown)); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn medium_stored_value_remains_inline() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("medium-inline.db"); + let value = random_value() + .get(..300) + .ok_or_else(|| io::Error::other("random test value is too short"))? + .to_vec(); + let mut tree = BPlusTree::new(); + tree.insert(7u32, value.clone()); + + let report = tree.store_verified(&path)?; + + assert_eq!(report.overflow_pages, 0); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(value)); + Ok(()) + } + + fn overflow_head(path: &Path, key: u32) -> io::Result { + let mut query = BPlusTreeQuery::>::try_new(path)?; + let leaf = query.locate_leaf(&key)?; + let next_page_id = query.header.next_page_id; + query.with_page(leaf, |bytes, _| { + let page = SlottedPage::open(bytes, leaf, next_page_id)?; + let index = search_leaf(&page, &key)?.map_err(|_| io::Error::other("test key is missing"))?; + match LeafCellRef::decode(page.cell(index)?, leaf, next_page_id)?.value { + LeafValueRef::Overflow { head, .. } => Ok(head), + LeafValueRef::Inline { .. } | LeafValueRef::Tombstone => { + Err(io::Error::other("test value is not overflow-backed")) + } + } + }) + } + + #[test] + fn overflow_update_reuses_head_and_frees_unused_tail() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("overflow-reuse.db"); + let large = random_value(); + let smaller = large + .get(..5_000) + .ok_or_else(|| io::Error::other("random test value is too short"))? + .to_vec(); + let mut tree = BPlusTree::new(); + tree.insert(7u32, large); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + let before_head = overflow_head(&path, 7)?; + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.update(&7, smaller.clone()).map_err(BPlusTreeError::to_io)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(fs::metadata(&path)?.len(), before_len); + assert_eq!(overflow_head(&path, 7)?, before_head); + assert_ne!(after.free_page_head, 0); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(smaller)); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn inline_to_overflow_update_allocates_new_pages() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("overflow-allocation.db"); + let mut tree = BPlusTree::new(); + tree.insert(7u32, vec![7]); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + let large = random_value(); + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.update(&7, large.clone()).map_err(BPlusTreeError::to_io)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert!(fs::metadata(&path)?.len() > before_len); + assert_ne!(overflow_head(&path, 7)?, 0); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(large)); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn new_key_is_inserted_into_existing_leaf_without_growth() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("leaf-insert.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(3u32, String::from("three")); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.upsert(&2, &String::from("two"))?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.iter().collect::>>()?, vec![(1, "one".into()), (2, "two".into()), (3, "three".into())]); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn leaf_split_repairs_siblings_and_creates_a_new_root() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("leaf-split.db"); + let value: String = random_value() + .get(..200) + .ok_or_else(|| io::Error::other("random test value is too short"))? + .iter() + .map(|byte| char::from(33 + byte % 90)) + .collect(); + let mut tree = BPlusTree::new(); + for key in 0..17u32 { + tree.insert(key, value.clone()); + } + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.upsert(&17, &value)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_ne!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len + 2 * u64::try_from(PAGE_SIZE).map_err(io::Error::other)?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + let report = verify_full(&mut query)?; + assert_eq!(report.live_entries, 18); + assert_eq!(query.query(&17).map_err(BPlusTreeError::to_io)?, Some(value)); + Ok(()) + } + + #[test] + fn splitting_a_non_rightmost_leaf_repairs_the_former_neighbor() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("middle-leaf-split.db"); + let value = random_value() + .get(..200) + .ok_or_else(|| io::Error::other("random test value is too short"))? + .to_vec(); + let mut tree = BPlusTree::new(); + for key in (0..800u32).step_by(2) { + tree.insert(key, value.clone()); + } + tree.store(&path)?; + let before = database_header(&path)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let left_leaf = query.locate_leaf(&0)?; + let former_right = query.with_page(left_leaf, |bytes, _| { + Ok(SlottedPage::open(bytes, left_leaf, before.next_page_id)?.header().right) + })?; + if former_right == 0 { + return Err(io::Error::other("fixture did not create a right leaf neighbor")); + } + drop(query); + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.upsert(&1, &value)?; + + let after = database_header(&path)?; + assert_eq!(after.root_page_id, before.root_page_id); + assert!(after.next_page_id > before.next_page_id); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let new_right = query.with_page(left_leaf, |bytes, _| { + Ok(SlottedPage::open(bytes, left_leaf, after.next_page_id)?.header().right) + })?; + assert_ne!(new_right, former_right); + query.with_page(new_right, |bytes, _| { + let page = SlottedPage::open(bytes, new_right, after.next_page_id)?; + assert_eq!(page.header().left, left_leaf); + assert_eq!(page.header().right, former_right); + Ok(()) + })?; + query.with_page(former_right, |bytes, _| { + let page = SlottedPage::open(bytes, former_right, after.next_page_id)?; + assert_eq!(page.header().left, new_right); + Ok(()) + })?; + let _ = verify_full(&mut query)?; + Ok(()) + } + + fn long_split_key(index: u32, marker: char) -> String { + format!("{index:04}{marker}{}", "k".repeat(1_880)) + } + + #[test] + fn leaf_promotion_recursively_splits_a_full_internal_root() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("internal-split.db"); + let value = String::from("value"); + let mut tree = BPlusTree::new(); + for index in 0..6u32 { + tree.insert(long_split_key(index, '-'), value.clone()); + } + tree.store(&path)?; + let before = database_header(&path)?; + let before_root = fs::read(&path)?; + let range = page_byte_range(before.root_page_id, before_root.len())?; + let root_page = SlottedPage::open( + before_root.get(range).ok_or_else(|| io::Error::other("root page is missing"))?, + before.root_page_id, + before.next_page_id, + )?; + assert_eq!(root_page.header().page_type, PageType::Internal); + assert_eq!(root_page.header().cell_count, 2); + + let inserted_key = long_split_key(1, 'z'); + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.upsert(&inserted_key, &value)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_ne!(after.root_page_id, before.root_page_id); + let mut query = BPlusTreeQuery::::try_new(&path)?; + let report = verify_full(&mut query)?; + assert_eq!(report.live_entries, 7); + assert_eq!(query.query(&inserted_key).map_err(BPlusTreeError::to_io)?, Some(value)); + Ok(()) + } + + #[test] + fn delete_writes_tombstone_and_reinsert_reuses_the_leaf() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("tombstone.db"); + let mut tree = BPlusTree::new(); + tree.insert(7u32, String::from("original")); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + assert!(updater.delete(&7)?); + assert_eq!(updater.query(&7).map_err(BPlusTreeError::to_io)?, None); + updater.upsert(&7, &String::from("restored"))?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 2); + assert_eq!(after.root_page_id, before.root_page_id); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&7).map_err(BPlusTreeError::to_io)?, Some(String::from("restored"))); + assert_eq!(verify_full(&mut query)?.live_entries, 1); + Ok(()) + } + + #[test] + fn freed_overflow_pages_are_reused_before_file_growth() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("free-reuse.db"); + let large = random_value(); + let mut tree = BPlusTree::new(); + tree.insert(1u32, large.clone()); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + assert!(updater.delete(&1)?); + assert_ne!(database_header(&path)?.free_page_head, 0); + updater.upsert(&2, &large)?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 2); + assert_eq!(fs::metadata(&path)?.len(), before_len); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, None); + assert_eq!(query.query(&2).map_err(BPlusTreeError::to_io)?, Some(large)); + assert_eq!(verify_full(&mut query)?.live_entries, 1); + Ok(()) + } + + #[test] + fn metadata_update_is_a_single_header_transaction() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("metadata.db"); + let mut tree = BPlusTree::::new(); + tree.insert(1, String::from("one")); + tree.store(&path)?; + let before = database_header(&path)?; + let before_len = fs::metadata(&path)?.len(); + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_metadata(&BPlusTreeMetadata::TargetIdMapping(42))?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + assert_eq!(after.metadata, BPlusTreeMetadata::TargetIdMapping(42)); + assert_eq!(fs::metadata(&path)?.len(), before_len); + assert_eq!(updater.get_metadata()?, BPlusTreeMetadata::TargetIdMapping(42)); + Ok(()) + } + + #[test] + fn identical_batch_metadata_is_a_true_noop() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("metadata-noop.db"); + let mut tree = BPlusTree::::new(); + tree.set_metadata(BPlusTreeMetadata::TargetIdMapping(42)); + tree.store(&path)?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.set_metadata(&BPlusTreeMetadata::TargetIdMapping(42))?; + assert!(updater.active.is_none()); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn immediate_commit_clears_wal_then_invalidates_sorted_index() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("immediate.db"); + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("old")); + tree.store(&path)?; + fs::write(&index_path, b"derived")?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.upsert(&1, &String::from("new"))?; + + assert!(!wal_path(&path).try_exists()?); + assert!(!wal_temporary_path(&path).try_exists()?); + assert!(!index_path.try_exists()?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(String::from("new"))); + Ok(()) + } + + #[test] + fn database_with_index_extension_is_never_deleted_as_derived_data() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("tree.idx"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("old")); + tree.store(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.upsert(&1, &String::from("new"))?; + + assert!(path.try_exists()?); + assert!(!wal_path(&path).try_exists()?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(String::from("new"))); + let _ = verify_full(&mut query)?; + Ok(()) + } + + #[test] + fn missing_delete_and_empty_commit_are_true_noops() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("no-op.db"); + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.store(&path)?; + fs::write(&index_path, b"still-valid")?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + assert!(!updater.delete(&2)?); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(index_path.try_exists()?); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn serial_writer_commits_a_batch_and_shuts_down_cleanly() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("serial-writer.db"); + BPlusTree::::new().store(&path)?; + + let writer = BPlusTreeSerialWriter::new(&path, FlushPolicy::Batch)?; + let one = String::from("one"); + let two = String::from("two"); + assert_ne!(writer.upsert(&[(&1, &one), (&2, &two)])?, 0); + writer.shutdown()?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.iter().collect::>>()?, vec![(1, one), (2, two)]); + Ok(()) + } + + #[test] + fn store_with_index_publishes_an_identity_bound_sorted_snapshot() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("indexed.db"); + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("ccc")); + tree.insert(2u32, String::from("a")); + tree.insert(3u32, String::from("bb")); + + assert_ne!(tree.store_with_index(&path, String::len)?, 0); + + let query = BPlusTreeQuery::::try_new(&path)?; + let mut sorted = crate::repository::bplustree::sorted_index::v4::OwnedIterator::::open( + query, + &index_path, + )?; + assert_eq!( + sorted.by_ref().collect::>>()?, + vec![(2, String::from("a")), (3, String::from("bb")), (1, String::from("ccc"))] + ); + assert_eq!(sorted.remaining(), 0); + assert!(!fs::read_dir(dir.path())? + .any(|entry| entry.is_ok_and(|entry| entry.file_name().to_string_lossy().ends_with(".v3.tmp")))); + Ok(()) + } + + #[test] + fn compact_rebuilds_only_live_data_with_a_fresh_identity() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("compact.db"); + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + let mut tree = BPlusTree::new(); + for key in 0..100u32 { + tree.insert( + key, + if key % 10 == 0 { random_value() } else { vec![u8::try_from(key).map_err(io::Error::other)?; 32] }, + ); + } + tree.store(&path)?; + let original_id = database_header(&path)?.database_id; + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + for key in (0..100u32).step_by(2) { + assert!(updater.delete(&key)?); + } + updater.commit()?; + let length_before = fs::metadata(&path)?.len(); + fs::write(&index_path, b"stale")?; + + updater.compact()?; + + let header = database_header(&path)?; + assert_ne!(header.database_id, original_id); + assert_eq!(header.generation, 1); + assert!(fs::metadata(&path)?.len() <= length_before); + assert!(!index_path.try_exists()?); + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let entries = query.iter().collect::>>()?; + assert_eq!(entries.len(), 50); + assert!(entries.iter().all(|(key, _)| key % 2 == 1)); + let report = verify_full(&mut query)?; + assert_eq!(report.live_entries, 50); + assert_eq!(report.free_pages, 0); + Ok(()) + } + + #[test] + fn compact_read_failure_preserves_database_and_index() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("compact-corrupt.db"); + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.store(&path)?; + let updater = BPlusTreeUpdate::::try_new(&path)?; + + let header = database_header(&path)?; + let mut corrupted = fs::read(&path)?; + let offset = usize::try_from(header.root_page_id) + .map_err(io::Error::other)? + .checked_mul(PAGE_SIZE) + .and_then(|start| start.checked_add(PAGE_HEADER_LEN)) + .ok_or_else(|| io::Error::other("corruption offset overflow"))?; + *corrupted.get_mut(offset).ok_or_else(|| io::Error::other("corruption offset outside database"))? ^= 0xff; + fs::write(&path, &corrupted)?; + fs::write(&index_path, b"still-valid")?; + + let mut updater = updater; + assert!(updater.compact().is_err()); + assert_eq!(fs::read(&path)?, corrupted); + assert_eq!(fs::read(&index_path)?, b"still-valid"); + Ok(()) + } + + #[test] + fn batch_overlay_is_visible_to_updater_and_commits_once() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("batch.db"); + let value: String = random_value() + .get(..200) + .ok_or_else(|| io::Error::other("random test value is too short"))? + .iter() + .map(|byte| char::from(33 + byte % 90)) + .collect(); + let mut tree = BPlusTree::new(); + for key in 0..17u32 { + tree.insert(key, value.clone()); + } + tree.store(&path)?; + let before = database_header(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&17, &value)?; + updater.upsert(&18, &value)?; + assert_eq!(updater.query(&18).map_err(BPlusTreeError::to_io)?, Some(value.clone())); + assert_eq!(database_header(&path)?.generation, before.generation); + + updater.commit()?; + + let after = database_header(&path)?; + assert_eq!(after.generation, before.generation + 1); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&17).map_err(BPlusTreeError::to_io)?, Some(value.clone())); + assert_eq!(query.query(&18).map_err(BPlusTreeError::to_io)?, Some(value)); + assert_eq!(verify_full(&mut query)?.live_entries, 19); + Ok(()) + } + + #[test] + fn dropped_uncommitted_batch_discards_overlay_without_wal() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("dropped-batch.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("old")); + tree.store(&path)?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&1, &String::from("uncommitted"))?; + assert_eq!(updater.query(&1).map_err(BPlusTreeError::to_io)?, Some(String::from("uncommitted"))); + drop(updater); + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(String::from("old"))); + Ok(()) + } + + #[test] + fn failed_batch_mutation_discards_the_whole_overlay() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("poisoned-batch.db"); + let mut tree = BPlusTree::new(); + tree.insert(String::from("key"), String::from("old")); + tree.store(&path)?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&String::from("key"), &String::from("staged"))?; + assert!(updater.upsert(&"x".repeat(2_100), &String::from("invalid")).is_err()); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&String::from("key")).map_err(BPlusTreeError::to_io)?, Some(String::from("old"))); + Ok(()) + } + + #[test] + fn final_overlay_validation_rejects_unordered_leaf_keys_before_wal() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("unordered-overlay.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.store(&path)?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&1, &String::from("staged"))?; + let active = updater.active.as_mut().ok_or_else(|| io::Error::other("test transaction is missing"))?; + let leaf_id = active.transaction.next_header.root_page_id; + let snapshot = active.transaction.page_copy(active.base.as_slice(), leaf_id)?; + let page = SlottedPage::open( + snapshot.as_slice(), + leaf_id, + active.transaction.next_header.next_page_id, + )?; + let mut cells = page.cells().map(|cell| cell.map(<[u8]>::to_vec)).collect::>>()?; + let duplicate = cells + .first() + .cloned() + .ok_or_else(|| io::Error::other("test leaf has no cells"))?; + *cells.get_mut(1).ok_or_else(|| io::Error::other("test leaf lacks a second cell"))? = duplicate; + let next_page_id = active.transaction.next_header.next_page_id; + let dirty = active.transaction.page_mut(active.base.as_slice(), leaf_id)?; + SlottedPage::open(dirty.as_mut_slice(), leaf_id, next_page_id)? + .rebuild_ordered(cells.iter().map(Vec::as_slice))?; + + let error = updater.commit().err().ok_or_else(|| io::Error::other("unordered overlay was committed"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn final_overlay_validation_rejects_shared_overflow_before_wal() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("shared-overflow-overlay.db"); + let first_value = random_value(); + let mut second_value = first_value.clone(); + second_value.reverse(); + let mut tree = BPlusTree::new(); + tree.insert(1u32, first_value); + tree.insert(2u32, second_value); + tree.store(&path)?; + let before = fs::read(&path)?; + + let mut updater = BPlusTreeUpdate::>::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.ensure_transaction()?; + let active = updater.active.as_mut().ok_or_else(|| io::Error::other("test transaction is missing"))?; + let leaf_id = active.transaction.next_header.root_page_id; + let next_page_id = active.transaction.next_header.next_page_id; + let snapshot = active.transaction.page_copy(active.base.as_slice(), leaf_id)?; + let page = SlottedPage::open(snapshot.as_slice(), leaf_id, next_page_id)?; + let first = LeafCellRef::decode(page.cell(0)?, leaf_id, next_page_id)?; + let second = LeafCellRef::decode(page.cell(1)?, leaf_id, next_page_id)?; + let LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } = first.value else { + return Err(io::Error::other("first test value is not overflow-backed")); + }; + let mut replacement = Vec::new(); + encode_overflow_leaf_cell( + second.key_bytes, + logical_len, + compression, + stored_len, + head, + crc32, + leaf_id, + next_page_id, + &mut replacement, + )?; + let dirty = active.transaction.page_mut(active.base.as_slice(), leaf_id)?; + SlottedPage::open(dirty.as_mut_slice(), leaf_id, next_page_id)?.replace_same_len(1, &replacement)?; + + let error = updater.commit().err().ok_or_else(|| io::Error::other("shared overflow was committed"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)] + struct ConditionalSerialize { + value: String, + fail: bool, + } + + impl Serialize for ConditionalSerialize { + fn serialize(&self, serializer: S) -> Result { + if self.fail { + return Err(serde::ser::Error::custom("injected serialization failure")); + } + self.value.serialize(serializer) + } + } + + impl<'de> Deserialize<'de> for ConditionalSerialize { + fn deserialize>(deserializer: D) -> Result { + String::deserialize(deserializer).map(|value| Self { value, fail: false }) + } + } + + #[test] + fn batch_serialization_error_discards_earlier_items() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("batch-serialization.db"); + let initial = ConditionalSerialize { value: String::from("initial"), fail: false }; + let mut tree = BPlusTree::new(); + tree.insert(0u32, initial); + tree.store(&path)?; + let before = fs::read(&path)?; + let first = ConditionalSerialize { value: String::from("first"), fail: false }; + let failing = ConditionalSerialize { value: String::from("second"), fail: true }; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + assert!(updater.upsert_batch(&[(&1, &first), (&2, &failing)]).is_err()); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, None); + assert_eq!(query.query(&2).map_err(BPlusTreeError::to_io)?, None); + Ok(()) + } + + #[test] + fn direct_serialization_error_aborts_an_existing_batch() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("direct-serialization.db"); + let initial = ConditionalSerialize { value: String::from("initial"), fail: false }; + let mut tree = BPlusTree::new(); + tree.insert(0u32, initial); + tree.store(&path)?; + let before = fs::read(&path)?; + let staged = ConditionalSerialize { value: String::from("staged"), fail: false }; + let failing = ConditionalSerialize { value: String::from("failing"), fail: true }; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&1, &staged)?; + assert!(updater.upsert(&2, &failing).is_err()); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn delete_key_serialization_error_aborts_an_existing_batch() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("delete-serialization.db"); + let initial = ConditionalSerialize { value: String::from("initial"), fail: false }; + let mut tree = BPlusTree::new(); + tree.insert(initial.clone(), String::from("old")); + tree.store(&path)?; + let before = fs::read(&path)?; + let staged = ConditionalSerialize { value: String::from("staged"), fail: false }; + let failing = ConditionalSerialize { value: String::from("failing"), fail: true }; + + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + updater.set_flush_policy(FlushPolicy::Batch); + updater.upsert(&staged, &String::from("new"))?; + assert!(updater.delete(&failing).is_err()); + updater.commit()?; + + assert_eq!(fs::read(&path)?, before); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn idle_updater_refreshes_after_full_replacement() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("replacement-refresh.db"); + let mut original = BPlusTree::new(); + original.insert(1u32, String::from("original")); + original.store(&path)?; + let mut updater = BPlusTreeUpdate::::try_new(&path)?; + let original_id = updater.database_id; + + let mut replacement = BPlusTree::new(); + replacement.insert(2u32, String::from("replacement")); + replacement.store(&path)?; + let replacement_id = database_header(&path)?.database_id; + assert_ne!(replacement_id, original_id); + + updater.upsert(&3, &String::from("updated"))?; + + assert_eq!(updater.database_id, replacement_id); + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, None); + assert_eq!(query.query(&2).map_err(BPlusTreeError::to_io)?, Some(String::from("replacement"))); + assert_eq!(query.query(&3).map_err(BPlusTreeError::to_io)?, Some(String::from("updated"))); + let _ = verify_full(&mut query)?; + Ok(()) + } + + fn pending_path(database: &Path, suffix: &str) -> PathBuf { + let mut name = database.as_os_str().to_os_string(); + name.push(suffix); + PathBuf::from(name) + } + + fn spawn_replacement_writer(path: PathBuf) -> io::Result<(Receiver>, JoinHandle<()>)> { + let (started_sender, started_receiver) = mpsc::channel(); + let (result_sender, result_receiver) = mpsc::channel(); + let handle = thread::spawn(move || { + let mut replacement = BPlusTree::new(); + replacement.insert(2u32, String::from("replacement")); + let _ = started_sender.send(()); + let _ = result_sender.send(replacement.store(&path)); + }); + started_receiver + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| io::Error::other(format!("replacement writer did not start: {error}")))?; + Ok((result_receiver, handle)) + } + + fn assert_writer_is_blocked(receiver: &Receiver>) -> io::Result<()> { + match receiver.recv_timeout(Duration::from_millis(100)) { + Err(RecvTimeoutError::Timeout) => Ok(()), + Err(RecvTimeoutError::Disconnected) => Err(io::Error::other("replacement writer disconnected")), + Ok(result) => { + let root = result?; + Err(io::Error::other(format!("replacement writer completed early with root page {root}"))) + } + } + } + + fn finish_writer(receiver: &Receiver>, handle: JoinHandle<()>) -> io::Result<()> { + receiver + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| io::Error::other(format!("replacement writer stayed blocked: {error}")))??; + handle.join().map_err(|_| io::Error::other("replacement writer panicked")) + } + + fn try_exclusive_sidecar(database: &Path) -> io::Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(crate::repository::bplustree::common::sidecar_lock_path(database))?; + match file.try_lock_exclusive() { + Ok(()) => { + fs2::FileExt::unlock(&file)?; + Ok(true) + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(false), + Err(error) => Err(error), + } + } + + fn run_exclusive_probe_child(database: &Path, expected: &str) -> io::Result<()> { + let status = Command::new(std::env::current_exe()?) + .arg("--exact") + .arg("repository::bplustree::v3::tree::tests::exclusive_sidecar_probe_child") + .arg("--nocapture") + .env("TULIPROX_V3_LOCK_PROBE_PATH", database) + .env("TULIPROX_V3_LOCK_PROBE_EXPECTED", expected) + .status()?; + if status.success() { + Ok(()) + } else { + Err(io::Error::other(format!("exclusive sidecar child probe failed with {status}"))) + } + } + + #[test] + fn exclusive_sidecar_probe_child() -> io::Result<()> { + let Some(path) = std::env::var_os("TULIPROX_V3_LOCK_PROBE_PATH") else { + return Ok(()); + }; + let expected = std::env::var("TULIPROX_V3_LOCK_PROBE_EXPECTED") + .map_err(|error| io::Error::other(format!("missing child probe expectation: {error}")))?; + let acquired = try_exclusive_sidecar(Path::new(&path))?; + match expected.as_str() { + "acquired" => assert!(acquired), + "blocked" => assert!(!acquired), + _ => return Err(io::Error::other(format!("unknown child probe expectation: {expected}"))), + } + Ok(()) + } + + #[test] + fn shared_queries_coexist_and_block_replacement_until_both_drop() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("shared-queries.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("original")); + tree.store(&path)?; + + let first = BPlusTreeQuery::::try_new(&path)?; + let second = first.try_clone()?; + assert!(!try_exclusive_sidecar(&path)?); + drop(first); + assert!(!try_exclusive_sidecar(&path)?); + let (receiver, handle) = spawn_replacement_writer(path.clone())?; + assert_writer_is_blocked(&receiver)?; + drop(second); + finish_writer(&receiver, handle)?; + assert!(try_exclusive_sidecar(&path)?); + Ok(()) + } + + #[test] + fn owned_iterator_keeps_shared_guard_after_query_is_consumed() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("owned-iterator.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("original")); + tree.store(&path)?; + + let query = BPlusTreeQuery::::try_new(&path)?; + let iterator = query.disk_iter(); + assert!(!try_exclusive_sidecar(&path)?); + let (receiver, handle) = spawn_replacement_writer(path.clone())?; + assert_writer_is_blocked(&receiver)?; + drop(iterator); + finish_writer(&receiver, handle)?; + assert!(try_exclusive_sidecar(&path)?); + Ok(()) + } + + #[test] + fn shared_query_blocks_exclusive_writer_in_another_process() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("two-process.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("original")); + tree.store(&path)?; + + let query = BPlusTreeQuery::::try_new(&path)?; + run_exclusive_probe_child(&path, "blocked")?; + drop(query); + run_exclusive_probe_child(&path, "acquired") + } + + #[test] + fn query_removes_abandoned_wal_temp_and_rejects_corrupt_active_wal() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("pending-recovery.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("original")); + tree.store(&path)?; + + let temporary = pending_path(&path, ".wal.tmp"); + fs::write(&temporary, b"not activated")?; + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, Some(String::from("original"))); + assert!(!temporary.try_exists()?); + drop(query); + + let active = pending_path(&path, ".wal"); + fs::write(&active, b"corrupt active WAL")?; + let active_before = fs::read(&active)?; + let error = BPlusTreeQuery::::try_new(&path) + .err() + .ok_or_else(|| io::Error::other("query accepted corrupt active WAL"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(fs::read(active)?, active_before); + Ok(()) + } + + #[test] + fn replacement_store_removes_wal_temp_and_preserves_corrupt_active_wal() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("pending-store.db"); + let mut original = BPlusTree::new(); + original.insert(1u32, String::from("original")); + original.store(&path)?; + let database_before = fs::read(&path)?; + + let temporary = pending_path(&path, ".wal.tmp"); + fs::write(&temporary, b"not activated")?; + let mut replacement = BPlusTree::new(); + replacement.insert(2u32, String::from("replacement")); + replacement.store(&path)?; + assert!(!temporary.try_exists()?); + + let published_before = fs::read(&path)?; + let active = pending_path(&path, ".wal"); + let active_before = b"corrupt active WAL".to_vec(); + fs::write(&active, &active_before)?; + let mut rejected = BPlusTree::new(); + rejected.insert(3u32, String::from("rejected")); + let error = rejected.store(&path).err().ok_or_else(|| io::Error::other("store accepted corrupt WAL"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_ne!(published_before, database_before); + assert_eq!(fs::read(&path)?, published_before); + assert_eq!(fs::read(active)?, active_before); + Ok(()) + } + + #[test] + fn clean_loaded_store_recovers_temp_but_rejects_corrupt_active_wal() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("pending-clean-store.db"); + let mut original = BPlusTree::new(); + original.insert(1u32, String::from("original")); + original.store(&path)?; + let database_before = fs::read(&path)?; + let mut loaded = BPlusTree::::load(&path)?; + + let temporary = pending_path(&path, ".wal.tmp"); + fs::write(&temporary, b"not activated")?; + assert_eq!(loaded.store(&path)?, 0); + assert!(!temporary.try_exists()?); + assert_eq!(fs::read(&path)?, database_before); + + let active = pending_path(&path, ".wal"); + let active_before = b"corrupt active WAL".to_vec(); + fs::write(&active, &active_before)?; + let error = loaded.store(&path).err().ok_or_else(|| io::Error::other("clean store accepted corrupt WAL"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(fs::read(&path)?, database_before); + assert_eq!(fs::read(active)?, active_before); + Ok(()) + } + + fn empty_leaf(page_id: u64, next_page_id: u64) -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0; PAGE_SIZE]; + PageHeader { + page_type: PageType::Leaf, + cell_count: 0, + free_start: u16::try_from(PAGE_HEADER_LEN).map_err(io::Error::other)?, + free_end: u16::try_from(PAGE_SIZE).map_err(io::Error::other)?, + left: 0, + right: 0, + } + .encode_into(&mut page, page_id, next_page_id)?; + Ok(page) + } + + fn leaf_cell_of_footprint(footprint: usize) -> io::Result> { + let key_length = footprint + .checked_sub(4 + 24) + .ok_or_else(|| io::Error::other("footprint is too small"))?; + if footprint <= 2032 { + let key = vec![b'k'; key_length]; + let mut cell = Vec::new(); + encode_tombstone_leaf_cell(&key, &mut cell)?; + Ok(cell) + } else { + Ok(vec![0; footprint - 4]) + } + } + + #[test] + fn typed_leaf_cells_round_trip_and_validate_value_crc() -> io::Result<()> { + let mut cell = Vec::new(); + encode_inline_leaf_cell(b"key", 5, Compression::None, b"value", &mut cell)?; + let decoded = LeafCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(decoded.key_bytes, b"key"); + match decoded.value { + LeafValueRef::Inline { compression, logical_len, stored, crc32 } => { + assert_eq!(compression, Compression::None); + assert_eq!(logical_len, 5); + assert_eq!(stored, b"value"); + assert_eq!(crc32, crc32fast::hash(b"value")); + } + _ => return Err(io::Error::other("expected inline value")), + } + + let last = cell.len().checked_sub(1).ok_or_else(|| io::Error::other("empty test cell"))?; + cell[last] ^= 1; + invalid_data(LeafCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)) + } + + #[test] + fn compressed_inline_tombstone_and_overflow_descriptors_round_trip() -> io::Result<()> { + let raw = [0u8; 128]; + let mut compression_scratch = Vec::new(); + let stored = encode_value(&raw, &mut compression_scratch)?; + assert_eq!(stored.compression(), Compression::Lz4); + + let mut cell = Vec::new(); + encode_inline_leaf_cell( + b"compressed", + u32::try_from(raw.len()).map_err(io::Error::other)?, + stored.compression(), + stored.as_slice(), + &mut cell, + )?; + let decoded = LeafCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)?; + let mut value_scratch = Vec::new(); + assert_eq!(read_leaf_value(&[], &decoded.value, NEXT_PAGE_ID, 1024, &mut value_scratch)?, Some(raw.as_slice())); + + encode_tombstone_leaf_cell(b"deleted", &mut cell)?; + assert!(matches!(LeafCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)?.value, LeafValueRef::Tombstone)); + + encode_overflow_leaf_cell(b"large", 5000, Compression::None, 5000, 2, 0x1234_5678, PAGE_ID, NEXT_PAGE_ID, &mut cell)?; + assert!(matches!(LeafCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)?.value, LeafValueRef::Overflow { head: 2, .. })); + Ok(()) + } + + #[test] + fn internal_separator_and_locator_have_exact_codecs() -> io::Result<()> { + let mut cell = Vec::new(); + encode_internal_cell(b"separator", 9, PAGE_ID, NEXT_PAGE_ID, &mut cell)?; + let decoded = InternalCellRef::decode(&cell, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(decoded.key_bytes, b"separator"); + assert_eq!(decoded.right_child, 9); + + let locator = Locator::for_key(7, 3, b"separator")?; + let encoded = locator.encode(); + assert_eq!(encoded.len(), 16); + assert_eq!(Locator::decode(&encoded)?, locator); + let mut corrupt = encoded; + corrupt[10] = 1; + invalid_data(Locator::decode(&corrupt)) + } + + #[test] + fn locator_rejects_wrong_key_crc_and_wrong_primary_key() -> io::Result<()> { + let key = binary_serialize(&42u32)?; + let mut cell = Vec::new(); + encode_inline_leaf_cell(&key, 1, Compression::None, &[7], &mut cell)?; + let mut bytes = empty_leaf(PAGE_ID, NEXT_PAGE_ID)?; + let mut page = SlottedPage::open(bytes.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + page.rebuild_ordered([cell.as_slice()])?; + + let locator = Locator::for_key(PAGE_ID, 0, &key)?; + validate_locator(&page, locator, &key)?; + let bad_crc = Locator { serialized_key_crc32: locator.serialized_key_crc32 ^ 1, ..locator }; + invalid_data(validate_locator(&page, bad_crc, &key))?; + invalid_data(validate_locator(&page, locator, &binary_serialize(&43u32)?)) + } + + #[test] + fn encoded_key_limit_is_2004_bytes() -> io::Result<()> { + let mut cell = Vec::new(); + encode_tombstone_leaf_cell(&vec![b'k'; 2004], &mut cell)?; + assert_eq!(cell.len() + 4, 2032); + invalid_input(encode_tombstone_leaf_cell(&vec![b'k'; 2005], &mut cell))?; + + encode_internal_cell(&vec![b'k'; 2004], 9, PAGE_ID, NEXT_PAGE_ID, &mut cell)?; + invalid_input(encode_internal_cell(&vec![b'k'; 2005], 9, PAGE_ID, NEXT_PAGE_ID, &mut cell)) + } + + #[test] + fn adversarial_leaf_splits_reject_old_limit_and_accept_capped_cells() -> io::Result<()> { + let first_witness = [ + leaf_cell_of_footprint(2026)?, + leaf_cell_of_footprint(2040)?, + leaf_cell_of_footprint(2038)?, + ]; + invalid_input(choose_leaf_split(&first_witness))?; + + let second_witness = [ + leaf_cell_of_footprint(1984)?, + leaf_cell_of_footprint(2296)?, + leaf_cell_of_footprint(1984)?, + ]; + invalid_input(choose_leaf_split(&second_witness))?; + + let capped = [ + leaf_cell_of_footprint(2026)?, + leaf_cell_of_footprint(2032)?, + leaf_cell_of_footprint(2032)?, + ]; + let split = choose_leaf_split(&capped)?; + assert_eq!(split, 2); + assert!(used_leaf_bytes(&capped[..split])? <= PAGE_SIZE); + assert!(used_leaf_bytes(&capped[split..])? <= PAGE_SIZE); + Ok(()) + } + + #[test] + fn leaf_split_selection_is_deterministic_for_variable_sizes() -> io::Result<()> { + let mut seed = 0x5eed_u64; + for count in 3..36 { + let mut cells = Vec::with_capacity(count); + for _ in 0..count { + seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + let footprint = 100 + usize::try_from(seed % 60).map_err(io::Error::other)?; + let insert_at = if cells.is_empty() { + 0 + } else { + usize::try_from(seed).map_err(io::Error::other)? % (cells.len() + 1) + }; + cells.insert(insert_at, leaf_cell_of_footprint(footprint)?); + } + if used_leaf_bytes(&cells)? > PAGE_SIZE { + let first = choose_leaf_split(&cells)?; + let second = choose_leaf_split(&cells)?; + assert_eq!(first, second); + assert!(used_leaf_bytes(&cells[..first])? <= PAGE_SIZE); + assert!(used_leaf_bytes(&cells[first..])? <= PAGE_SIZE); + } + } + Ok(()) + } + + #[test] + fn internal_split_promotes_exact_separator_and_child() -> io::Result<()> { + let mut cells = Vec::new(); + for (key_length, child) in [(1000, 2), (1000, 3), (1000, 4), (1000, 5)] { + let mut cell = Vec::new(); + encode_internal_cell(&vec![b'k'; key_length], child, PAGE_ID, NEXT_PAGE_ID, &mut cell)?; + cells.push(cell); + } + let split = choose_internal_split(&cells, PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(split.promoted_index, 1); + assert_eq!(split.promoted.key_bytes.len(), 1000); + assert_eq!(split.right_leftmost_child, 3); + assert_eq!(split.left_cells, &cells[..1]); + assert_eq!(split.right_cells, &cells[2..]); + Ok(()) + } + + #[test] + fn typed_search_decodes_only_binary_search_candidates() -> io::Result<()> { + let mut cells = Vec::new(); + for key in 0u8..8 { + let mut cell = Vec::new(); + encode_inline_leaf_cell(&binary_serialize(&key)?, 1, Compression::None, &[key], &mut cell)?; + cells.push(cell); + } + if let Some(byte) = cells.get_mut(0).and_then(|cell| cell.get_mut(24)) { + *byte = 0xc1; + } + + let mut bytes = empty_leaf(PAGE_ID, NEXT_PAGE_ID)?; + let mut mutable = SlottedPage::open(bytes.as_mut_slice(), PAGE_ID, NEXT_PAGE_ID)?; + mutable.rebuild_ordered(cells.iter().map(Vec::as_slice))?; + let page = SlottedPage::open(bytes.as_slice(), PAGE_ID, NEXT_PAGE_ID)?; + assert_eq!(search_leaf(&page, &7u8)?, Ok(7)); + invalid_data(search_leaf(&page, &0u8)) + } + + #[test] + fn overflow_chain_round_trip_is_bounded_and_checksum_validated() -> io::Result<()> { + let stored = vec![0x5a; 5000]; + let mut database = vec![0; PAGE_SIZE * 4]; + let first = encode_overflow_page(2, NEXT_PAGE_ID, 3, &stored[..4056])?; + let second = encode_overflow_page(3, NEXT_PAGE_ID, 0, &stored[4056..])?; + database[PAGE_SIZE * 2..PAGE_SIZE * 3].copy_from_slice(&first); + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&second); + + let value = LeafValueRef::Overflow { + compression: Compression::None, + logical_len: 5000, + stored_len: 5000, + head: 2, + crc32: crc32fast::hash(&stored), + }; + let mut scratch = Vec::new(); + assert_eq!(read_leaf_value(&database, &value, NEXT_PAGE_ID, 6000, &mut scratch)?, Some(stored.as_slice())); + + let free = encode_free_page(3, NEXT_PAGE_ID, 0)?; + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&free); + invalid_data(read_leaf_value(&database, &value, NEXT_PAGE_ID, 6000, &mut scratch)) + } + + #[test] + fn compressed_overflow_chain_reuses_scratch_and_round_trips() -> io::Result<()> { + let mut raw = Vec::with_capacity(10_000); + let mut state = 0x51a7_9e2d_4c83_b6f0_u64; + for _ in 0..5000 { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + raw.push(state.to_le_bytes()[0]); + } + raw.resize(10_000, 0); + let mut encoded = Vec::new(); + let stored = encode_value(&raw, &mut encoded)?; + if stored.compression() != Compression::Lz4 || stored.as_slice().len() <= OVERFLOW_PAYLOAD_LEN { + return Err(io::Error::other("test value must span compressed overflow pages")); + } + let stored = stored.as_slice().to_vec(); + let mut corrupt = stored.clone(); + corrupt + .get_mut(..4) + .ok_or_else(|| io::Error::other("missing test LZ4 length"))? + .copy_from_slice(&1u32.to_le_bytes()); + invalid_data(decompress_value_in_place( + &mut corrupt, + u32::try_from(raw.len()).map_err(io::Error::other)?, + raw.len(), + ))?; + let page_count = stored.len().div_ceil(OVERFLOW_PAYLOAD_LEN); + assert!(page_count > 1); + let next_page_id = u64::try_from(page_count) + .map_err(io::Error::other)? + .checked_add(2) + .ok_or_else(|| io::Error::other("test page count overflow"))?; + let mut database = vec![0; usize::try_from(next_page_id).map_err(io::Error::other)? * PAGE_SIZE]; + for (index, payload) in stored.chunks(OVERFLOW_PAYLOAD_LEN).enumerate() { + let page_id = u64::try_from(index).map_err(io::Error::other)? + 1; + let next = if index + 1 == page_count { 0 } else { page_id + 1 }; + let page = encode_overflow_page(page_id, next_page_id, next, payload)?; + let start = usize::try_from(page_id).map_err(io::Error::other)? * PAGE_SIZE; + database[start..start + PAGE_SIZE].copy_from_slice(&page); + } + let value = LeafValueRef::Overflow { + compression: Compression::Lz4, + logical_len: u32::try_from(raw.len()).map_err(io::Error::other)?, + stored_len: u32::try_from(stored.len()).map_err(io::Error::other)?, + head: 1, + crc32: crc32fast::hash(&stored), + }; + let mut scratch = Vec::new(); + assert_eq!( + read_leaf_value(&database, &value, next_page_id, raw.len(), &mut scratch)?, + Some(raw.as_slice()) + ); + Ok(()) + } + + #[test] + fn overflow_chain_rejects_truncation_cycle_oversize_and_bad_crc() -> io::Result<()> { + let stored = vec![0x5a; 5000]; + let value = LeafValueRef::Overflow { + compression: Compression::None, + logical_len: 5000, + stored_len: 5000, + head: 2, + crc32: crc32fast::hash(&stored), + }; + let mut database = vec![0; PAGE_SIZE * 4]; + let first = encode_overflow_page(2, NEXT_PAGE_ID, 3, &stored[..4056])?; + let short = encode_overflow_page(3, NEXT_PAGE_ID, 0, &stored[4056..4999])?; + database[PAGE_SIZE * 2..PAGE_SIZE * 3].copy_from_slice(&first); + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&short); + let mut scratch = Vec::new(); + invalid_data(read_leaf_value(&database, &value, NEXT_PAGE_ID, 6000, &mut scratch))?; + + let cycle = encode_overflow_page(3, NEXT_PAGE_ID, 2, &stored[4056..])?; + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&cycle); + invalid_data(read_leaf_value(&database, &value, NEXT_PAGE_ID, 6000, &mut scratch))?; + invalid_data(read_leaf_value(&database, &value, NEXT_PAGE_ID, 4999, &mut scratch))?; + + let last = encode_overflow_page(3, NEXT_PAGE_ID, 0, &stored[4056..])?; + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&last); + let LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32 } = value else { + return Err(io::Error::other("expected overflow value")); + }; + let bad_crc = LeafValueRef::Overflow { compression, logical_len, stored_len, head, crc32: crc32 ^ 1 }; + invalid_data(read_leaf_value(&database, &bad_crc, NEXT_PAGE_ID, 6000, &mut scratch))?; + + database[PAGE_SIZE * 2 + 40] ^= 1; + invalid_data(read_leaf_value(&database, &value, NEXT_PAGE_ID, 6000, &mut scratch))?; + + let empty = encode_overflow_page(2, NEXT_PAGE_ID, 3, &[])?; + let one_byte = encode_overflow_page(3, NEXT_PAGE_ID, 0, b"x")?; + database[PAGE_SIZE * 2..PAGE_SIZE * 3].copy_from_slice(&empty); + database[PAGE_SIZE * 3..PAGE_SIZE * 4].copy_from_slice(&one_byte); + let empty_chain_value = LeafValueRef::Overflow { + compression: Compression::None, + logical_len: 1, + stored_len: 1, + head: 2, + crc32: crc32fast::hash(b"x"), + }; + invalid_data(read_leaf_value( + &database, + &empty_chain_value, + NEXT_PAGE_ID, + 6000, + &mut scratch, + )) + } + + fn page_range(page_id: u64) -> io::Result> { + let start = usize::try_from(page_id) + .map_err(io::Error::other)? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| io::Error::other("test page offset overflow"))?; + Ok(start..start + PAGE_SIZE) + } + + fn rewrite_page_checksum(database: &mut [u8], page_id: u64) -> io::Result<()> { + let range = page_range(page_id)?; + crate::repository::bplustree::v3::format::write_page_checksum( + database.get_mut(range).ok_or_else(|| io::Error::other("test page missing"))?, + ) + } + + fn set_page_reference(database: &mut [u8], page_id: u64, offset: usize, reference: u64) -> io::Result<()> { + let page = page_range(page_id)?; + let start = page.start + offset; + database + .get_mut(start..start + 8) + .ok_or_else(|| io::Error::other("test reference missing"))? + .copy_from_slice(&reference.to_le_bytes()); + rewrite_page_checksum(database, page_id) + } + + fn verify_rejects(database: &[u8]) -> io::Result<()> + where + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, + V: Serialize + for<'de> Deserialize<'de> + Clone, + { + let dir = tempfile::tempdir()?; + let path = dir.path().join("corrupt.db"); + fs::write(&path, database)?; + let mut query = BPlusTreeQuery::::try_new(&path)?; + invalid_data(verify_full(&mut query)) + } + + fn stored_tree_fixture() -> io::Result<(Vec, u64, u64, u64, u64)> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("tree.db"); + let mut tree = BPlusTree::new(); + for key in 1_000..1_700u32 { + tree.insert(key, "same-value".to_string()); + } + tree.store(&path)?; + let mut query = BPlusTreeQuery::::try_new(&path)?; + let root = query.header.root_page_id; + let first = query.leftmost_leaf()?; + let next_page_id = query.header.next_page_id; + let second = query.with_page(first, |bytes, _| { + SlottedPage::open(bytes, first, next_page_id).map(|page| page.header().right) + })?; + if second == 0 { + return Err(io::Error::other("test tree did not split")); + } + let mut last = second; + loop { + let right = query.with_page(last, |bytes, _| { + SlottedPage::open(bytes, last, next_page_id).map(|page| page.header().right) + })?; + if right == 0 { + break; + } + last = right; + } + drop(query); + Ok((fs::read(path)?, root, first, second, last)) + } + + #[test] + fn iterator_yields_second_leaf_corruption_once_then_fuses() -> io::Result<()> { + let (mut database, _, _, second, _) = stored_tree_fixture()?; + let byte = page_range(second)?.start + 100; + *database.get_mut(byte).ok_or_else(|| io::Error::other("test corruption byte missing"))? ^= 1; + let dir = tempfile::tempdir()?; + let path = dir.path().join("iterator-corrupt.db"); + fs::write(&path, database)?; + let mut query = BPlusTreeQuery::::try_new(&path)?; + let mut iterator = query.iter(); + let mut yielded = 0usize; + let mut saw_error = false; + for item in iterator.by_ref() { + match item { + Ok(_) => yielded += 1, + Err(err) => { + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + saw_error = true; + break; + } + } + } + if !saw_error { + return Err(io::Error::other("corruption was hidden as end-of-stream")); + } + assert!(iterator.next().is_none()); + assert!(yielded > 0); + Ok(()) + } + + #[test] + fn iterator_skips_corrupt_value_and_continues_with_next_cell() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("iterator-corrupt-value.db"); + let mut tree = BPlusTree::new(); + for key in 1..=3u32 { + tree.insert(key, format!("value-{key}")); + } + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + let (page_id, range) = query + .locate_cell(&2)? + .ok_or_else(|| io::Error::other("test key missing"))?; + let next_page_id = query.header.next_page_id; + let stored_offset = query.with_page(page_id, |bytes, _| { + let cell = LeafCellRef::decode( + bytes.get(range).ok_or_else(|| io::Error::other("test cell range missing"))?, + page_id, + next_page_id, + )?; + let LeafValueRef::Inline { stored, .. } = cell.value else { + return Err(io::Error::other("test value is not inline")); + }; + Ok(stored.as_ptr() as usize - bytes.as_ptr() as usize) + })?; + drop(query); + + let mut database = fs::read(&path)?; + let absolute = page_range(page_id)?.start + stored_offset; + *database + .get_mut(absolute) + .ok_or_else(|| io::Error::other("test value byte missing"))? = 0xc1; + rewrite_page_checksum(&mut database, page_id)?; + fs::write(&path, database)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + let mut iterator = query.iter(); + assert_eq!(iterator.next().transpose()?, Some((1, String::from("value-1")))); + assert!(iterator.next().is_some_and(|entry| entry.is_err())); + assert_eq!(iterator.next().transpose()?, Some((3, String::from("value-3")))); + assert!(iterator.next().is_none()); + Ok(()) + } + + #[test] + fn iterator_rejects_leaf_cycle_before_yielding_duplicate_entries() -> io::Result<()> { + let (mut database, _, first, second, _) = stored_tree_fixture()?; + set_page_reference(&mut database, second, 16, first)?; + set_page_reference(&mut database, first, 8, second)?; + let dir = tempfile::tempdir()?; + let path = dir.path().join("iterator-cycle.db"); + fs::write(&path, database)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + let mut iterator = query.iter(); + let mut seen = HashSet::new(); + let mut errors = 0usize; + for item in iterator.by_ref() { + match item { + Ok((key, _)) => assert!(seen.insert(key), "iterator yielded key {key} twice"), + Err(err) => { + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + errors += 1; + } + } + } + assert_eq!(errors, 1); + assert!(iterator.next().is_none()); + Ok(()) + } + + #[test] + fn range_page_does_not_preallocate_the_requested_limit() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("range-limit.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + let (entries, has_more) = query + .range_page(Bound::Unbounded, Bound::Unbounded, 0, usize::MAX) + .map_err(BPlusTreeError::to_io)?; + assert_eq!(entries, vec![(1, String::from("one")), (2, String::from("two"))]); + assert!(!has_more); + Ok(()) + } + + #[derive(Debug)] + struct RejectValueDeserialize; + + impl<'de> Deserialize<'de> for RejectValueDeserialize { + fn deserialize>(_deserializer: D) -> Result { + Err(serde::de::Error::custom("value deserialization must not run")) + } + } + + #[test] + fn contains_and_len_read_only_leaf_descriptors() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("descriptor-only.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.store(&path)?; + + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert!(query.contains_live_key(&1).map_err(BPlusTreeError::to_io)?); + assert!(!query.contains_live_key(&3).map_err(BPlusTreeError::to_io)?); + assert_eq!(query.len().map_err(BPlusTreeError::to_io)?, 2); + Ok(()) + } + + #[test] + fn full_verifier_rejects_tree_sibling_free_and_orphan_corruption() -> io::Result<()> { + let (database, root, first, second, last) = stored_tree_fixture()?; + + let mut child_cycle = database.clone(); + set_page_reference(&mut child_cycle, root, 32, root)?; + verify_rejects::(&child_cycle)?; + + let mut duplicate_child = database.clone(); + let root_start = page_range(root)?.start; + let leftmost = u64::from_le_bytes( + duplicate_child[root_start + 32..root_start + 40] + .try_into() + .map_err(io::Error::other)?, + ); + let first_cell = u16::from_le_bytes( + duplicate_child[root_start + 40..root_start + 42] + .try_into() + .map_err(io::Error::other)?, + ); + let child_offset = root_start + usize::from(first_cell) + 4; + duplicate_child[child_offset..child_offset + 8].copy_from_slice(&leftmost.to_le_bytes()); + rewrite_page_checksum(&mut duplicate_child, root)?; + verify_rejects::(&duplicate_child)?; + + let mut wrong_type = database.clone(); + let header = DatabaseHeader::decode(&wrong_type[..PAGE_SIZE])?; + let overflow = encode_overflow_page(leftmost, header.next_page_id, 0, b"wrong type")?; + let range = page_range(leftmost)?; + wrong_type[range].copy_from_slice(&overflow); + verify_rejects::(&wrong_type)?; + + let mut asymmetric = database.clone(); + set_page_reference(&mut asymmetric, second, 8, 0)?; + verify_rejects::(&asymmetric)?; + + let mut sibling_cycle = database.clone(); + set_page_reference(&mut sibling_cycle, last, 16, first)?; + set_page_reference(&mut sibling_cycle, first, 8, last)?; + verify_rejects::(&sibling_cycle)?; + + let mut inverted = database.clone(); + let first_start = page_range(first)?.start; + let second_start = page_range(second)?.start; + let first_count = u16::from_le_bytes( + inverted[first_start + 2..first_start + 4] + .try_into() + .map_err(io::Error::other)?, + ); + let last_slot = first_start + PAGE_HEADER_LEN + (usize::from(first_count) - 1) * SLOT_LEN; + let left_cell = usize::from(u16::from_le_bytes( + inverted[last_slot..last_slot + 2].try_into().map_err(io::Error::other)?, + )); + let right_cell = usize::from(u16::from_le_bytes( + inverted[second_start + PAGE_HEADER_LEN..second_start + PAGE_HEADER_LEN + 2] + .try_into() + .map_err(io::Error::other)?, + )); + for index in 0..3 { + inverted.swap(first_start + left_cell + 24 + index, second_start + right_cell + 24 + index); + } + rewrite_page_checksum(&mut inverted, first)?; + rewrite_page_checksum(&mut inverted, second)?; + verify_rejects::(&inverted)?; + + let mut live_and_free = database.clone(); + let mut header = DatabaseHeader::decode(&live_and_free[..PAGE_SIZE])?; + header.free_page_head = first; + live_and_free[..PAGE_SIZE].copy_from_slice(&header.encode()?); + verify_rejects::(&live_and_free)?; + + let mut orphan = database.clone(); + let mut header = DatabaseHeader::decode(&orphan[..PAGE_SIZE])?; + let orphan_id = header.next_page_id; + header.next_page_id += 1; + orphan[..PAGE_SIZE].copy_from_slice(&header.encode()?); + orphan.extend_from_slice(&encode_free_page(orphan_id, header.next_page_id, 0)?); + verify_rejects::(&orphan)?; + + let mut duplicate_free = database; + let mut header = DatabaseHeader::decode(&duplicate_free[..PAGE_SIZE])?; + let first_free = header.next_page_id; + let second_free = first_free + 1; + header.next_page_id += 2; + header.free_page_head = first_free; + duplicate_free[..PAGE_SIZE].copy_from_slice(&header.encode()?); + duplicate_free.extend_from_slice(&encode_free_page(first_free, header.next_page_id, second_free)?); + duplicate_free.extend_from_slice(&encode_free_page(second_free, header.next_page_id, first_free)?); + verify_rejects::(&duplicate_free) + } + + fn random_value() -> Vec { + let mut state = 0x1234_5678_9abc_def0u64; + (0..12_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.to_le_bytes()[0] + }) + .collect() + } + + fn overflow_heads(query: &mut BPlusTreeQuery>) -> io::Result<(u64, u64, u64)> { + let mut result = Vec::new(); + for key in [1, 2] { + let leaf = query.locate_leaf(&key)?; + let next_page_id = query.header.next_page_id; + let (head, cell_offset) = query.with_page(leaf, |bytes, _| { + let page = SlottedPage::open(bytes, leaf, next_page_id)?; + let index = search_leaf(&page, &key)?.map_err(|_| io::Error::other("test key missing"))?; + let cell = LeafCellRef::decode(page.cell(index)?, leaf, next_page_id)?; + let LeafValueRef::Overflow { head, .. } = cell.value else { + return Err(io::Error::other("test value is not overflow-backed")); + }; + let slot = PAGE_HEADER_LEN + index * SLOT_LEN; + let offset = u16::from_le_bytes(bytes[slot..slot + 2].try_into().map_err(io::Error::other)?); + Ok((head, u64::from(offset))) + })?; + result.push((leaf, head, cell_offset)); + } + let first = result.first().ok_or_else(|| io::Error::other("first test overflow missing"))?; + let second = result.get(1).ok_or_else(|| io::Error::other("second test overflow missing"))?; + Ok((first.1, second.0, second.2)) + } + + #[test] + fn full_verifier_rejects_overflow_cycle_and_shared_chain() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("overflow.db"); + let value = random_value(); + let mut tree = BPlusTree::new(); + tree.insert(1u32, value.clone()); + tree.insert(2u32, value); + tree.store(&path)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let (first_head, second_leaf, second_cell_offset) = overflow_heads(&mut query)?; + let next_page_id = query.header.next_page_id; + drop(query); + let database = fs::read(&path)?; + + let mut cycle = database.clone(); + let mut last = first_head; + loop { + let start = page_range(last)?.start; + let next = u64::from_le_bytes(cycle[start + 16..start + 24].try_into().map_err(io::Error::other)?); + if next == 0 { + break; + } + last = next; + } + set_page_reference(&mut cycle, last, 16, first_head)?; + verify_rejects::>(&cycle)?; + + let mut shared = database; + let start = page_range(second_leaf)?.start + + usize::try_from(second_cell_offset).map_err(io::Error::other)? + + 12; + shared[start..start + 8].copy_from_slice(&first_head.to_le_bytes()); + rewrite_page_checksum(&mut shared, second_leaf)?; + let _ = next_page_id; + verify_rejects::>(&shared) + } + + #[test] + fn full_verifier_rejects_oversized_overflow_descriptor_before_reserve() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("oversized-overflow.db"); + let value = random_value(); + let mut tree = BPlusTree::new(); + tree.insert(1u32, value.clone()); + tree.insert(2u32, value); + tree.store(&path)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let (_, second_leaf, second_cell_offset) = overflow_heads(&mut query)?; + drop(query); + + let mut database = fs::read(&path)?; + let descriptor = page_range(second_leaf)?.start + + usize::try_from(second_cell_offset).map_err(io::Error::other)?; + database[descriptor + 4..descriptor + 8].copy_from_slice(&u32::MAX.to_le_bytes()); + database[descriptor + 8..descriptor + 12].copy_from_slice(&u32::MAX.to_le_bytes()); + rewrite_page_checksum(&mut database, second_leaf)?; + fs::write(&path, database)?; + + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let error = verify_full(&mut query).err().ok_or_else(|| io::Error::other("oversized value accepted"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert!(error.to_string().contains("overflow value exceeds allocation limit")); + Ok(()) + } + + #[test] + fn full_verifier_rejects_uncompressed_overflow_length_mismatch() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("overflow-length.db"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, random_value()); + tree.store(&path)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + let leaf = query.locate_leaf(&1)?; + let next_page_id = query.header.next_page_id; + let cell_offset = query.with_page(leaf, |bytes, _| { + let page = SlottedPage::open(bytes, leaf, next_page_id)?; + let index = search_leaf(&page, &1)?.map_err(|_| io::Error::other("test key missing"))?; + let slot = PAGE_HEADER_LEN + .checked_add(index.checked_mul(SLOT_LEN).ok_or_else(|| io::Error::other("test slot overflow"))?) + .ok_or_else(|| io::Error::other("test slot overflow"))?; + Ok(u16::from_le_bytes( + bytes[slot..slot + 2].try_into().map_err(io::Error::other)?, + )) + })?; + drop(query); + + let mut database = fs::read(&path)?; + let descriptor = page_range(leaf)?.start + usize::from(cell_offset); + database[descriptor + 4..descriptor + 8].copy_from_slice(&1u32.to_le_bytes()); + rewrite_page_checksum(&mut database, leaf)?; + fs::write(&path, database)?; + let mut query = BPlusTreeQuery::>::try_new(&path)?; + invalid_data(verify_full(&mut query)) + } + + #[test] + fn publish_reports_post_commit_directory_sync_failure() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let destination = dir.path().join("database.db"); + let temporary = dir.path().join("database.tx.v3.tmp"); + let mut old = BPlusTree::new(); + old.insert(1u32, vec![1]); + old.store(&destination)?; + let mut new = BPlusTree::new(); + new.insert(2u32, vec![2]); + new.store(&temporary)?; + + let error = publish_database(&temporary, &destination, |_| { + Err(io::Error::other("injected directory sync failure")) + }) + .err() + .ok_or_else(|| io::Error::other("post-commit sync failure was hidden"))?; + assert!(error.to_string().contains("database published but directory sync failed; durability unknown")); + assert!(!temporary.exists()); + let mut query = BPlusTreeQuery::>::try_new(&destination)?; + assert_eq!(query.query(&1).map_err(BPlusTreeError::to_io)?, None); + assert_eq!(query.query(&2).map_err(BPlusTreeError::to_io)?, Some(vec![2])); + Ok(()) + } + + #[test] + fn publish_failure_removes_the_temporary_file() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let destination = dir.path().join("destination-directory"); + fs::create_dir(&destination)?; + let temporary = dir.path().join("database.tx.v3.tmp"); + fs::write(&temporary, b"new")?; + + assert!(publish_database(&temporary, &destination, sync_parent_directory).is_err()); + assert!(!temporary.exists()); + Ok(()) + } +} diff --git a/backend/src/repository/bplustree/v3/wal.rs b/backend/src/repository/bplustree/v3/wal.rs new file mode 100644 index 000000000..ca41c9e30 --- /dev/null +++ b/backend/src/repository/bplustree/v3/wal.rs @@ -0,0 +1,1992 @@ +use super::{ + format::{DatabaseHeader, PAGE_SIZE}, + page::SlottedPage, +}; +use crate::repository::bplustree::common::sidecar_lock_path; +use fs2::FileExt as _; +use log::info; +use std::{ + collections::HashSet, + error::Error, + ffi::OsString, + fmt, + fs::{File, OpenOptions}, + io::{self, Read, Seek, SeekFrom, Write}, + path::{Path, PathBuf}, +}; + +const WAL_HEADER_LEN: usize = 64; +const RECORD_HEADER_LEN: usize = 16; +const BEFORE_IMAGE_PAYLOAD_LEN: usize = 16 + PAGE_SIZE; +const COMMIT_PAYLOAD_LEN: usize = 32; +const WAL_HEADER_LEN_U32: u32 = 64; +const WAL_HEADER_LEN_U64: u64 = 64; +const RECORD_HEADER_LEN_U64: u64 = 16; +const PAGE_SIZE_U32: u32 = 4096; +const PAGE_SIZE_U64: u64 = 4096; +const WAL_MAGIC: &[u8; 4] = b"BTW3"; +const WAL_VERSION: u32 = 1; + +fn invalid_data(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, message.into()) +} + +fn invalid_input(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message.into()) +} + +fn out_of_memory(message: &'static str) -> impl FnOnce(std::collections::TryReserveError) -> io::Error { + move |error| io::Error::new(io::ErrorKind::OutOfMemory, format!("{message}: {error}")) +} + +fn checked_end(offset: usize, length: usize) -> io::Result { + offset.checked_add(length).ok_or_else(|| invalid_data("WAL offset overflow")) +} + +fn bytes_at(bytes: &[u8], offset: usize) -> io::Result<[u8; N]> { + let end = checked_end(offset, N)?; + bytes + .get(offset..end) + .ok_or_else(|| invalid_data("truncated WAL field"))? + .try_into() + .map_err(|_| invalid_data("invalid WAL field length")) +} + +fn read_u32(bytes: &[u8], offset: usize) -> io::Result { Ok(u32::from_le_bytes(bytes_at(bytes, offset)?)) } + +fn read_u64(bytes: &[u8], offset: usize) -> io::Result { Ok(u64::from_le_bytes(bytes_at(bytes, offset)?)) } + +fn require_zero(bytes: &[u8], message: &'static str) -> io::Result<()> { + if bytes.iter().all(|byte| *byte == 0) { Ok(()) } else { Err(invalid_data(message)) } +} + +fn crc_with_zeroed_field(bytes: &[u8], offset: usize) -> io::Result { + let end = checked_end(offset, 4)?; + let mut hasher = crc32fast::Hasher::new(); + hasher.update(bytes.get(..offset).ok_or_else(|| invalid_data("missing checksum prefix"))?); + hasher.update(&[0; 4]); + hasher.update(bytes.get(end..).ok_or_else(|| invalid_data("missing checksum suffix"))?); + Ok(hasher.finalize()) +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct WalHeader { + database_id: [u8; 16], + transaction_id: u64, + original_database_len: u64, + original_generation: u64, +} + +impl WalHeader { + fn encode(&self) -> io::Result<[u8; WAL_HEADER_LEN]> { + self.validate()?; + let mut encoded = [0u8; WAL_HEADER_LEN]; + encoded[0..4].copy_from_slice(WAL_MAGIC); + encoded[4..8].copy_from_slice(&WAL_VERSION.to_le_bytes()); + encoded[8..12].copy_from_slice(&WAL_HEADER_LEN_U32.to_le_bytes()); + encoded[12..16].copy_from_slice(&PAGE_SIZE_U32.to_le_bytes()); + encoded[16..32].copy_from_slice(&self.database_id); + encoded[32..40].copy_from_slice(&self.transaction_id.to_le_bytes()); + encoded[40..48].copy_from_slice(&self.original_database_len.to_le_bytes()); + encoded[48..56].copy_from_slice(&self.original_generation.to_le_bytes()); + let checksum = crc_with_zeroed_field(&encoded, 56)?; + encoded[56..60].copy_from_slice(&checksum.to_le_bytes()); + Ok(encoded) + } + + fn decode(encoded: &[u8]) -> io::Result { + if encoded.len() != WAL_HEADER_LEN { + return Err(invalid_data("WAL header must be exactly 64 bytes")); + } + if bytes_at::<4>(encoded, 0)? != *WAL_MAGIC { + return Err(invalid_data("invalid WAL magic")); + } + if read_u32(encoded, 4)? != WAL_VERSION { + return Err(invalid_data("unsupported WAL version")); + } + if read_u32(encoded, 8)? != WAL_HEADER_LEN_U32 { + return Err(invalid_data("invalid WAL header length")); + } + if read_u32(encoded, 12)? != PAGE_SIZE_U32 { + return Err(invalid_data("invalid WAL page size")); + } + require_zero(&encoded[60..64], "WAL header reserved bytes must be zero")?; + if read_u32(encoded, 56)? != crc_with_zeroed_field(encoded, 56)? { + return Err(invalid_data("WAL header checksum mismatch")); + } + let header = Self { + database_id: bytes_at(encoded, 16)?, + transaction_id: read_u64(encoded, 32)?, + original_database_len: read_u64(encoded, 40)?, + original_generation: read_u64(encoded, 48)?, + }; + header.validate()?; + Ok(header) + } + + fn validate(&self) -> io::Result<()> { + if self.database_id.iter().all(|byte| *byte == 0) { + return Err(invalid_data("WAL database identity must be nonzero")); + } + if self.transaction_id == 0 { + return Err(invalid_data("WAL transaction id must be nonzero")); + } + if self.original_database_len < PAGE_SIZE_U64 + || !self.original_database_len.is_multiple_of(PAGE_SIZE_U64) + { + return Err(invalid_data("invalid original database length")); + } + if self.original_generation == 0 { + return Err(invalid_data("WAL original generation must be nonzero")); + } + Ok(()) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct BeforeImage { + page_id: u64, + page: [u8; PAGE_SIZE], +} + +impl BeforeImage { + #[cfg(test)] + fn encode_record(&self) -> io::Result> { + let mut encoded = Vec::new(); + self.encode_record_into(&mut encoded)?; + Ok(encoded) + } + + fn encode_record_into(&self, encoded: &mut Vec) -> io::Result<()> { + prepare_record(encoded, 1, BEFORE_IMAGE_PAYLOAD_LEN)?; + encoded[16..24].copy_from_slice(&self.page_id.to_le_bytes()); + encoded[24..28].copy_from_slice(&crc32fast::hash(&self.page).to_le_bytes()); + encoded[32..].copy_from_slice(&self.page); + finish_record(encoded) + } + + fn decode_payload(payload: &[u8]) -> io::Result { + if payload.len() != BEFORE_IMAGE_PAYLOAD_LEN { + return Err(invalid_data("invalid before-image payload length")); + } + require_zero(&payload[12..16], "before-image reserved bytes must be zero")?; + let page: [u8; PAGE_SIZE] = payload[16..] + .try_into() + .map_err(|_| invalid_data("truncated before-image page"))?; + if read_u32(payload, 8)? != crc32fast::hash(&page) { + return Err(invalid_data("before-image page checksum mismatch")); + } + Ok(Self { page_id: read_u64(payload, 0)?, page }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct CommitRecord { + transaction_id: u64, + new_generation: u64, + database_length: u64, + database_header_crc32: u32, +} + +impl CommitRecord { + #[cfg(test)] + fn encode_record(self) -> io::Result> { + let mut encoded = Vec::new(); + self.encode_record_into(&mut encoded)?; + Ok(encoded) + } + + fn encode_record_into(self, encoded: &mut Vec) -> io::Result<()> { + prepare_record(encoded, 2, COMMIT_PAYLOAD_LEN)?; + encoded[16..24].copy_from_slice(&self.transaction_id.to_le_bytes()); + encoded[24..32].copy_from_slice(&self.new_generation.to_le_bytes()); + encoded[32..40].copy_from_slice(&self.database_length.to_le_bytes()); + encoded[40..44].copy_from_slice(&self.database_header_crc32.to_le_bytes()); + finish_record(encoded) + } + + fn decode_payload(payload: &[u8]) -> io::Result { + if payload.len() != COMMIT_PAYLOAD_LEN { + return Err(invalid_data("invalid commit payload length")); + } + require_zero(&payload[28..32], "commit reserved bytes must be zero")?; + let commit = Self { + transaction_id: read_u64(payload, 0)?, + new_generation: read_u64(payload, 8)?, + database_length: read_u64(payload, 16)?, + database_header_crc32: read_u32(payload, 24)?, + }; + if commit.transaction_id == 0 || commit.new_generation == 0 { + return Err(invalid_data("invalid commit identity or generation")); + } + if commit.database_length < PAGE_SIZE_U64 || !commit.database_length.is_multiple_of(PAGE_SIZE_U64) { + return Err(invalid_data("invalid committed database length")); + } + Ok(commit) + } +} + +fn prepare_record(encoded: &mut Vec, kind: u8, payload_length: usize) -> io::Result<()> { + let payload_length_u32 = u32::try_from(payload_length).map_err(|_| invalid_input("WAL payload exceeds u32"))?; + let record_length = RECORD_HEADER_LEN + .checked_add(payload_length) + .ok_or_else(|| invalid_input("WAL record length overflow"))?; + encoded.clear(); + encoded.try_reserve_exact(record_length).map_err(out_of_memory("WAL record allocation failed"))?; + encoded.resize(record_length, 0); + encoded[0] = kind; + encoded[4..8].copy_from_slice(&payload_length_u32.to_le_bytes()); + Ok(()) +} + +fn finish_record(encoded: &mut [u8]) -> io::Result<()> { + let checksum = crc_with_zeroed_field(encoded, 8)?; + encoded[8..12].copy_from_slice(&checksum.to_le_bytes()); + Ok(()) +} + +#[derive(Debug)] +struct ParsedWal { + header: WalHeader, + before_images: Vec, + commit: Option, + torn_tail: bool, +} + +fn read_wal(path: &Path) -> io::Result { + let mut file = File::open(path)?; + let file_len = file.metadata()?.len(); + if file_len < WAL_HEADER_LEN_U64 { + return Err(invalid_data("active WAL has a truncated header")); + } + let mut header_bytes = [0u8; WAL_HEADER_LEN]; + file.read_exact(&mut header_bytes)?; + let header = WalHeader::decode(&header_bytes)?; + let original_page_count = header.original_database_len / PAGE_SIZE_U64; + let mut before_images = Vec::new(); + let mut page_ids = HashSet::new(); + let mut commit = None; + let mut offset = WAL_HEADER_LEN_U64; + let mut torn_tail = false; + + while offset < file_len { + if commit.is_some() { + return Err(invalid_data("WAL record appears after commit")); + } + let remaining = file_len - offset; + if remaining < RECORD_HEADER_LEN_U64 { + torn_tail = true; + break; + } + let mut record_header = [0u8; RECORD_HEADER_LEN]; + file.read_exact(&mut record_header)?; + let kind = record_header[0]; + if record_header[1] != 0 { + return Err(invalid_data("unknown WAL record flags")); + } + require_zero(&record_header[2..4], "WAL record reserved bytes must be zero")?; + require_zero(&record_header[12..16], "WAL record reserved bytes must be zero")?; + let payload_length = usize::try_from(read_u32(&record_header, 4)?) + .map_err(|_| invalid_data("WAL payload length exceeds usize"))?; + let expected_length = match kind { + 1 => BEFORE_IMAGE_PAYLOAD_LEN, + 2 => COMMIT_PAYLOAD_LEN, + _ => return Err(invalid_data("unknown WAL record kind")), + }; + if payload_length != expected_length { + return Err(invalid_data("WAL record payload length mismatch")); + } + let total_length = RECORD_HEADER_LEN + .checked_add(payload_length) + .ok_or_else(|| invalid_data("WAL record length overflow"))?; + let total_length_u64 = u64::try_from(total_length).map_err(|_| invalid_data("WAL record exceeds u64"))?; + if remaining < total_length_u64 { + torn_tail = true; + break; + } + let mut payload = vec![0u8; payload_length]; + file.read_exact(&mut payload)?; + let mut hasher = crc32fast::Hasher::new(); + hasher.update(&record_header[..8]); + hasher.update(&[0; 4]); + hasher.update(&record_header[12..]); + hasher.update(&payload); + if read_u32(&record_header, 8)? != hasher.finalize() { + return Err(invalid_data("WAL record checksum mismatch")); + } + + match kind { + 1 => { + let image = BeforeImage::decode_payload(&payload)?; + if image.page_id >= original_page_count { + return Err(invalid_data("WAL before-image page id exceeds original database")); + } + page_ids.try_reserve(1).map_err(out_of_memory("WAL page-id set allocation failed"))?; + if !page_ids.insert(image.page_id) { + return Err(invalid_data("duplicate WAL before-image page id")); + } + before_images.try_reserve(1).map_err(out_of_memory("WAL before-image allocation failed"))?; + before_images.push(image); + } + 2 => { + let record = CommitRecord::decode_payload(&payload)?; + if record.transaction_id != header.transaction_id { + return Err(invalid_data("commit transaction id does not match WAL header")); + } + if header.original_generation.checked_add(1) != Some(record.new_generation) { + return Err(invalid_data("commit generation must advance WAL generation exactly once")); + } + commit = Some(record); + } + _ => return Err(invalid_data("unknown WAL record kind")), + } + offset = offset + .checked_add(total_length_u64) + .ok_or_else(|| invalid_data("WAL record offset overflow"))?; + } + Ok(ParsedWal { header, before_images, commit, torn_tail }) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum CommitBoundary { + WalTempSynced, + WalActivated, + BeforeImagesSynced, + DatabaseWritten, + DatabaseSynced, + CommitAppended, + CommitSynced, + WalCleared, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RecoveryBoundary { + PageRestored(u64), + HeaderRestored, + DatabaseTruncated, + DatabaseSynced, + WalRemoved, + ParentDirectorySynced, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum WalOutcome { + RecoveryPending, + CommittedCleanupPending, +} + +#[derive(Debug)] +pub(crate) struct WalOperationError { + outcome: WalOutcome, + database: PathBuf, + wal: PathBuf, + transaction_id: u64, + phase: &'static str, + cause: io::Error, +} + +impl WalOperationError { + pub(crate) fn outcome(&self) -> WalOutcome { self.outcome } + + #[cfg(test)] + pub(crate) fn database_path(&self) -> &Path { &self.database } + + #[cfg(test)] + pub(crate) fn wal_path(&self) -> &Path { &self.wal } + + #[cfg(test)] + pub(crate) fn transaction_id(&self) -> u64 { self.transaction_id } + + #[cfg(test)] + pub(crate) fn phase(&self) -> &'static str { self.phase } +} + +impl fmt::Display for WalOperationError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "WAL operation for {} failed during {} with {:?}: WAL {} transaction {}: {}", + self.database.display(), + self.phase, + self.outcome, + self.wal.display(), + self.transaction_id, + self.cause + ) + } +} + +impl Error for WalOperationError { + fn source(&self) -> Option<&(dyn Error + 'static)> { Some(&self.cause) } +} + +fn wal_operation_error<'a>( + outcome: WalOutcome, + database: &'a Path, + wal: &'a Path, + transaction_id: u64, + phase: &'static str, +) -> impl FnOnce(io::Error) -> io::Error + 'a { + move |cause| { + let kind = cause.kind(); + io::Error::new( + kind, + WalOperationError { + outcome, + database: database.to_path_buf(), + wal: wal.to_path_buf(), + transaction_id, + phase, + cause, + }, + ) + } +} + +#[derive(Clone, Copy)] +struct WalErrorContext<'a> { + outcome: WalOutcome, + database: &'a Path, + wal: &'a Path, + transaction_id: u64, +} + +impl WalErrorContext<'_> { + fn wrap(&self, phase: &'static str, result: io::Result) -> io::Result { + result.map_err(wal_operation_error( + self.outcome, + self.database, + self.wal, + self.transaction_id, + phase, + )) + } +} + +#[derive(Debug)] +pub(crate) struct WalReadError { + database: PathBuf, + wal: PathBuf, + wal_database_id: Option<[u8; 16]>, + current_database_id: Option<[u8; 16]>, + cause: io::Error, +} + +impl WalReadError { + #[cfg(test)] + pub(crate) fn database_path(&self) -> &Path { &self.database } + + #[cfg(test)] + pub(crate) fn wal_path(&self) -> &Path { &self.wal } + + #[cfg(test)] + pub(crate) fn wal_database_id(&self) -> Option<[u8; 16]> { self.wal_database_id } + + #[cfg(test)] + pub(crate) fn current_database_id(&self) -> Option<[u8; 16]> { self.current_database_id } +} + +impl fmt::Display for WalReadError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "failed to parse WAL {} for database {}", + self.wal.display(), + self.database.display() + )?; + if let Some(identity) = self.wal_database_id { + write!(formatter, ", WAL database identity {identity:02x?}")?; + } else { + write!(formatter, ", WAL database identity unavailable")?; + } + if let Some(identity) = self.current_database_id { + write!(formatter, ", current database identity {identity:02x?}")?; + } else { + write!(formatter, ", current database identity unavailable")?; + } + write!(formatter, ": {}", self.cause) + } +} + +impl Error for WalReadError { + fn source(&self) -> Option<&(dyn Error + 'static)> { Some(&self.cause) } +} + +fn raw_database_id(path: &Path, offset: u64) -> Option<[u8; 16]> { + let mut file = File::open(path).ok()?; + file.seek(SeekFrom::Start(offset)).ok()?; + let mut identity = [0u8; 16]; + file.read_exact(&mut identity).ok()?; + Some(identity) +} + +fn wal_read_error(database: &Path, wal: &Path, cause: io::Error) -> io::Error { + let kind = cause.kind(); + io::Error::new( + kind, + WalReadError { + database: database.to_path_buf(), + wal: wal.to_path_buf(), + wal_database_id: raw_database_id(wal, 16), + current_database_id: raw_database_id(database, 48), + cause, + }, + ) +} + +fn page_offset(page_id: u64) -> io::Result { + page_id.checked_mul(PAGE_SIZE_U64).ok_or_else(|| invalid_data("database page offset overflow")) +} + +fn read_page(file: &mut File, page_id: u64) -> io::Result<[u8; PAGE_SIZE]> { + let mut page = [0u8; PAGE_SIZE]; + file.seek(SeekFrom::Start(page_offset(page_id)?))?; + file.read_exact(&mut page)?; + Ok(page) +} + +fn write_page(file: &mut File, page_id: u64, page: &[u8; PAGE_SIZE]) -> io::Result<()> { + file.seek(SeekFrom::Start(page_offset(page_id)?))?; + file.write_all(page) +} + +fn database_length(next_page_id: u64) -> io::Result { + next_page_id + .checked_mul(PAGE_SIZE_U64) + .ok_or_else(|| invalid_input("database length overflow")) +} + +fn database_header_crc32(page: &[u8; PAGE_SIZE]) -> io::Result { read_u32(page, 72) } + +fn new_transaction_id() -> io::Result { + let bytes = uuid::Uuid::new_v4().as_u128().to_le_bytes(); + let low: [u8; 8] = bytes[..8] + .try_into() + .map_err(|_| io::Error::other("UUID transaction id conversion failed"))?; + let transaction_id = u64::from_le_bytes(low); + Ok(if transaction_id == 0 { 1 } else { transaction_id }) +} + +#[cfg(unix)] +pub(super) fn sync_parent_directory(path: &Path) -> io::Result<()> { + let parent = path.parent().filter(|parent| !parent.as_os_str().is_empty()).unwrap_or_else(|| Path::new(".")); + File::open(parent)?.sync_all() +} + +#[cfg(not(unix))] +pub(super) fn sync_parent_directory(_path: &Path) -> io::Result<()> { + Ok(()) +} + +pub(crate) fn invalidate_sorted_index(database: &Path) -> io::Result<()> { + let index = crate::repository::storage::get_file_path_for_db_index(database); + if index == database { + return Ok(()); + } + match std::fs::remove_file(&index) { + Ok(()) => sync_parent_directory(&index), + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error), + } +} + +fn clear_wal(database: &Path, transaction_id: u64, outcome: WalOutcome) -> io::Result<()> { + clear_wal_with_hook(database, transaction_id, outcome, &mut |_| Ok(())) +} + +fn clear_wal_with_hook( + database: &Path, + transaction_id: u64, + outcome: WalOutcome, + hook: &mut H, +) -> io::Result<()> +where + H: FnMut(RecoveryBoundary) -> io::Result<()>, +{ + let active = wal_path(database); + if active.try_exists().map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "inspect active WAL before cleanup", + ))? { + std::fs::remove_file(&active).map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "remove active WAL", + ))?; + hook(RecoveryBoundary::WalRemoved).map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "WalRemoved hook", + ))?; + } + let temporary = wal_temporary_path(database); + if temporary.try_exists().map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "inspect temporary WAL before cleanup", + ))? { + std::fs::remove_file(temporary).map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "remove temporary WAL", + ))?; + } + sync_parent_directory(database).map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "sync parent directory after WAL removal", + ))?; + hook(RecoveryBoundary::ParentDirectorySynced).map_err(wal_operation_error( + outcome, + database, + &active, + transaction_id, + "ParentDirectorySynced hook", + )) +} + +#[cfg(test)] +fn ordered_prepared_pages( + prepared: &[(u64, [u8; PAGE_SIZE])], +) -> io::Result> { + let mut ordered = Vec::new(); + ordered.try_reserve_exact(prepared.len()).map_err(out_of_memory("prepared-page ordering allocation failed"))?; + ordered.extend(prepared.iter().map(|(page_id, page)| (*page_id, page))); + ordered.sort_unstable_by_key(|(page_id, _)| *page_id); + if ordered.windows(2).any(|pair| pair[0].0 == pair[1].0) { + return Err(invalid_input("duplicate prepared page id")); + } + Ok(ordered) +} + +fn validate_prepared_pages( + database: &mut File, + ordered: &[(u64, &[u8; PAGE_SIZE])], +) -> io::Result<(DatabaseHeader, DatabaseHeader, u64)> { + let original_length = database.metadata()?.len(); + if original_length < PAGE_SIZE_U64 || !original_length.is_multiple_of(PAGE_SIZE_U64) { + return Err(invalid_data("database length is not page aligned")); + } + let original_page = read_page(database, 0)?; + let original = DatabaseHeader::decode(&original_page)?; + if database_length(original.next_page_id)? != original_length { + return Err(invalid_data("database length does not match original header")); + } + let (_, final_header_page) = ordered + .first() + .copied() + .filter(|(page_id, _)| *page_id == 0) + .ok_or_else(|| invalid_input("prepared transaction must include database header page 0"))?; + let final_header = DatabaseHeader::decode(final_header_page)?; + if final_header.database_id != original.database_id { + return Err(invalid_input("prepared database identity differs from current database")); + } + if original.generation.checked_add(1) != Some(final_header.generation) { + return Err(invalid_input("prepared generation must advance exactly once")); + } + let final_length = database_length(final_header.next_page_id)?; + if final_length < original_length { + return Err(invalid_input("prepared transaction cannot shrink the database")); + } + for (page_id, page) in ordered { + if *page_id >= final_header.next_page_id { + return Err(invalid_input("prepared page id exceeds final database bounds")); + } + if *page_id == 0 { + DatabaseHeader::decode(*page)?; + } else { + SlottedPage::open(page.as_slice(), *page_id, final_header.next_page_id)?; + } + } + for page_id in original.next_page_id..final_header.next_page_id { + if ordered.binary_search_by_key(&page_id, |(candidate, _)| *candidate).is_err() { + return Err(invalid_input("prepared transaction omits an appended page")); + } + } + Ok((original, final_header, final_length)) +} + +#[cfg(test)] +pub(crate) fn commit_prepared_pages( + database: &Path, + prepared: &[(u64, [u8; PAGE_SIZE])], +) -> io::Result<()> { + with_exclusive_sidecar(database, || { + recover_pending_under_existing_lock(database)?; + commit_prepared_pages_under_existing_lock(database, prepared) + }) +} + +#[cfg(test)] +pub(crate) fn commit_prepared_pages_under_existing_lock( + database: &Path, + prepared: &[(u64, [u8; PAGE_SIZE])], +) -> io::Result<()> { + let ordered = ordered_prepared_pages(prepared)?; + commit_ordered_page_refs_with_hook_under_existing_lock(database, &ordered, |_| Ok(())) +} + +pub(crate) fn commit_ordered_page_refs_under_existing_lock( + database: &Path, + prepared: &[(u64, &[u8; PAGE_SIZE])], +) -> io::Result<()> { + if prepared.windows(2).any(|pair| pair[0].0 >= pair[1].0) { + return Err(invalid_input("prepared page ids must be strictly increasing")); + } + commit_ordered_page_refs_with_hook_under_existing_lock(database, prepared, |_| Ok(())) +} + +#[cfg(test)] +fn commit_prepared_pages_with_hook( + database: &Path, + prepared: &[(u64, [u8; PAGE_SIZE])], + hook: impl FnMut(CommitBoundary) -> io::Result<()>, +) -> io::Result<()> { + with_exclusive_sidecar(database, || { + recover_pending_under_existing_lock(database)?; + commit_prepared_pages_with_hook_under_existing_lock(database, prepared, hook) + }) +} + +#[cfg(test)] +fn commit_prepared_pages_with_hook_under_existing_lock( + database_path: &Path, + prepared: &[(u64, [u8; PAGE_SIZE])], + hook: impl FnMut(CommitBoundary) -> io::Result<()>, +) -> io::Result<()> { + let ordered = ordered_prepared_pages(prepared)?; + commit_ordered_page_refs_with_hook_under_existing_lock(database_path, &ordered, hook) +} + +fn commit_ordered_page_refs_with_hook_under_existing_lock( + database_path: &Path, + ordered: &[(u64, &[u8; PAGE_SIZE])], + mut hook: impl FnMut(CommitBoundary) -> io::Result<()>, +) -> io::Result<()> { + if ordered.is_empty() { + return Err(invalid_input("prepared transaction is empty")); + } + let active_path = wal_path(database_path); + let temporary_path = wal_temporary_path(database_path); + if active_path.try_exists()? || temporary_path.try_exists()? { + return Err(invalid_data("pending WAL must be recovered before commit")); + } + let mut database = OpenOptions::new().read(true).write(true).open(database_path)?; + let (original_header, final_header, final_length) = validate_prepared_pages(&mut database, ordered)?; + let original_length = database.metadata()?.len(); + let original_page_count = original_length / PAGE_SIZE_U64; + let mut before_images = Vec::new(); + before_images.try_reserve(ordered.len()).map_err(out_of_memory("before-image allocation failed"))?; + for (page_id, _) in ordered { + if *page_id < original_page_count { + let page = read_page(&mut database, *page_id)?; + if *page_id == 0 { + DatabaseHeader::decode(&page)?; + } else { + SlottedPage::open(page.as_slice(), *page_id, original_header.next_page_id)?; + } + before_images.push(BeforeImage { page_id: *page_id, page }); + } + } + let transaction_id = new_transaction_id()?; + let wal_header = WalHeader { + database_id: original_header.database_id, + transaction_id, + original_database_len: original_length, + original_generation: original_header.generation, + }; + let mut temporary = OpenOptions::new().write(true).create_new(true).open(&temporary_path)?; + temporary.write_all(&wal_header.encode()?)?; + temporary.sync_all()?; + hook(CommitBoundary::WalTempSynced)?; + drop(temporary); + + std::fs::rename(&temporary_path, &active_path)?; + let mut context = WalErrorContext { + outcome: WalOutcome::RecoveryPending, + database: database_path, + wal: &active_path, + transaction_id, + }; + context.wrap("sync parent directory after WAL activation", sync_parent_directory(&active_path))?; + context.wrap("WalActivated hook", hook(CommitBoundary::WalActivated))?; + + let mut wal = context.wrap("open activated WAL", OpenOptions::new().append(true).open(&active_path))?; + let mut record = Vec::new(); + let reserve = record + .try_reserve_exact(RECORD_HEADER_LEN + BEFORE_IMAGE_PAYLOAD_LEN) + .map_err(out_of_memory("WAL record allocation failed")); + context.wrap("allocate WAL record buffer", reserve)?; + for image in &before_images { + context.wrap("encode WAL before-image", image.encode_record_into(&mut record))?; + context.wrap("append WAL before-image", wal.write_all(&record))?; + } + context.wrap("sync WAL before-images", wal.sync_data())?; + context.wrap("BeforeImagesSynced hook", hook(CommitBoundary::BeforeImagesSynced))?; + + context.wrap("resize database", database.set_len(final_length))?; + for (page_id, page) in ordered { + context.wrap("write database page", write_page(&mut database, *page_id, page))?; + } + context.wrap("DatabaseWritten hook", hook(CommitBoundary::DatabaseWritten))?; + context.wrap("sync database", database.sync_all())?; + context.wrap("DatabaseSynced hook", hook(CommitBoundary::DatabaseSynced))?; + + let final_header_page = ordered + .first() + .map(|(_, page)| page) + .ok_or_else(|| invalid_input("prepared header page is missing")); + let final_header_page = context.wrap("locate prepared database header", final_header_page)?; + let commit = CommitRecord { + transaction_id, + new_generation: final_header.generation, + database_length: final_length, + database_header_crc32: context.wrap( + "read prepared database header checksum", + database_header_crc32(final_header_page), + )?, + }; + context.wrap("encode WAL commit record", commit.encode_record_into(&mut record))?; + context.wrap("append WAL commit record", wal.write_all(&record))?; + context.wrap("CommitAppended hook", hook(CommitBoundary::CommitAppended))?; + context.wrap("sync WAL commit record", wal.sync_data())?; + context.outcome = WalOutcome::CommittedCleanupPending; + context.wrap("CommitSynced hook", hook(CommitBoundary::CommitSynced))?; + drop(wal); + + context.wrap("invalidate sorted index after commit", invalidate_sorted_index(database_path))?; + + clear_wal(database_path, transaction_id, context.outcome)?; + context.wrap("WalCleared hook", hook(CommitBoundary::WalCleared)) +} + +#[derive(Debug)] +pub(crate) struct RecoveryRequired { + database: PathBuf, + pending: PathBuf, + cause: Option, +} + +impl fmt::Display for RecoveryRequired { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "recovery required for {} because {} exists", + self.database.display(), + self.pending.display() + )?; + if let Some(cause) = &self.cause { + write!(formatter, ": {cause}")?; + } + Ok(()) + } +} + +impl Error for RecoveryRequired { + fn source(&self) -> Option<&(dyn Error + 'static)> { self.cause.as_ref().map(|cause| cause as &(dyn Error + 'static)) } +} + +fn append_suffix(path: &Path, suffix: &str) -> PathBuf { + let mut name = OsString::from(path.as_os_str()); + name.push(suffix); + PathBuf::from(name) +} + +pub(crate) fn wal_path(database: &Path) -> PathBuf { append_suffix(database, ".wal") } + +pub(crate) fn wal_temporary_path(database: &Path) -> PathBuf { append_suffix(database, ".wal.tmp") } + +pub(crate) fn recovery_required(database: &Path, cause: io::Error) -> io::Error { + let active = wal_path(database); + let pending = if active.exists() { active } else { wal_temporary_path(database) }; + io::Error::other(RecoveryRequired { database: database.to_path_buf(), pending, cause: Some(cause) }) +} + +fn open_sidecar(database: &Path) -> io::Result { + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(sidecar_lock_path(database)) +} + +pub(crate) struct SharedSidecarGuard { + _file: File, +} + +impl SharedSidecarGuard { + pub(crate) fn acquire(database: &Path) -> io::Result { + let file = open_sidecar(database)?; + file.lock_shared()?; + Ok(Self { _file: file }) + } +} + +pub(crate) struct ExclusiveSidecarGuard { + _file: File, +} + +impl ExclusiveSidecarGuard { + pub(crate) fn acquire(database: &Path) -> io::Result { + let file = open_sidecar(database)?; + file.lock_exclusive()?; + Ok(Self { _file: file }) + } +} + +pub(crate) fn with_exclusive_sidecar(database: &Path, operation: impl FnOnce() -> io::Result) -> io::Result { + let _guard = ExclusiveSidecarGuard::acquire(database)?; + operation() +} + +fn validate_recovery_before_images( + parsed: &ParsedWal, + current_header_page: &[u8; PAGE_SIZE], +) -> io::Result { + let original_header = match parsed.before_images.iter().find(|image| image.page_id == 0) { + Some(image) => DatabaseHeader::decode(&image.page)?, + None if parsed.before_images.is_empty() => DatabaseHeader::decode(current_header_page)?, + None => return Err(invalid_data("WAL before-images omit database header page 0")), + }; + if original_header.database_id != parsed.header.database_id { + return Err(invalid_data("WAL before-image database identity does not match WAL header")); + } + if original_header.generation != parsed.header.original_generation { + return Err(invalid_data("WAL before-image generation does not match WAL header")); + } + if database_length(original_header.next_page_id)? != parsed.header.original_database_len { + return Err(invalid_data("WAL original length does not match before-image header")); + } + let original_page_count = parsed.header.original_database_len / PAGE_SIZE_U64; + for image in &parsed.before_images { + if image.page_id >= original_page_count { + return Err(invalid_data("WAL before-image page id exceeds original database")); + } + if image.page_id != 0 { + SlottedPage::open(image.page.as_slice(), image.page_id, original_header.next_page_id)?; + } + } + Ok(original_header) +} + +fn validate_current_identity( + database_path: &Path, + current_header_page: &[u8; PAGE_SIZE], + expected: [u8; 16], +) -> io::Result<()> { + let actual = bytes_at::<16>(current_header_page, 48)?; + if actual == expected { + Ok(()) + } else { + Err(invalid_data(format!( + "WAL recovery for {} refused: WAL database identity {expected:02x?} does not match current database identity {actual:02x?}", + database_path.display() + ))) + } +} + +fn validate_committed_database( + database: &mut File, + parsed: &ParsedWal, + commit: CommitRecord, + current_header_page: &[u8; PAGE_SIZE], +) -> io::Result<()> { + if parsed.torn_tail { + return Err(invalid_data("committed WAL has a torn trailing record")); + } + let current_length = database.metadata()?.len(); + if current_length != commit.database_length { + return Err(invalid_data("committed WAL database length mismatch")); + } + let current_header = DatabaseHeader::decode(current_header_page)?; + if current_header.database_id != parsed.header.database_id { + return Err(invalid_data("committed WAL database identity mismatch")); + } + if current_header.generation != commit.new_generation { + return Err(invalid_data("committed WAL database generation mismatch")); + } + if database_length(current_header.next_page_id)? != commit.database_length { + return Err(invalid_data("committed WAL header length mismatch")); + } + if database_header_crc32(current_header_page)? != commit.database_header_crc32 { + return Err(invalid_data("committed WAL database header checksum mismatch")); + } + Ok(()) +} + +fn rollback_uncommitted( + database: &mut File, + database_path: &Path, + parsed: &ParsedWal, + hook: &mut H, +) -> io::Result<()> +where + H: FnMut(RecoveryBoundary) -> io::Result<()>, +{ + let active_path = wal_path(database_path); + let outcome = WalOutcome::RecoveryPending; + let transaction_id = parsed.header.transaction_id; + for image in parsed.before_images.iter().filter(|image| image.page_id != 0) { + write_page(database, image.page_id, &image.page).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "restore database page", + ))?; + hook(RecoveryBoundary::PageRestored(image.page_id)).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "PageRestored hook", + ))?; + } + if let Some(header) = parsed.before_images.iter().find(|image| image.page_id == 0) { + write_page(database, 0, &header.page).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "restore database header", + ))?; + hook(RecoveryBoundary::HeaderRestored).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "HeaderRestored hook", + ))?; + } + database.set_len(parsed.header.original_database_len).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "truncate database after rollback", + ))?; + hook(RecoveryBoundary::DatabaseTruncated).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "DatabaseTruncated hook", + ))?; + database.sync_all().map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "sync rolled-back database", + ))?; + hook(RecoveryBoundary::DatabaseSynced).map_err(wal_operation_error( + outcome, + database_path, + &active_path, + transaction_id, + "DatabaseSynced recovery hook", + ))?; + clear_wal_with_hook(database_path, transaction_id, outcome, hook)?; + info!( + "rolled back uncommitted WAL database={} transaction_id={}", + database_path.display(), + parsed.header.transaction_id + ); + Ok(()) +} + +pub(crate) fn recover_pending(database: &Path) -> io::Result<()> { + with_exclusive_sidecar(database, || recover_pending_under_existing_lock(database)) +} + +pub(crate) fn recover_pending_under_existing_lock(database_path: &Path) -> io::Result<()> { + recover_pending_with_hook_under_existing_lock(database_path, |_| Ok(())) +} + +#[cfg(test)] +fn recover_pending_with_hook( + database: &Path, + hook: impl FnMut(RecoveryBoundary) -> io::Result<()>, +) -> io::Result<()> { + with_exclusive_sidecar(database, || recover_pending_with_hook_under_existing_lock(database, hook)) +} + +fn recover_pending_with_hook_under_existing_lock(database_path: &Path, mut hook: H) -> io::Result<()> +where + H: FnMut(RecoveryBoundary) -> io::Result<()>, +{ + let active_path = wal_path(database_path); + let temporary_path = wal_temporary_path(database_path); + if !active_path.try_exists()? { + if temporary_path.try_exists()? { + std::fs::remove_file(temporary_path)?; + sync_parent_directory(database_path)?; + } + return Ok(()); + } + + let parsed = read_wal(&active_path).map_err(|cause| wal_read_error(database_path, &active_path, cause))?; + let mut database = OpenOptions::new().read(true).write(true).open(database_path)?; + let current_length = database.metadata()?.len(); + if current_length < PAGE_SIZE_U64 || current_length < parsed.header.original_database_len { + return Err(invalid_data("database is shorter than WAL recovery requires")); + } + let current_header_page = read_page(&mut database, 0)?; + validate_current_identity(database_path, ¤t_header_page, parsed.header.database_id)?; + let original_header = validate_recovery_before_images(&parsed, ¤t_header_page)?; + + if let Some(commit) = parsed.commit { + validate_committed_database(&mut database, &parsed, commit, ¤t_header_page)?; + invalidate_sorted_index(database_path).map_err(wal_operation_error( + WalOutcome::CommittedCleanupPending, + database_path, + &active_path, + parsed.header.transaction_id, + "invalidate sorted index during committed recovery", + ))?; + clear_wal_with_hook( + database_path, + parsed.header.transaction_id, + WalOutcome::CommittedCleanupPending, + &mut hook, + )?; + info!( + "cleared committed WAL database={} transaction_id={}", + database_path.display(), + parsed.header.transaction_id + ); + Ok(()) + } else { + if parsed.before_images.is_empty() + && (current_length != parsed.header.original_database_len + || original_header.generation != parsed.header.original_generation) + { + return Err(invalid_data("header-only WAL does not match unchanged database")); + } + rollback_uncommitted(&mut database, database_path, &parsed, &mut hook) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::repository::bplustree::v3::{ + format::{write_page_checksum, DatabaseHeader, PAGE_SIZE}, + page::encode_free_page, + tree::{BPlusTree, BPlusTreeQuery}, + }; + use std::{ + fs, + io::{Read, Seek, SeekFrom, Write}, + }; + + const DATABASE_ID: [u8; 16] = [ + 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, + ]; + const TRANSACTION_ID: u64 = 0x0102_0304_0506_0708; + + fn checksum_with_zeroed_field(bytes: &[u8], offset: usize) -> io::Result { + let end = offset.checked_add(4).ok_or_else(|| io::Error::other("checksum offset overflow"))?; + let before = bytes.get(..offset).ok_or_else(|| io::Error::other("missing checksum prefix"))?; + let after = bytes.get(end..).ok_or_else(|| io::Error::other("missing checksum suffix"))?; + let mut hasher = crc32fast::Hasher::new(); + hasher.update(before); + hasher.update(&[0; 4]); + hasher.update(after); + Ok(hasher.finalize()) + } + + fn read_u32(bytes: &[u8], offset: usize) -> io::Result { + let end = offset.checked_add(4).ok_or_else(|| io::Error::other("u32 offset overflow"))?; + let encoded: [u8; 4] = bytes + .get(offset..end) + .ok_or_else(|| io::Error::other("truncated u32"))? + .try_into() + .map_err(io::Error::other)?; + Ok(u32::from_le_bytes(encoded)) + } + + fn sample_header() -> WalHeader { + WalHeader { + database_id: DATABASE_ID, + transaction_id: TRANSACTION_ID, + original_database_len: 3 * PAGE_SIZE as u64, + original_generation: 9, + } + } + + fn sample_before_image() -> BeforeImage { + let mut page = [0u8; PAGE_SIZE]; + for (index, byte) in page.iter_mut().enumerate() { + *byte = index.to_le_bytes()[0]; + } + BeforeImage { page_id: 2, page } + } + + fn sample_commit() -> CommitRecord { + CommitRecord { + transaction_id: TRANSACTION_ID, + new_generation: 10, + database_length: 4 * PAGE_SIZE as u64, + database_header_crc32: 0x4433_2211, + } + } + + fn create_database(name: &str) -> io::Result<(tempfile::TempDir, PathBuf, DatabaseHeader, [u8; PAGE_SIZE])> { + let directory = tempfile::tempdir()?; + let path = directory.path().join(name); + let mut tree = BPlusTree::::new(); + tree.insert(1, String::from("original")); + tree.store(&path)?; + let mut original_page = [0u8; PAGE_SIZE]; + File::open(&path)?.read_exact(&mut original_page)?; + let original_header = DatabaseHeader::decode(&original_page)?; + Ok((directory, path, original_header, original_page)) + } + + fn prepared_header_page(original: &DatabaseHeader) -> io::Result<[u8; PAGE_SIZE]> { + let mut updated = original.clone(); + updated.generation = updated + .generation + .checked_add(1) + .ok_or_else(|| io::Error::other("test generation overflow"))?; + updated.encode() + } + + fn database_header(path: &Path) -> io::Result { + let mut page = [0u8; PAGE_SIZE]; + File::open(path)?.read_exact(&mut page)?; + DatabaseHeader::decode(&page) + } + + fn fail_at(target: CommitBoundary) -> impl FnMut(CommitBoundary) -> io::Result<()> { + move |actual| { + if actual == target { + Err(io::Error::other(format!("fault at {actual:?}"))) + } else { + Ok(()) + } + } + } + + fn fail_recovery_at(target: RecoveryBoundary) -> impl FnMut(RecoveryBoundary) -> io::Result<()> { + move |actual| { + if actual == target { + Err(io::Error::other(format!("recovery fault at {actual:?}"))) + } else { + Ok(()) + } + } + } + + fn create_multi_page_database( + name: &str, + ) -> io::Result<(tempfile::TempDir, PathBuf, DatabaseHeader, Vec)> { + let directory = tempfile::tempdir()?; + let path = directory.path().join(name); + let mut tree = BPlusTree::::new(); + for key in 0..160 { + tree.insert(key, format!("value-{key:03}-{}", "x".repeat(96))); + } + tree.store(&path)?; + let bytes = fs::read(&path)?; + let header_page: [u8; PAGE_SIZE] = bytes + .get(..PAGE_SIZE) + .ok_or_else(|| io::Error::other("multi-page database lacks a header"))? + .try_into() + .map_err(io::Error::other)?; + let header = DatabaseHeader::decode(&header_page)?; + if header.next_page_id < 3 { + return Err(io::Error::other("multi-page fixture did not create two data pages")); + } + Ok((directory, path, header, bytes)) + } + + fn mutate_fixture_page( + mut page: [u8; PAGE_SIZE], + page_id: u64, + next_page_id: u64, + marker: u8, + ) -> io::Result<[u8; PAGE_SIZE]> { + let header = SlottedPage::open(page.as_slice(), page_id, next_page_id)?.header(); + let free_start = usize::from(header.free_start); + if free_start >= usize::from(header.free_end) { + return Err(io::Error::other("fixture data page has no unused byte")); + } + let byte = page + .get_mut(free_start) + .ok_or_else(|| io::Error::other("fixture free-space offset is outside the page"))?; + *byte ^= marker; + write_page_checksum(&mut page)?; + SlottedPage::open(page.as_slice(), page_id, next_page_id)?; + Ok(page) + } + + fn create_uncommitted_multi_page_wal( + name: &str, + ) -> io::Result<(tempfile::TempDir, PathBuf, Vec, Vec)> { + let (directory, path, original, original_bytes) = create_multi_page_database(name)?; + let appended_page_id = original.next_page_id; + let mut updated = original.clone(); + updated.generation = updated + .generation + .checked_add(1) + .ok_or_else(|| io::Error::other("test generation overflow"))?; + updated.next_page_id = updated + .next_page_id + .checked_add(1) + .ok_or_else(|| io::Error::other("test page-id overflow"))?; + updated.free_page_head = appended_page_id; + + let prepared_len = usize::try_from(updated.next_page_id).map_err(io::Error::other)?; + let mut prepared = Vec::with_capacity(prepared_len); + prepared.push((0, updated.encode()?)); + for page_id in 1..original.next_page_id { + let start = usize::try_from(page_id) + .map_err(io::Error::other)? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| io::Error::other("test page offset overflow"))?; + let end = start + .checked_add(PAGE_SIZE) + .ok_or_else(|| io::Error::other("test page end overflow"))?; + let page: [u8; PAGE_SIZE] = original_bytes + .get(start..end) + .ok_or_else(|| io::Error::other("multi-page fixture is truncated"))? + .try_into() + .map_err(io::Error::other)?; + let page = match page_id { + 1 => mutate_fixture_page(page, page_id, updated.next_page_id, 0x5a)?, + 2 => mutate_fixture_page(page, page_id, updated.next_page_id, 0xa5)?, + _ => page, + }; + prepared.push((page_id, page)); + } + prepared.push(( + appended_page_id, + encode_free_page(appended_page_id, updated.next_page_id, 0)?, + )); + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::DatabaseWritten)); + let active = wal_path(&path); + if !active.try_exists()? { + return Err(io::Error::other("uncommitted multi-page WAL was not activated")); + } + Ok((directory, path, original_bytes, fs::read(active)?)) + } + + fn fixture_page(bytes: &[u8], page_id: u64) -> io::Result<&[u8]> { + let start = usize::try_from(page_id) + .map_err(io::Error::other)? + .checked_mul(PAGE_SIZE) + .ok_or_else(|| io::Error::other("fixture page offset overflow"))?; + let end = start + .checked_add(PAGE_SIZE) + .ok_or_else(|| io::Error::other("fixture page end overflow"))?; + bytes.get(start..end).ok_or_else(|| io::Error::other("fixture page is truncated")) + } + + fn assert_wal_outcome(error: &io::Error, expected: WalOutcome, path: &Path) -> io::Result<()> { + let operation = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .ok_or_else(|| io::Error::other("error lacked typed WAL outcome"))?; + assert_eq!(operation.outcome(), expected); + assert_eq!(operation.database_path(), path); + assert_eq!(operation.wal_path(), wal_path(path)); + assert_ne!(operation.transaction_id(), 0); + assert!(!operation.phase().is_empty()); + assert!(operation.source().is_some()); + Ok(()) + } + + fn write_active_wal(path: &Path, header: &WalHeader, records: &[Vec]) -> io::Result<()> { + let mut file = File::create(wal_path(path))?; + file.write_all(&header.encode()?)?; + for record in records { + file.write_all(record)?; + } + file.sync_all() + } + + fn rebuild_record_crc(record: &mut [u8]) -> io::Result<()> { + let payload_length = usize::try_from(read_u32(record, 4)?).map_err(io::Error::other)?; + let record_length = 16usize + .checked_add(payload_length) + .ok_or_else(|| io::Error::other("record length overflow"))?; + let bytes = record + .get(..record_length) + .ok_or_else(|| io::Error::other("truncated record fixture"))?; + let checksum = checksum_with_zeroed_field(bytes, 8)?; + record + .get_mut(8..12) + .ok_or_else(|| io::Error::other("missing record checksum"))? + .copy_from_slice(&checksum.to_le_bytes()); + Ok(()) + } + + #[test] + fn wal_header_golden_layout_and_round_trip() -> io::Result<()> { + let header = sample_header(); + let encoded = header.encode()?; + assert_eq!(encoded.len(), 64); + assert_eq!(&encoded[0..4], b"BTW3"); + assert_eq!(&encoded[4..8], &1u32.to_le_bytes()); + assert_eq!(&encoded[8..12], &64u32.to_le_bytes()); + assert_eq!(&encoded[12..16], &4096u32.to_le_bytes()); + assert_eq!(&encoded[16..32], &DATABASE_ID); + assert_eq!(&encoded[32..40], &TRANSACTION_ID.to_le_bytes()); + assert_eq!(&encoded[40..48], &(3 * PAGE_SIZE as u64).to_le_bytes()); + assert_eq!(&encoded[48..56], &9u64.to_le_bytes()); + assert_eq!(read_u32(&encoded, 56)?, 0x3ad1_eb4d); + assert_eq!(read_u32(&encoded, 56)?, checksum_with_zeroed_field(&encoded, 56)?); + assert_eq!(&encoded[60..64], &[0; 4]); + assert_eq!(WalHeader::decode(&encoded)?, header); + Ok(()) + } + + #[test] + fn before_image_golden_layout_and_round_trip() -> io::Result<()> { + let image = sample_before_image(); + let encoded = image.encode_record()?; + assert_eq!(encoded.len(), 16 + 4112); + assert_eq!(encoded[0], 1); + assert_eq!(encoded[1], 0); + assert_eq!(&encoded[2..4], &[0; 2]); + assert_eq!(&encoded[4..8], &4112u32.to_le_bytes()); + assert_eq!(&encoded[12..16], &[0; 4]); + assert_eq!(&encoded[16..24], &2u64.to_le_bytes()); + assert_eq!(read_u32(&encoded, 24)?, 0xa291_2082); + assert_eq!(read_u32(&encoded, 24)?, crc32fast::hash(&image.page)); + assert_eq!(&encoded[28..32], &[0; 4]); + assert_eq!(&encoded[32..], &image.page); + assert_eq!(read_u32(&encoded, 8)?, 0x9ea7_c8fc); + assert_eq!(read_u32(&encoded, 8)?, checksum_with_zeroed_field(&encoded, 8)?); + assert_eq!(BeforeImage::decode_payload(&encoded[16..])?, image); + Ok(()) + } + + #[test] + fn commit_record_golden_layout_and_round_trip() -> io::Result<()> { + let commit = sample_commit(); + let encoded = commit.encode_record()?; + assert_eq!(encoded.len(), 48); + assert_eq!(encoded[0], 2); + assert_eq!(encoded[1], 0); + assert_eq!(&encoded[2..4], &[0; 2]); + assert_eq!(&encoded[4..8], &32u32.to_le_bytes()); + assert_eq!(&encoded[12..16], &[0; 4]); + assert_eq!(&encoded[16..24], &TRANSACTION_ID.to_le_bytes()); + assert_eq!(&encoded[24..32], &10u64.to_le_bytes()); + assert_eq!(&encoded[32..40], &(4 * PAGE_SIZE as u64).to_le_bytes()); + assert_eq!(&encoded[40..44], &0x4433_2211u32.to_le_bytes()); + assert_eq!(&encoded[44..48], &[0; 4]); + assert_eq!(read_u32(&encoded, 8)?, 0xb9d2_319c); + assert_eq!(read_u32(&encoded, 8)?, checksum_with_zeroed_field(&encoded, 8)?); + assert_eq!(CommitRecord::decode_payload(&encoded[16..])?, commit); + Ok(()) + } + + #[test] + fn rejects_invalid_header_fields_and_checksum() -> io::Result<()> { + for offset in [0usize, 4, 8, 12, 16, 32, 40, 48, 56, 60] { + let mut encoded = sample_header().encode()?; + let byte = encoded.get_mut(offset).ok_or_else(|| io::Error::other("bad fixture offset"))?; + *byte ^= 0x5a; + assert!(WalHeader::decode(&encoded).is_err(), "offset {offset} was accepted"); + } + Ok(()) + } + + #[test] + fn rejects_invalid_record_fields_and_duplicates() -> io::Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("invalid-record.db"); + let header = sample_header(); + let before = sample_before_image().encode_record()?; + let commit = sample_commit().encode_record()?; + + for offset in [0usize, 1, 2, 4, 8, 12] { + let mut invalid = before.clone(); + let byte = invalid.get_mut(offset).ok_or_else(|| io::Error::other("bad record fixture offset"))?; + *byte ^= 0x7f; + write_active_wal(&path, &header, &[invalid])?; + assert!(read_wal(&wal_path(&path)).is_err(), "record offset {offset} was accepted"); + } + + write_active_wal(&path, &header, &[before.clone(), before.clone()])?; + assert!(read_wal(&wal_path(&path)).is_err()); + + let mut lower = sample_before_image(); + lower.page_id = 1; + write_active_wal(&path, &header, &[before.clone(), lower.encode_record()?])?; + assert_eq!(read_wal(&wal_path(&path))?.before_images.len(), 2); + + write_active_wal(&path, &header, &[commit.clone(), before.clone()])?; + assert!(read_wal(&wal_path(&path)).is_err()); + write_active_wal(&path, &header, &[commit.clone(), commit])?; + assert!(read_wal(&wal_path(&path)).is_err()); + + let mut skipped_generation = sample_commit(); + skipped_generation.new_generation += 1; + write_active_wal(&path, &header, &[skipped_generation.encode_record()?])?; + assert!(read_wal(&wal_path(&path)).is_err()); + + let mut bad_page_crc = before; + bad_page_crc[24] ^= 1; + rebuild_record_crc(&mut bad_page_crc)?; + write_active_wal(&path, &header, &[bad_page_crc])?; + assert!(read_wal(&wal_path(&path)).is_err()); + + let mut out_of_bounds = sample_before_image(); + out_of_bounds.page_id = 3; + write_active_wal(&path, &header, &[out_of_bounds.encode_record()?])?; + assert!(read_wal(&wal_path(&path)).is_err()); + Ok(()) + } + + #[test] + fn torn_final_record_rolls_back_complete_before_images() -> io::Result<()> { + let (_directory, path, original, _) = create_database("torn-tail.db")?; + let updated_page = prepared_header_page(&original)?; + let prepared = [(0, updated_page)]; + let error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::DatabaseWritten)) + .err() + .ok_or_else(|| io::Error::other("fault did not interrupt commit"))?; + assert!(error.to_string().contains("DatabaseWritten")); + let mut wal = OpenOptions::new().append(true).open(wal_path(&path))?; + wal.write_all(&[2, 0, 0, 0, 32])?; + wal.sync_all()?; + drop(wal); + + assert_eq!(database_header(&path)?.generation, original.generation + 1); + recover_pending(&path)?; + assert_eq!(database_header(&path)?.generation, original.generation); + assert!(!wal_path(&path).try_exists()?); + recover_pending(&path)?; + Ok(()) + } + + #[test] + fn every_commit_boundary_recovers_idempotently() -> io::Result<()> { + for boundary in [ + CommitBoundary::WalTempSynced, + CommitBoundary::WalActivated, + CommitBoundary::BeforeImagesSynced, + CommitBoundary::DatabaseWritten, + CommitBoundary::DatabaseSynced, + CommitBoundary::CommitAppended, + CommitBoundary::CommitSynced, + CommitBoundary::WalCleared, + ] { + let (_directory, path, original, _) = create_database(&format!("boundary-{boundary:?}.db"))?; + let updated_page = prepared_header_page(&original)?; + let prepared = [(0, updated_page)]; + let error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(boundary)) + .err() + .ok_or_else(|| io::Error::other(format!("{boundary:?} did not interrupt commit")))?; + assert!(error.to_string().contains(&format!("{boundary:?}"))); + + recover_pending(&path)?; + recover_pending(&path)?; + let generation = database_header(&path)?.generation; + let committed = matches!( + boundary, + CommitBoundary::CommitAppended | CommitBoundary::CommitSynced | CommitBoundary::WalCleared + ); + assert_eq!(generation, original.generation + u64::from(committed), "boundary {boundary:?}"); + assert!(!wal_path(&path).try_exists()?); + assert!(!wal_temporary_path(&path).try_exists()?); + } + Ok(()) + } + + #[test] + fn committed_recovery_invalidates_sorted_index_before_clearing_wal() -> io::Result<()> { + let (_directory, path, original, _) = create_database("committed-index-recovery.db")?; + let index_path = crate::repository::storage::get_file_path_for_db_index(&path); + fs::write(&index_path, b"stale sorted index")?; + let prepared = [(0, prepared_header_page(&original)?)]; + + let error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::CommitSynced)) + .err() + .ok_or_else(|| io::Error::other("CommitSynced did not interrupt commit"))?; + assert_wal_outcome(&error, WalOutcome::CommittedCleanupPending, &path)?; + assert!(index_path.try_exists()?); + assert!(wal_path(&path).try_exists()?); + + recover_pending(&path)?; + assert_eq!(database_header(&path)?.generation, original.generation + 1); + assert!(!index_path.try_exists()?); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn late_commit_failures_report_the_durable_wal_outcome() -> io::Result<()> { + for (boundary, expected) in [ + (CommitBoundary::WalActivated, WalOutcome::RecoveryPending), + (CommitBoundary::CommitAppended, WalOutcome::RecoveryPending), + (CommitBoundary::CommitSynced, WalOutcome::CommittedCleanupPending), + (CommitBoundary::WalCleared, WalOutcome::CommittedCleanupPending), + ] { + let (_directory, path, original, _) = create_database(&format!("outcome-{boundary:?}.db"))?; + let prepared = [(0, prepared_header_page(&original)?)]; + let error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(boundary)) + .err() + .ok_or_else(|| io::Error::other(format!("{boundary:?} did not interrupt commit")))?; + let operation = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .ok_or_else(|| io::Error::other(format!("{boundary:?} lacked typed WAL outcome")))?; + assert_eq!(operation.outcome(), expected); + assert_eq!(operation.database_path(), path); + assert_eq!(operation.wal_path(), wal_path(&path)); + assert_ne!(operation.transaction_id(), 0); + assert!(!operation.phase().is_empty()); + assert!(operation.source().is_some()); + } + + let (_directory, path, original, _) = create_database("pre-rename-error.db")?; + let prepared = [(0, prepared_header_page(&original)?)]; + let error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::WalTempSynced)) + .err() + .ok_or_else(|| io::Error::other("pre-rename fault did not interrupt commit"))?; + assert!(error.get_ref().and_then(|source| source.downcast_ref::()).is_none()); + Ok(()) + } + + #[test] + fn committed_mismatch_leaves_database_and_wal_untouched() -> io::Result<()> { + for mismatch in ["length", "identity", "generation", "header-page-crc", "commit-header-crc", "transaction"] { + let (_directory, path, original, _) = create_database(&format!("mismatch-{mismatch}.db"))?; + let updated_page = prepared_header_page(&original)?; + let prepared = [(0, updated_page)]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::CommitSynced)); + + match mismatch { + "length" => { + let file = OpenOptions::new().write(true).open(&path)?; + file.set_len(file.metadata()?.len() + PAGE_SIZE as u64)?; + } + "identity" => { + let mut file = OpenOptions::new().read(true).write(true).open(&path)?; + file.seek(SeekFrom::Start(48))?; + file.write_all(&[0x77; 16])?; + file.sync_all()?; + } + "generation" => { + let mut page = [0u8; PAGE_SIZE]; + let mut file = OpenOptions::new().read(true).write(true).open(&path)?; + file.read_exact(&mut page)?; + let mut header = DatabaseHeader::decode(&page)?; + header.generation += 1; + file.seek(SeekFrom::Start(0))?; + file.write_all(&header.encode()?)?; + file.sync_all()?; + } + "header-page-crc" => { + let mut file = OpenOptions::new().read(true).write(true).open(&path)?; + file.seek(SeekFrom::Start(72))?; + file.write_all(&[0; 4])?; + file.sync_all()?; + } + "commit-header-crc" => { + let wal_path = wal_path(&path); + let mut bytes = fs::read(&wal_path)?; + let record_offset = 64 + 16 + 4112; + let crc = bytes + .get_mut(record_offset + 40..record_offset + 44) + .ok_or_else(|| io::Error::other("commit header CRC fixture is truncated"))?; + crc.copy_from_slice(&0x1122_3344u32.to_le_bytes()); + rebuild_record_crc( + bytes + .get_mut(record_offset..) + .ok_or_else(|| io::Error::other("commit record fixture is missing"))?, + )?; + fs::write(wal_path, bytes)?; + } + "transaction" => { + let wal_path = wal_path(&path); + let mut bytes = fs::read(&wal_path)?; + let record_offset = 64 + 16 + 4112; + let tx = bytes + .get_mut(record_offset + 16..record_offset + 24) + .ok_or_else(|| io::Error::other("commit transaction fixture is truncated"))?; + tx.copy_from_slice(&0x9999u64.to_le_bytes()); + rebuild_record_crc( + bytes + .get_mut(record_offset..) + .ok_or_else(|| io::Error::other("commit record fixture is missing"))?, + )?; + fs::write(wal_path, bytes)?; + } + _ => return Err(io::Error::other("unknown mismatch fixture")), + } + let database_before = fs::read(&path)?; + let wal_before = fs::read(wal_path(&path))?; + assert!(recover_pending(&path).is_err(), "{mismatch} was accepted"); + assert_eq!(fs::read(&path)?, database_before, "database changed for {mismatch}"); + assert_eq!(fs::read(wal_path(&path))?, wal_before, "WAL changed for {mismatch}"); + } + Ok(()) + } + + #[test] + fn abandoned_wal_temp_is_removed_without_touching_database() -> io::Result<()> { + let (_directory, path, _, _) = create_database("abandoned-temp.db")?; + let database_before = fs::read(&path)?; + fs::write(wal_temporary_path(&path), b"not activated")?; + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, database_before); + assert!(!wal_temporary_path(&path).try_exists()?); + Ok(()) + } + + #[test] + fn foreign_uncommitted_wal_leaves_database_and_wal_untouched() -> io::Result<()> { + let (_directory, path, original, _) = create_database("foreign-wal.db")?; + let prepared = [(0, prepared_header_page(&original)?)]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::BeforeImagesSynced)); + let active = wal_path(&path); + let mut bytes = fs::read(&active)?; + bytes + .get_mut(16..32) + .ok_or_else(|| io::Error::other("WAL header identity fixture is truncated"))? + .copy_from_slice(&[0x88; 16]); + let checksum = checksum_with_zeroed_field( + bytes.get(..64).ok_or_else(|| io::Error::other("WAL header fixture is truncated"))?, + 56, + )?; + bytes + .get_mut(56..60) + .ok_or_else(|| io::Error::other("WAL header checksum fixture is truncated"))? + .copy_from_slice(&checksum.to_le_bytes()); + fs::write(&active, bytes)?; + let database_before = fs::read(&path)?; + let wal_before = fs::read(&active)?; + + let error = recover_pending(&path) + .err() + .ok_or_else(|| io::Error::other("foreign WAL was accepted"))?; + let message = error.to_string(); + assert!(message.contains(&path.display().to_string())); + assert!(message.contains(&format!("{:02x?}", [0x88; 16]))); + assert!(message.contains(&format!("{:02x?}", original.database_id))); + assert_eq!(fs::read(&path)?, database_before); + assert_eq!(fs::read(&active)?, wal_before); + Ok(()) + } + + #[test] + fn corrupt_record_reports_recovery_paths_identities_and_source_without_mutation() -> io::Result<()> { + let (_directory, path, original, _) = create_database("corrupt-record-context.db")?; + let prepared = [(0, prepared_header_page(&original)?)]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::BeforeImagesSynced)); + let active = wal_path(&path); + let mut bytes = fs::read(&active)?; + let checksum_byte = bytes + .get_mut(WAL_HEADER_LEN + 8) + .ok_or_else(|| io::Error::other("WAL record fixture is truncated"))?; + *checksum_byte ^= 1; + fs::write(&active, bytes)?; + let database_before = fs::read(&path)?; + let wal_before = fs::read(&active)?; + + let error = recover_pending(&path) + .err() + .ok_or_else(|| io::Error::other("corrupt WAL record was accepted"))?; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + let context = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .ok_or_else(|| io::Error::other("WAL parser error lacked recovery context"))?; + assert_eq!(context.database_path(), path); + assert_eq!(context.wal_path(), active); + assert_eq!(context.wal_database_id(), Some(original.database_id)); + assert_eq!(context.current_database_id(), Some(original.database_id)); + assert!(context.source().and_then(|source| source.downcast_ref::()).is_some()); + let message = error.to_string(); + assert!(message.contains(&path.display().to_string())); + assert!(message.contains(&active.display().to_string())); + assert!(message.contains(&format!("WAL database identity {:02x?}", original.database_id))); + assert!(message.contains(&format!("current database identity {:02x?}", original.database_id))); + assert_eq!(fs::read(&path)?, database_before); + assert_eq!(fs::read(&active)?, wal_before); + Ok(()) + } + + #[test] + fn rollback_faults_leave_recovery_idempotent_at_every_database_boundary() -> io::Result<()> { + for boundary in [ + RecoveryBoundary::PageRestored(1), + RecoveryBoundary::HeaderRestored, + RecoveryBoundary::DatabaseTruncated, + RecoveryBoundary::DatabaseSynced, + ] { + let (_directory, path, original_bytes, _) = + create_uncommitted_multi_page_wal(&format!("rollback-{boundary:?}.db"))?; + let modified_bytes = fs::read(&path)?; + assert_ne!(fixture_page(&modified_bytes, 1)?, fixture_page(&original_bytes, 1)?); + assert_ne!(fixture_page(&modified_bytes, 2)?, fixture_page(&original_bytes, 2)?); + let error = recover_pending_with_hook(&path, fail_recovery_at(boundary)) + .err() + .ok_or_else(|| io::Error::other(format!("{boundary:?} did not interrupt recovery")))?; + assert_wal_outcome(&error, WalOutcome::RecoveryPending, &path)?; + assert!(wal_path(&path).try_exists()?); + if boundary == RecoveryBoundary::PageRestored(1) { + let partially_restored = fs::read(&path)?; + assert_eq!(fixture_page(&partially_restored, 1)?, fixture_page(&original_bytes, 1)?); + assert_eq!(fixture_page(&partially_restored, 2)?, fixture_page(&modified_bytes, 2)?); + } + + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, original_bytes, "second recovery failed after {boundary:?}"); + assert!(!wal_path(&path).try_exists()?); + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, original_bytes, "third recovery changed data after {boundary:?}"); + } + Ok(()) + } + + #[test] + fn rollback_cleanup_faults_report_recovery_pending_and_allow_reappeared_wal() -> io::Result<()> { + for boundary in [RecoveryBoundary::WalRemoved, RecoveryBoundary::ParentDirectorySynced] { + let (_directory, path, original_bytes, wal_bytes) = + create_uncommitted_multi_page_wal(&format!("rollback-cleanup-{boundary:?}.db"))?; + let error = recover_pending_with_hook(&path, fail_recovery_at(boundary)) + .err() + .ok_or_else(|| io::Error::other(format!("{boundary:?} did not interrupt rollback cleanup")))?; + assert_wal_outcome(&error, WalOutcome::RecoveryPending, &path)?; + assert_eq!(fs::read(&path)?, original_bytes); + assert!(!wal_path(&path).try_exists()?); + + fs::write(wal_path(&path), &wal_bytes)?; + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, original_bytes); + assert!(!wal_path(&path).try_exists()?); + recover_pending(&path)?; + } + Ok(()) + } + + #[test] + fn committed_cleanup_faults_preserve_commit_and_allow_reappeared_wal() -> io::Result<()> { + for boundary in [RecoveryBoundary::WalRemoved, RecoveryBoundary::ParentDirectorySynced] { + let (_directory, path, original, _) = create_database(&format!("commit-cleanup-{boundary:?}.db"))?; + let prepared = [(0, prepared_header_page(&original)?)]; + let commit_error = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::CommitSynced)) + .err() + .ok_or_else(|| io::Error::other("CommitSynced did not preserve committed WAL"))?; + assert_wal_outcome(&commit_error, WalOutcome::CommittedCleanupPending, &path)?; + let committed_bytes = fs::read(&path)?; + let active = wal_path(&path); + let wal_bytes = fs::read(&active)?; + + let error = recover_pending_with_hook(&path, fail_recovery_at(boundary)) + .err() + .ok_or_else(|| io::Error::other(format!("{boundary:?} did not interrupt committed cleanup")))?; + assert_wal_outcome(&error, WalOutcome::CommittedCleanupPending, &path)?; + assert_eq!(fs::read(&path)?, committed_bytes); + assert!(!active.try_exists()?); + + fs::write(&active, &wal_bytes)?; + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, committed_bytes); + assert!(!active.try_exists()?); + recover_pending(&path)?; + } + Ok(()) + } + + #[test] + fn rollback_truncates_new_pages_and_second_recovery_is_noop() -> io::Result<()> { + let (_directory, path, original, _) = create_database("truncate-appended.db")?; + let original_length = fs::metadata(&path)?.len(); + let appended_page_id = original.next_page_id; + let mut updated = original.clone(); + updated.generation += 1; + updated.next_page_id += 1; + updated.free_page_head = appended_page_id; + let prepared = [ + (0, updated.encode()?), + (appended_page_id, encode_free_page(appended_page_id, updated.next_page_id, 0)?), + ]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::DatabaseWritten)); + assert_eq!(fs::metadata(&path)?.len(), original_length + PAGE_SIZE as u64); + + recover_pending(&path)?; + assert_eq!(fs::metadata(&path)?.len(), original_length); + let after_first = fs::read(&path)?; + recover_pending(&path)?; + assert_eq!(fs::read(&path)?, after_first); + Ok(()) + } + + #[test] + fn query_open_recovers_uncommitted_wal_before_mapping() -> io::Result<()> { + let (_directory, path, original, _) = create_database("query-recovers.db")?; + let prepared = [(0, prepared_header_page(&original)?)]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::DatabaseWritten)); + assert_eq!(database_header(&path)?.generation, original.generation + 1); + + let mut query = BPlusTreeQuery::::try_new(&path)?; + assert_eq!(query.query(&1).map_err(crate::repository::bplustree::common::BPlusTreeError::to_io)?, Some(String::from("original"))); + assert_eq!(database_header(&path)?.generation, original.generation); + assert!(!wal_path(&path).try_exists()?); + Ok(()) + } + + #[cfg(unix)] + #[test] + fn query_reports_recovery_required_when_database_is_not_writable() -> io::Result<()> { + use std::os::unix::fs::PermissionsExt as _; + + let (_directory, path, original, _) = create_database("readonly-recovery.db")?; + let prepared = [(0, prepared_header_page(&original)?)]; + let _ = commit_prepared_pages_with_hook(&path, &prepared, fail_at(CommitBoundary::DatabaseWritten)); + let original_permissions = fs::metadata(&path)?.permissions(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o444))?; + let result = BPlusTreeQuery::::try_new(&path); + fs::set_permissions(&path, original_permissions)?; + + let error = result.err().ok_or_else(|| io::Error::other("query recovered read-only database"))?; + assert!(error.get_ref().and_then(|source| source.downcast_ref::()).is_some()); + assert!(wal_path(&path).try_exists()?); + recover_pending(&path)?; + Ok(()) + } + + #[test] + fn commit_rejects_corrupt_original_page_before_creating_wal() -> io::Result<()> { + let (_directory, path, original, _) = create_database("corrupt-original.db")?; + let mut valid_leaf = [0u8; PAGE_SIZE]; + let mut file = OpenOptions::new().read(true).write(true).open(&path)?; + file.seek(SeekFrom::Start(PAGE_SIZE_U64))?; + file.read_exact(&mut valid_leaf)?; + file.seek(SeekFrom::Start(PAGE_SIZE_U64 + 24))?; + file.write_all(&[valid_leaf[24] ^ 1])?; + file.sync_all()?; + drop(file); + let database_before = fs::read(&path)?; + let prepared = [(0, prepared_header_page(&original)?), (1, valid_leaf)]; + + assert!(commit_prepared_pages(&path, &prepared).is_err()); + assert_eq!(fs::read(&path)?, database_before); + assert!(!wal_path(&path).try_exists()?); + assert!(!wal_temporary_path(&path).try_exists()?); + Ok(()) + } +} diff --git a/backend/src/repository/epg_repository.rs b/backend/src/repository/epg_repository.rs index e9ce9c608..536021525 100644 --- a/backend/src/repository/epg_repository.rs +++ b/backend/src/repository/epg_repository.rs @@ -468,7 +468,13 @@ mod tests { .expect("case-preserving EPG should be written"); let mut query = BPlusTreeQuery::, EpgChannel>::try_new(&path).expect("EPG DB should open"); - let stored_ids = query.iter().map(|(_, channel)| channel.id).collect::>(); + let stored_ids = query + .iter() + .collect::>>() + .expect("EPG entries should be readable") + .into_iter() + .map(|(_, channel)| channel.id) + .collect::>(); assert_eq!(stored_ids.iter().map(AsRef::as_ref).collect::>(), vec!["Z.Channel", "a.channel"]); } diff --git a/backend/src/repository/library_repository.rs b/backend/src/repository/library_repository.rs index 9895c265b..acd7ce817 100644 --- a/backend/src/repository/library_repository.rs +++ b/backend/src/repository/library_repository.rs @@ -32,10 +32,11 @@ pub async fn persist_input_library_playlist( // Keep previously probed technical metadata for unchanged local files. let mut existing_by_uuid: HashMap = HashMap::new(); if library_path.exists() { - if let Ok(mut query) = BPlusTreeQuery::::try_new(&library_path) { - for (uuid, item) in query.iter() { - existing_by_uuid.insert(uuid, item); - } + let mut query = BPlusTreeQuery::::try_new(&library_path) + .map_err(|error| repository_read_error(&library_path, error))?; + for entry in query.iter() { + let (uuid, item) = entry.map_err(|error| repository_read_error(&library_path, error))?; + existing_by_uuid.insert(uuid, item); } } @@ -111,6 +112,10 @@ fn preserve_local_probe_state_if_unchanged(new_item: &mut XtreamPlaylistItem, ol } } +fn repository_read_error(path: &Path, error: impl std::fmt::Display) -> TuliproxError { + TuliproxError::Repository(format!("failed to read B+Tree database {}: {error}", path.display())) +} + pub async fn load_input_local_library_playlist(app_config: &Arc, lib_path: &Path) -> Result, TuliproxError> { if file_exists_async(lib_path).await { let file_lock = app_config.file_locks.read_lock(lib_path).await; @@ -120,23 +125,24 @@ pub async fn load_input_local_library_playlist(app_config: &Arc, lib_ let groups = task::spawn_blocking(move || -> Result, TuliproxError> { let _guard = file_lock; let mut groups: IndexMap = IndexMap::new(); - if let Ok(mut query) = BPlusTreeQuery::::try_new(&lib_path) { - let mut group_cnt = 0; - for (_, ref item) in query.iter() { - let cluster = XtreamCluster::try_from(item.item_type).unwrap_or(XtreamCluster::Live); - let key = (cluster, item.group.clone()); - groups.entry(key) - .or_insert_with(|| { - group_cnt += 1; - PlaylistGroup { - id: group_cnt, - title: item.group.clone(), - channels: Vec::new(), - xtream_cluster: cluster, - } - }) - .channels.push(PlaylistItem::from(item)); - } + let mut query = BPlusTreeQuery::::try_new(&lib_path) + .map_err(|error| repository_read_error(&lib_path, error))?; + let mut group_cnt = 0; + for entry in query.iter() { + let (_, item) = entry.map_err(|error| repository_read_error(&lib_path, error))?; + let cluster = XtreamCluster::try_from(item.item_type).unwrap_or(XtreamCluster::Live); + let key = (cluster, item.group.clone()); + groups.entry(key) + .or_insert_with(|| { + group_cnt += 1; + PlaylistGroup { + id: group_cnt, + title: item.group.clone(), + channels: Vec::new(), + xtream_cluster: cluster, + } + }) + .channels.push(PlaylistItem::from(&item)); } Ok(groups.into_values().collect()) }) @@ -155,6 +161,12 @@ mod tests { use shared::utils::Internable; use shared::model::{EpisodeStreamProperties, VideoStreamDetailProperties, VideoStreamProperties}; + #[test] + fn bplustree_read_errors_are_generic_repository_errors() { + let error = repository_read_error(Path::new("library.db"), std::io::Error::other("corrupt page")); + assert!(matches!(error, TuliproxError::Repository(message) if message.contains("library.db") && message.contains("corrupt page"))); + } + fn video_item( url: &str, added: &str, diff --git a/backend/src/repository/m3u_playlist_iterator.rs b/backend/src/repository/m3u_playlist_iterator.rs index a5aa1dff2..9f8143367 100644 --- a/backend/src/repository/m3u_playlist_iterator.rs +++ b/backend/src/repository/m3u_playlist_iterator.rs @@ -23,7 +23,7 @@ use tokio::task; create_bitset!(u8, M3uPlaylistIteratorFlags, MaskRedirectUrl, IncludeTypeInUrl, RewriteResource); pub struct M3uPlaylistIterator { - inner: LockedReceiverStream<(M3uPlaylistItem, bool)>, + inner: LockedReceiverStream>, } struct UrlRewriteContext<'a> { @@ -252,10 +252,11 @@ impl M3uPlaylistIterator { let m3u_path = m3u_path.clone(); let index_path = get_file_path_for_db_index(&m3u_path); - let (tx, rx) = mpsc::channel::<(M3uPlaylistItem, bool)>(256); + let (tx, rx) = mpsc::channel::>(256); let m3u_path_for_log = m3u_path.clone(); let index_path_for_log = index_path.clone(); + let join_error_tx = tx.clone(); let handle = task::spawn_blocking(move || { let _guard = bg_lock; let reader = match open_playlist_reader::( @@ -266,6 +267,7 @@ impl M3uPlaylistIterator { Ok(reader) => reader, Err(err) => { error!("Failed to open M3U playlist DB {}: {err}", m3u_path.display()); + let _ = tx.blocking_send(Err(err)); return; } }; @@ -275,7 +277,7 @@ impl M3uPlaylistIterator { let item = match entry { Ok((_, item)) => item, Err(err) => { - error!("Iterator error: {err}"); + error!("Skipping unreadable M3U playlist entry: {err}"); continue; } }; @@ -308,14 +310,14 @@ impl M3uPlaylistIterator { ); if let Some(prev) = pending.replace(item) { - if tx.blocking_send((prev, true)).is_err() { + if tx.blocking_send(Ok((prev, true))).is_err() { return; } } } if let Some(last) = pending { - let _ = tx.blocking_send((last, false)); + let _ = tx.blocking_send(Ok((last, false))); } }); tokio::spawn(async move { @@ -325,6 +327,12 @@ impl M3uPlaylistIterator { m3u_path_for_log.display(), index_path_for_log.display() ); + let _ = join_error_tx + .send(Err(TuliproxError::RepositoryM3u(format!( + "M3U playlist iterator task failed for {}: {err}", + m3u_path_for_log.display() + )))) + .await; } }); @@ -335,7 +343,7 @@ impl M3uPlaylistIterator { } impl Stream for M3uPlaylistIterator { - type Item = (M3uPlaylistItem, bool); + type Item = Result<(M3uPlaylistItem, bool), TuliproxError>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_next(cx) @@ -363,19 +371,20 @@ impl M3uPlaylistM3uTextIterator { } impl Stream for M3uPlaylistM3uTextIterator { - type Item = String; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if !self.started { self.started = true; - return Poll::Ready(Some("#EXTM3U".to_string())); + return Poll::Ready(Some(Ok("#EXTM3U".to_string()))); } match Pin::new(&mut self.inner).poll_next(cx) { - Poll::Ready(Some((m3u_pli, _has_next))) => { + Poll::Ready(Some(Ok((m3u_pli, _has_next)))) => { let target_options = self.target_options.as_ref(); - Poll::Ready(Some(m3u_pli.to_m3u(target_options, true))) + Poll::Ready(Some(Ok(m3u_pli.to_m3u(target_options, true)))) } + Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error.to_string()))), Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, } @@ -545,15 +554,28 @@ mod tests { // no item lines, which lets the assertion focus on the header line only. #[tokio::test] async fn m3u_text_iterator_emits_bare_extm3u_header_without_proxying_source_url_tvg() { - let (tx, rx) = mpsc::channel::<(M3uPlaylistItem, bool)>(1); + let (tx, rx) = mpsc::channel::>(1); drop(tx); let inner_iter = M3uPlaylistIterator { inner: LockedReceiverStream::new_empty(rx) }; let mut text_iter = M3uPlaylistM3uTextIterator { inner: inner_iter, started: false, target_options: None }; let first = text_iter.next().await; - assert_eq!(first.as_deref(), Some("#EXTM3U"), "EXTM3U header must be bare, never contain url-tvg"); + assert_eq!(first.as_ref().and_then(|line| line.as_deref().ok()), Some("#EXTM3U"), "EXTM3U header must be bare, never contain url-tvg"); let second = text_iter.next().await; assert_eq!(second, None, "inner channel is closed, so no item lines should follow"); } + + #[tokio::test] + async fn iterator_forwards_one_storage_error_then_ends() { + let (tx, rx) = mpsc::channel(2); + assert!(tx.send(Ok((m3u_item("http://example.test/live.ts"), true))).await.is_ok()); + assert!(tx.send(Err(shared::error::TuliproxError::RepositoryM3u("corrupt page".into()))).await.is_ok()); + drop(tx); + + let mut iterator = M3uPlaylistIterator { inner: LockedReceiverStream::new_empty(rx) }; + assert!(iterator.next().await.is_some_and(|entry| entry.is_ok())); + assert!(iterator.next().await.is_some_and(|entry| entry.is_err())); + assert!(iterator.next().await.is_none()); + } } diff --git a/backend/src/repository/m3u_repository.rs b/backend/src/repository/m3u_repository.rs index 392ea8873..8d6e43bf0 100644 --- a/backend/src/repository/m3u_repository.rs +++ b/backend/src/repository/m3u_repository.rs @@ -384,7 +384,7 @@ async fn iter_raw_m3u_playlist( ) -> Option> + Send + Unpin>> where ItemKey: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static, - SortKey: for<'de> Deserialize<'de> + Send + 'static, + SortKey: Ord + for<'de> Deserialize<'de> + Send + 'static, { // Two read locks: iter_lock is held by LockedReceiverStream for the consumer lifetime, // while bg_lock is moved into spawn_blocking to guard the on-disk reader. @@ -419,7 +419,7 @@ where let item = match entry { Ok((_, item)) => item, Err(err) => { - error!("M3U playlist reader error: {err}"); + error!("Skipping unreadable M3U playlist entry: {err}"); continue; } }; @@ -487,25 +487,26 @@ pub async fn load_input_m3u_playlist( let groups = task::spawn_blocking(move || -> Result, TuliproxError> { let _guard = file_lock; let mut groups: IndexMap = IndexMap::new(); - if let Ok(mut query) = BPlusTreeQuery::, M3uPlaylistItem>::try_new(&m3u_path) { - let mut group_cnt = 0; - for (_, ref item) in query.iter() { - let cluster = XtreamCluster::try_from(item.item_type).unwrap_or(XtreamCluster::Live); - let key = (cluster, item.group.clone()); - groups - .entry(key) - .or_insert_with(|| { - group_cnt += 1; - PlaylistGroup { - id: group_cnt, - title: item.group.clone(), - channels: Vec::new(), - xtream_cluster: cluster, - } - }) - .channels - .push(PlaylistItem::from(item)); - } + let mut query = BPlusTreeQuery::, M3uPlaylistItem>::try_new(&m3u_path) + .map_err(|error| TuliproxError::RepositoryM3u(error.to_string()))?; + let mut group_cnt = 0; + for entry in query.iter() { + let (_, item) = entry.map_err(|error| TuliproxError::RepositoryM3u(error.to_string()))?; + let cluster = XtreamCluster::try_from(item.item_type).unwrap_or(XtreamCluster::Live); + let key = (cluster, item.group.clone()); + groups + .entry(key) + .or_insert_with(|| { + group_cnt += 1; + PlaylistGroup { + id: group_cnt, + title: item.group.clone(), + channels: Vec::new(), + xtream_cluster: cluster, + } + }) + .channels + .push(PlaylistItem::from(&item)); } Ok(groups.into_values().collect()) }) diff --git a/backend/src/repository/mod.rs b/backend/src/repository/mod.rs index bee6e1631..1888553df 100644 --- a/backend/src/repository/mod.rs +++ b/backend/src/repository/mod.rs @@ -1,7 +1,6 @@ mod storage; mod target_id_mapping; mod bplustree; -mod bplustree_migration; mod playlist_repository; mod m3u_repository; mod xtream_repository; @@ -14,7 +13,6 @@ pub mod storage_const; mod playlist_scratch; mod playlist_source; mod library_repository; -mod sorted_index; mod alias_repository; mod playlist_stream; mod provider_dns_repository; @@ -24,7 +22,6 @@ mod qos_snapshot_repository; pub use storage::*; pub use target_id_mapping::*; pub use bplustree::*; -pub use bplustree_migration::*; pub use playlist_repository::*; pub use m3u_repository::*; pub use xtream_repository::*; @@ -34,7 +31,6 @@ pub use m3u_playlist_iterator::*; pub use xtream_playlist_iterator::*; pub use user_repository::*; pub use storage_const::*; -pub use sorted_index::*; pub use alias_repository::*; pub use playlist_source::*; pub use library_repository::*; @@ -42,6 +38,3 @@ pub use playlist_stream::*; pub use provider_dns_repository::*; pub use stream_history::*; pub use qos_snapshot_repository::*; - -#[cfg(test)] -mod bplustree_stress; diff --git a/backend/src/repository/playlist_repository.rs b/backend/src/repository/playlist_repository.rs index 1cac3dace..de1c3f45f 100644 --- a/backend/src/repository/playlist_repository.rs +++ b/backend/src/repository/playlist_repository.rs @@ -684,7 +684,7 @@ pub async fn load_input_playlist(ctx: &PlaylistProcessingContext, input: &Config let clusters_to_load = clusters.unwrap_or(&XTREAM_CLUSTER); if disk_based_processing { let source = PlaylistSource::xtream_disk( - XtreamDiskPlaylistSource::new(app_config, &storage_path).await, + XtreamDiskPlaylistSource::new(app_config, &storage_path).await?, ); Ok(PlaylistSource::filtered(source, skipped_clusters(clusters_to_load))) } else { @@ -697,7 +697,7 @@ pub async fn load_input_playlist(ctx: &PlaylistProcessingContext, input: &Config let file_path = get_input_m3u_playlist_file_path(&storage_path, &input.name); if disk_based_processing && file_path.exists() { Ok(PlaylistSource::m3u_disk( - M3uDiskPlaylistSource::new(app_config, &file_path).await, + M3uDiskPlaylistSource::new(app_config, &file_path).await?, )) } else { let groups = load_input_m3u_playlist(app_config, &file_path).await?; @@ -708,7 +708,7 @@ pub async fn load_input_playlist(ctx: &PlaylistProcessingContext, input: &Config let file_path = get_input_local_library_playlist_file_path(&storage_path, &input.name); if disk_based_processing && file_path.exists() { Ok(PlaylistSource::local_library_disk( - LocalLibraryDiskPlaylistSource::new(app_config, &file_path).await, + LocalLibraryDiskPlaylistSource::new(app_config, &file_path).await?, )) } else { let groups = load_input_local_library_playlist(app_config, &file_path).await?; @@ -719,7 +719,7 @@ pub async fn load_input_playlist(ctx: &PlaylistProcessingContext, input: &Config let file_path = get_input_media_server_playlist_file_path(&storage_path, &input.name); if disk_based_processing && file_path.exists() { Ok(PlaylistSource::media_server_disk( - MediaServerDiskPlaylistSource::new(app_config, &file_path).await, + MediaServerDiskPlaylistSource::new(app_config, &file_path).await?, )) } else { let groups = load_input_media_server_playlist(app_config, &file_path).await?; diff --git a/backend/src/repository/playlist_source.rs b/backend/src/repository/playlist_source.rs index 65af774a4..c6dc42aca 100644 --- a/backend/src/repository/playlist_source.rs +++ b/backend/src/repository/playlist_source.rs @@ -457,8 +457,18 @@ pub struct XtreamDiskPlaylistSource { series: Option, } +fn log_and_skip_btree_error(entry: std::io::Result) -> Option { + match entry { + Ok(entry) => Some(entry), + Err(error) => { + error!("Skipping unreadable B+Tree playlist entry; iteration continues when possible: {error}"); + None + } + } +} + impl XtreamDiskPlaylistSource { - pub(crate) async fn new(app_config: &Arc, storage_path: &Path) -> Self { + pub(crate) async fn new(app_config: &Arc, storage_path: &Path) -> Result { let mut source = XtreamDiskPlaylistSource { app_config: Arc::clone(app_config), storage_path: storage_path.to_path_buf(), @@ -466,27 +476,31 @@ impl XtreamDiskPlaylistSource { vod: None, series: None, }; - source.reload().await; - source + source.reload().await?; + Ok(source) } - async fn reload(&mut self) { + async fn reload(&mut self) -> Result<(), TuliproxError> { if self.live.is_none() { let live_path = xtream_get_file_path(&self.storage_path, XtreamCluster::Live); self.live = load_bplustree_query::(&self.app_config, &live_path).await + ? .map(|(query, guard)| (query, Arc::new(guard))); } if self.vod.is_none() { let vod_path = xtream_get_file_path(&self.storage_path, XtreamCluster::Video); self.vod = load_bplustree_query::(&self.app_config, &vod_path).await + ? .map(|(query, guard)| (query, Arc::new(guard))); } if self.series.is_none() { let series_path = xtream_get_file_path(&self.storage_path, XtreamCluster::Series); self.series = load_bplustree_query::(&self.app_config, &series_path).await + ? .map(|(query, guard)| (query, Arc::new(guard))); } + Ok(()) } } @@ -501,7 +515,11 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { fn get_group_count(&mut self) -> usize { fn collect_groups(query: &mut Option<(BPlusTreeQuery, Q)>, groups: &mut HashSet>) { - if let Some((query, _)) = query { for (_, item) in query.iter() { groups.insert(item.group.clone()); } } + if let Some((query, _)) = query { + for (_, item) in query.iter().filter_map(log_and_skip_btree_error) { + groups.insert(item.group.clone()); + } + } } let mut groups = HashSet::new(); @@ -541,7 +559,7 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { return; } if let Some((query, _)) = query { - for (_, item) in query.iter() { + for (_, item) in query.iter().filter_map(log_and_skip_btree_error) { groups.insert((cluster, Arc::clone(&item.group))); } } @@ -561,16 +579,16 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { } fn into_items(&mut self) -> Box + Send + '_> { - let live = self.live.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| PlaylistItem::from(&item)); - let vod = self.vod.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| PlaylistItem::from(&item)); - let series = self.series.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| PlaylistItem::from(&item)); + let live = self.live.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| PlaylistItem::from(&item)); + let vod = self.vod.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| PlaylistItem::from(&item)); + let series = self.series.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| PlaylistItem::from(&item)); Box::new(live.chain(vod).chain(series)) } fn items<'a>(&'a mut self) -> Box> + Send + 'a> { - let live = self.live.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); - let vod = self.vod.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); - let series = self.series.as_mut().into_iter().flat_map(|(q, _)| q.iter()).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); + let live = self.live.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); + let vod = self.vod.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); + let series = self.series.as_mut().into_iter().flat_map(|(q, _)| q.iter()).filter_map(log_and_skip_btree_error).map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); Box::new(live.chain(vod).chain(series)) } @@ -602,7 +620,7 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { fn get_missing_vod_info_count(&mut self) -> usize { self.vod.as_mut().map_or(0, |(query, _)| { - query.iter().filter(|(_, item)| { + query.iter().filter_map(log_and_skip_btree_error).filter(|(_, item)| { item.item_type == PlaylistItemType::Video && item.provider_id > 0 && !item.has_details() }).count() }) @@ -610,7 +628,7 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { fn get_missing_series_info_count(&mut self) -> usize { self.series.as_mut().map_or(0, |(query, _)| { - query.iter().filter(|(_, item)| { + query.iter().filter_map(log_and_skip_btree_error).filter(|(_, item)| { item.item_type == PlaylistItemType::SeriesInfo && item.provider_id > 0 && !item.has_details() }).count() }) @@ -641,13 +659,13 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { let mut groups_map: IndexMap<(XtreamCluster, u32), PlaylistGroup> = IndexMap::new(); let mut iters: Vec<(XtreamCluster, Box + Send>)> = vec![]; if let Some((q, _)) = self.live.as_mut() { - iters.push((XtreamCluster::Live, Box::new(q.iter().map(|(_, item)| item)))); + iters.push((XtreamCluster::Live, Box::new(q.iter().filter_map(log_and_skip_btree_error).map(|(_, item)| item)))); } if let Some((q, _)) = self.vod.as_mut() { - iters.push((XtreamCluster::Video, Box::new(q.iter().map(|(_, item)| item)))); + iters.push((XtreamCluster::Video, Box::new(q.iter().filter_map(log_and_skip_btree_error).map(|(_, item)| item)))); } if let Some((q, _)) = self.series.as_mut() { - iters.push((XtreamCluster::Series, Box::new(q.iter().map(|(_, item)| item)))); + iters.push((XtreamCluster::Series, Box::new(q.iter().filter_map(log_and_skip_btree_error).map(|(_, item)| item)))); } for (cluster, iter) in iters { @@ -700,7 +718,9 @@ impl PlaylistSourceOps for XtreamDiskPlaylistSource { fn obtain_resources(&mut self) -> BoxFuture<'_, ()> { Box::pin(async move { - self.reload().await; + if let Err(error) = self.reload().await { + error!("Failed to reload Xtream disk playlist source: {error}"); + } }) } fn sort_by_provider_ordinal(&mut self) { @@ -719,24 +739,26 @@ macro_rules! impl_single_file_disk_source { } impl [<$name DiskPlaylistSource>] { - pub(crate) async fn new(app_config: &Arc, file_path: &Path) -> Self { + pub(crate) async fn new(app_config: &Arc, file_path: &Path) -> Result { let mut source = Self { app_config: Arc::clone(app_config), file_path: file_path.to_path_buf(), playlist: None, guard: None, }; - source.reload().await; - source + source.reload().await?; + Ok(source) } - async fn reload(&mut self) { + async fn reload(&mut self) -> Result<(), TuliproxError> { self.guard = None; self.playlist = load_bplustree_query::<$key_type, $entry_type>(&self.app_config, &self.file_path).await + ? .map(|(query, guard)| { self.guard = Some(Arc::new(guard)); query }); + Ok(()) } } @@ -748,7 +770,11 @@ macro_rules! impl_single_file_disk_source { fn get_group_count(&mut self) -> usize { let mut groups = HashSet::new(); - if let Some(query) = self.playlist.as_mut() { for (_, item) in query.iter() { groups.insert(item.group.clone()); } } + if let Some(query) = self.playlist.as_mut() { + for (_, item) in query.iter().filter_map(log_and_skip_btree_error) { + groups.insert(item.group.clone()); + } + } groups.len() } @@ -756,6 +782,7 @@ macro_rules! impl_single_file_disk_source { self.playlist.as_mut().map_or(0, |query| { query .iter() + .filter_map(log_and_skip_btree_error) .filter(|(_, item)| !skip_set.contains(&cluster_from_item_type(item.item_type))) .count() }) @@ -764,7 +791,7 @@ macro_rules! impl_single_file_disk_source { fn get_group_count_excluding_clusters(&mut self, skip_set: &HashSet) -> usize { let mut groups = HashSet::<(XtreamCluster, Arc)>::new(); if let Some(query) = self.playlist.as_mut() { - for (_, item) in query.iter() { + for (_, item) in query.iter().filter_map(log_and_skip_btree_error) { let cluster = cluster_from_item_type(item.item_type); if !skip_set.contains(&cluster) { groups.insert((cluster, Arc::clone(&item.group))); @@ -778,7 +805,7 @@ macro_rules! impl_single_file_disk_source { fn into_items(&mut self) -> Box + Send + '_> { if let Some(q) = self.playlist.as_mut() { - Box::new(q.iter().map(|(_, item)| PlaylistItem::from(&item))) + Box::new(q.iter().filter_map(log_and_skip_btree_error).map(|(_, item)| PlaylistItem::from(&item))) } else { Box::new(std::iter::empty()) } @@ -786,7 +813,10 @@ macro_rules! impl_single_file_disk_source { fn items<'a>(&'a mut self) -> Box> + Send + 'a> { if let Some(pl) = self.playlist.as_mut() { - let iter = pl.iter().map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); + let iter = pl + .iter() + .filter_map(log_and_skip_btree_error) + .map(|(_, item)| Cow::Owned(PlaylistItem::from(&item))); Box::new(iter) } else { Box::new(std::iter::empty()) @@ -815,7 +845,7 @@ macro_rules! impl_single_file_disk_source { // Build groups on-the-fly using disk iterator (streams one leaf at a time) if let Some(q) = self.playlist.as_mut() { let mut groups_map: IndexMap<(XtreamCluster, Arc), PlaylistGroup> = IndexMap::new(); - for (_, item) in q.iter() { + for (_, item) in q.iter().filter_map(log_and_skip_btree_error) { let cluster = XtreamCluster::try_from(item.item_type).unwrap_or(XtreamCluster::Live); let normalized_group = shared::utils::deunicode_string(&item.group).to_lowercase().intern(); let key = (cluster, normalized_group); @@ -873,7 +903,9 @@ macro_rules! impl_single_file_disk_source { fn obtain_resources(&mut self) -> BoxFuture<'_, ()> { Box::pin(async move { - self.reload().await; + if let Err(error) = self.reload().await { + error!("Failed to reload {} disk playlist source: {error}", stringify!($name)); + } }) } @@ -1064,9 +1096,12 @@ impl PlaylistSourceOps for MemoryPlaylistSource { } } -async fn load_bplustree_query(app_config: &Arc, file_path: &Path) -> Option<(BPlusTreeQuery, FileReadGuard)> +async fn load_bplustree_query( + app_config: &Arc, + file_path: &Path, +) -> Result, FileReadGuard)>, TuliproxError> where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, + K: Ord + Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static, P: Serialize + for<'de> Deserialize<'de> + Clone + Send + 'static, { if file_path.exists() { @@ -1078,30 +1113,29 @@ where }) .await { - Ok(Ok((query, guard))) => Some((query, guard)), - Ok(Err(err)) => { - error!( - "I/O error loading disk playlist {} via BPlusTreeQuery::try_new: {err}", - file_path_err.display() - ); - None - } - Err(err) => { - error!( - "Task/JoinError loading disk playlist {}: {err} (panic={}, cancelled={})", - file_path_err.display(), - err.is_panic(), - err.is_cancelled() - ); - None - } + Ok(Ok((query, guard))) => Ok(Some((query, guard))), + Ok(Err(err)) => Err(TuliproxError::RepositoryPlaylist(format!( + "Failed to open disk playlist {}: {err}", + file_path_err.display() + ))), + Err(err) => Err(TuliproxError::RepositoryPlaylist(format!( + "Failed to load disk playlist {}: {err} (panic={}, cancelled={})", + file_path_err.display(), + err.is_panic(), + err.is_cancelled() + ))), } - } else { None } + } else { + Ok(None) + } } #[cfg(test)] mod tests { - use super::{MemoryPlaylistSource, PlaylistGroup, PlaylistItem, PlaylistSource, XtreamCluster, XtreamDiskPlaylistSource}; + use super::{ + load_bplustree_query, MemoryPlaylistSource, PlaylistGroup, PlaylistItem, PlaylistSource, XtreamCluster, + XtreamDiskPlaylistSource, + }; use crate::model::{AppConfig, Config, MediaToolCapabilities, SourcesConfig}; use crate::repository::BPlusTreeQuery; use crate::utils::FileLockManager; @@ -1146,6 +1180,32 @@ mod tests { }) } + #[tokio::test] + async fn invalid_disk_playlist_is_an_error_not_an_absent_source() { + let directory = tempfile::tempdir(); + assert!(directory.is_ok()); + let Some(directory) = directory.ok() else { return }; + let path = directory.path().join("invalid.db"); + assert!(std::fs::write(&path, b"not a btree").is_ok()); + + let result = load_bplustree_query::(&test_app_config(), &path).await; + assert!(matches!(result, Err(TuliproxError::RepositoryPlaylist(_)))); + } + + #[tokio::test] + async fn absent_disk_playlist_remains_absent() { + let directory = tempfile::tempdir(); + assert!(directory.is_ok()); + let Some(directory) = directory.ok() else { return }; + + let result = load_bplustree_query::( + &test_app_config(), + &directory.path().join("missing.db"), + ) + .await; + assert!(matches!(result, Ok(None))); + } + fn make_item(title: &str, group: &str, category_id: u32) -> PlaylistItem { PlaylistItem { header: PlaylistItemHeader { @@ -1238,7 +1298,7 @@ mod tests { result, Err(TuliproxError::RepositoryPlaylist(message)) if message.contains("Failed to clone live disk playlist query") - && message.contains("No data source available to clone") + && message.contains("mapped query without a path cannot be cloned") )); } diff --git a/backend/src/repository/playlist_stream.rs b/backend/src/repository/playlist_stream.rs index 1cb1d65bf..574dcd137 100644 --- a/backend/src/repository/playlist_stream.rs +++ b/backend/src/repository/playlist_stream.rs @@ -3,13 +3,56 @@ use futures::Stream; use log::error; use serde::{Deserialize, Serialize}; use shared::error::{ TuliproxError}; -use std::path::Path; +use std::{io, path::Path}; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; -use super::{BPlusTreeQuery, BPlusTreeSortedIteratorOwned, PlaylistIteratorReader, SortedIndexReader}; +use super::{BPlusTreeDiskIteratorOwned, BPlusTreeQuery}; + +pub(crate) enum PlaylistIteratorReader { + Sorted { + iterator: super::bplustree::sorted_index::v4::OwnedIterator, + fallback_path: std::path::PathBuf, + yielded: bool, + }, + Unsorted(BPlusTreeDiskIteratorOwned), +} + +impl Iterator for PlaylistIteratorReader +where + K: Ord + for<'de> Deserialize<'de>, + V: for<'de> Deserialize<'de>, + SortKey: Ord + for<'de> Deserialize<'de>, +{ + type Item = io::Result<(K, V)>; + + fn next(&mut self) -> Option { + let fallback = match self { + Self::Sorted { iterator, fallback_path, yielded } => match iterator.next() { + Some(Ok(entry)) => { + *yielded = true; + return Some(Ok(entry)); + } + Some(Err(error)) if !*yielded => Some((fallback_path.clone(), error)), + other => return other, + }, + Self::Unsorted(iterator) => return iterator.next(), + }; + let (path, index_error) = fallback?; + match BPlusTreeQuery::try_new(&path) { + Ok(query) => { + *self = Self::Unsorted(query.disk_iter()); + self.next() + } + Err(tree_error) => Some(Err(io::Error::new( + tree_error.kind(), + format!("sorted index failed before its first entry ({index_error}); tree fallback failed: {tree_error}"), + ))), + } + } +} /// Stream wrapper that holds a file read lock for the lifetime of the stream. pub struct LockedReceiverStream { @@ -51,7 +94,7 @@ impl Stream for LockedReceiverStream { /// Open a playlist reader with sorted-index fallback. /// /// NOTE: This performs disk I/O and should be used inside `spawn_blocking`. -pub fn open_playlist_reader( +pub(crate) fn open_playlist_reader( path: &Path, index_path: &Path, sorted_err_prefix: Option<&str>, @@ -59,7 +102,7 @@ pub fn open_playlist_reader( where K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, V: Serialize + for<'de> Deserialize<'de> + Clone, - SortKey: for<'de> Deserialize<'de>, + SortKey: Ord + for<'de> Deserialize<'de>, { let query = BPlusTreeQuery::::try_new(path) .map_err(|err| TuliproxError::Config(format!( @@ -68,11 +111,13 @@ where )))?; if index_path.exists() { - match SortedIndexReader::::open(index_path) { - Ok(index_reader) => { - let (filepath, file, mmap) = query.into_sorted_parts(); - let reader = BPlusTreeSortedIteratorOwned::from_index_reader(index_reader, filepath, file, mmap); - return Ok(PlaylistIteratorReader::Sorted(reader)); + match super::bplustree::sorted_index::v4::OwnedIterator::open(query, index_path) { + Ok(iterator) => { + return Ok(PlaylistIteratorReader::Sorted { + iterator, + fallback_path: path.to_path_buf(), + yielded: false, + }); } Err(err) => { if let Some(prefix) = sorted_err_prefix { @@ -82,5 +127,66 @@ where } } + let query = BPlusTreeQuery::::try_new(path).map_err(|err| { + TuliproxError::Config(format!("Could not reopen BPlusTreeQuery {} - {err}", path.display())) + })?; Ok(PlaylistIteratorReader::Unsorted(query.disk_iter())) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::repository::{ + bplustree::v3::{BPlusTree, Locator}, + sorted_index::v4, + }; + use std::{fs, io}; + + #[test] + fn corrupt_first_index_entry_falls_back_before_yielding() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let database = dir.path().join("playlist.db"); + let index = dir.path().join("playlist.idx"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("ccc")); + tree.insert(2u32, String::from("a")); + tree.store_with_index(&database, String::len)?; + let mut bytes = fs::read(&index)?; + *bytes.get_mut(72).ok_or_else(|| io::Error::other("index body missing"))? ^= 1; + fs::write(&index, bytes)?; + + let reader = open_playlist_reader::(&database, &index, None) + .map_err(io::Error::other)?; + assert_eq!(reader.collect::>>()?, vec![(1, String::from("ccc")), (2, String::from("a"))]); + Ok(()) + } + + #[test] + fn corrupt_late_index_entry_does_not_hide_following_entries() -> io::Result<()> { + let dir = tempfile::tempdir()?; + let database = dir.path().join("playlist-late.db"); + let index = dir.path().join("playlist-late.idx"); + let mut tree = BPlusTree::new(); + tree.insert(1u32, String::from("one")); + tree.insert(2u32, String::from("two")); + tree.insert(3u32, String::from("three")); + tree.store(&database)?; + let mut query = BPlusTreeQuery::::try_new(&database)?; + let entries = query.collect_with_locators()?; + let (database_id, generation) = query.snapshot_identity(); + drop(query); + let mut writer = v4::Writer::::new(&index, database_id, generation)?; + writer.push(&1, &entries[0].0, entries[0].2)?; + writer.push(&2, &entries[1].0, Locator { slot_index: u16::MAX, ..entries[1].2 })?; + writer.push(&3, &entries[2].0, entries[2].2)?; + writer.finish()?; + + let mut reader = open_playlist_reader::(&database, &index, None) + .map_err(io::Error::other)?; + assert_eq!(reader.next().transpose()?, Some((1, String::from("one")))); + assert!(reader.next().is_some_and(|entry| entry.is_err())); + assert_eq!(reader.next().transpose()?, Some((3, String::from("three")))); + assert!(reader.next().is_none()); + Ok(()) + } +} diff --git a/backend/src/repository/qos_snapshot_repository.rs b/backend/src/repository/qos_snapshot_repository.rs index 301087b33..e9c23e483 100644 --- a/backend/src/repository/qos_snapshot_repository.rs +++ b/backend/src/repository/qos_snapshot_repository.rs @@ -300,11 +300,9 @@ impl QosSnapshotRepository { F: FnMut(&QosSnapshotRecord), { let tree = load_snapshot_tree(&self.snapshot_path)?; - tree.traverse(|_keys, values| { - for snapshot in values { - visit(snapshot); - } - }); + for (_, snapshot) in &tree { + visit(snapshot); + } Ok(()) } @@ -320,11 +318,9 @@ impl QosSnapshotRepository { { let snapshot_path = Self::snapshot_db_path(storage_dir); let tree = load_snapshot_tree(&snapshot_path)?; - tree.traverse(|_keys, values| { - for snapshot in values { - visit(snapshot); - } - }); + for (_, snapshot) in &tree { + visit(snapshot); + } Ok(()) } diff --git a/backend/src/repository/sorted_index.rs b/backend/src/repository/sorted_index.rs deleted file mode 100644 index f1179f931..000000000 --- a/backend/src/repository/sorted_index.rs +++ /dev/null @@ -1,788 +0,0 @@ -//! Sorted Index for `BPlusTree` -//! -//! This module provides a compact sorted index file that enables iteration -//! over a `BPlusTree` in order of a secondary sort key rather than the primary key. -//! -//! # File Format (v3) -//! ```text -//! [magic: 4 bytes]["SIDX"] -//! [version: u32] -//! [count: u64] -//! [entry0][entry1]... -//! -//! Entry format (with value location for O(1) access): -//! [sort_key_len: u32][sort_key_bytes][primary_key_len: u32][primary_key_bytes][value_location] -//! -//! Value location format: -//! - Single mode: [mode: u8 = 0][offset: u64][length: u32] -//! - Packed mode: [mode: u8 = 1][block_offset: u64][index: u16][length: u32] -//! ``` -//! -//! # Optimization -//! By storing the value location directly in the index, iteration can -//! read values in O(1) time by seeking directly to the offset, avoiding O(log n) -//! tree traversal for each lookup. -//! -//! > **Important**: The index is tightly coupled to the B+Tree file structure. -//! > It must be rebuilt after any operation that changes value offsets (e.g., `compact()`). - -use crate::repository::bplustree::{BPlusTreeQuery, COMPRESSION_FLAG_LZ4, PAGE_SIZE_USIZE}; -use crate::repository::storage::get_file_path_for_db_index; -use crate::utils::{binary_deserialize, binary_serialize}; -use indexmap::IndexMap; -use serde::{Deserialize, Serialize}; -use std::fs::{File, OpenOptions}; -use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write}; -use std::marker::PhantomData; -use std::path::{Path, PathBuf}; - -const MAGIC: &[u8; 4] = b"SIDX"; -const VERSION: u32 = 3; // Bumped for new format with flexible value location -const HEADER_SIZE: usize = 16; // 4 (magic) + 4 (version) + 8 (count) - -const MODE_SINGLE: u8 = 0; -const MODE_PACKED: u8 = 1; - - -/// Represents how a value can be located and read from the tree file. -#[derive(Debug, Clone, Copy)] -pub enum ValueLocation { - /// Single value stored at a specific offset with a given length. - Single { offset: u64, length: u32 }, - /// Value packed in a block with other values, identified by block offset and index. - Packed { block_offset: u64, index: u16, length: u32 }, -} - -impl ValueLocation { - /// Serialize the value location to bytes. - fn to_bytes(self) -> Vec { - match self { - ValueLocation::Single { offset, length } => { - let mut bytes = vec![MODE_SINGLE]; - bytes.extend_from_slice(&offset.to_le_bytes()); - bytes.extend_from_slice(&length.to_le_bytes()); - bytes - } - ValueLocation::Packed { block_offset, index, length } => { - let mut bytes = vec![MODE_PACKED]; - bytes.extend_from_slice(&block_offset.to_le_bytes()); - bytes.extend_from_slice(&index.to_le_bytes()); - bytes.extend_from_slice(&length.to_le_bytes()); - bytes - } - } - } - - /// Deserialize a value location from a reader. - fn from_reader(reader: &mut R) -> io::Result { - let mut mode = [0u8; 1]; - reader.read_exact(&mut mode)?; - - match mode[0] { - MODE_SINGLE => { - let mut offset_buf = [0u8; 8]; - let mut length_buf = [0u8; 4]; - reader.read_exact(&mut offset_buf)?; - reader.read_exact(&mut length_buf)?; - Ok(ValueLocation::Single { - offset: u64::from_le_bytes(offset_buf), - length: u32::from_le_bytes(length_buf), - }) - } - MODE_PACKED => { - let mut block_offset_buf = [0u8; 8]; - let mut index_buf = [0u8; 2]; - let mut length_buf = [0u8; 4]; - reader.read_exact(&mut block_offset_buf)?; - reader.read_exact(&mut index_buf)?; - reader.read_exact(&mut length_buf)?; - Ok(ValueLocation::Packed { - block_offset: u64::from_le_bytes(block_offset_buf), - index: u16::from_le_bytes(index_buf), - length: u32::from_le_bytes(length_buf), - }) - } - _ => Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Unknown value location mode: {}", mode[0]), - )), - } - } -} - -/// Entry containing value location information for direct access. -#[derive(Debug, Clone)] -pub struct IndexEntry { - pub sort_key: SortKey, - pub primary_key: K, - pub location: ValueLocation, -} - -/// Writer for building a sorted index file. -/// -/// Entries must be pushed in sorted order. The writer buffers writes -/// and flushes on `finish()`. -pub struct SortedIndexWriter { - writer: BufWriter, - count: u64, - _marker: PhantomData<(SortKey, K)>, -} - -impl SortedIndexWriter -where - SortKey: Serialize, - K: Serialize, -{ - /// Create a new index writer at the given path. - /// Overwrites any existing file. - pub fn new(path: &Path) -> io::Result { - let file = OpenOptions::new() - .write(true) - .create(true) - .truncate(true) - .open(path)?; - - let mut writer = BufWriter::new(file); - - // Write header (count will be updated on finish) - writer.write_all(MAGIC)?; - writer.write_all(&VERSION.to_le_bytes())?; - writer.write_all(&0u64.to_le_bytes())?; // placeholder count - - Ok(Self { - writer, - count: 0, - _marker: PhantomData, - }) - } - - /// Append an entry to the index with value location for O(1) access. - /// Caller must ensure entries are pushed in sorted order. - pub fn push( - &mut self, - sort_key: &SortKey, - primary_key: &K, - location: ValueLocation, - ) -> io::Result<()> { - let sk_bytes = binary_serialize(sort_key)?; - let pk_bytes = binary_serialize(primary_key)?; - - // Write sort key - let sk_len = u32::try_from(sk_bytes.len()) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - self.writer.write_all(&sk_len.to_le_bytes())?; - self.writer.write_all(&sk_bytes)?; - - // Write primary key - let pk_len = u32::try_from(pk_bytes.len()) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - self.writer.write_all(&pk_len.to_le_bytes())?; - self.writer.write_all(&pk_bytes)?; - - // Write value location - self.writer.write_all(&location.to_bytes())?; - - self.count += 1; - Ok(()) - } - - /// Finalize the index file by writing the count to the header. - pub fn finish(mut self) -> io::Result { - self.writer.flush()?; - - // Seek back to count position and write final count - let mut file = self.writer.into_inner()?; - file.seek(SeekFrom::Start(8))?; // After magic + version - file.write_all(&self.count.to_le_bytes())?; - file.sync_all()?; - - Ok(self.count) - } -} - -/// Reader for iterating over a sorted index file. -pub struct SortedIndexReader { - reader: BufReader, - remaining: u64, - _marker: PhantomData<(SortKey, K)>, -} - -impl SortedIndexReader -where - SortKey: for<'de> Deserialize<'de>, - K: for<'de> Deserialize<'de>, -{ - /// Open an existing index file for reading. - pub fn open(path: &Path) -> io::Result { - let file = File::open(path)?; - let mut reader = BufReader::new(file); - - // Verify header - let mut header = [0u8; HEADER_SIZE]; - reader.read_exact(&mut header)?; - - if &header[0..4] != MAGIC { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "Invalid sorted index magic", - )); - } - - let version = u32::from_le_bytes( - header[4..8] - .try_into() - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("Invalid version bytes: {e}")))?, - ); - if version != VERSION { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Unsupported sorted index version: {version} (expected {VERSION})"), - )); - } - - let count = match header[8..16].try_into() { - Ok(count) => u64::from_le_bytes(count), - Err(e) => { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid count {e}"), - )); - } - }; - - Ok(Self { - reader, - remaining: count, - _marker: PhantomData, - }) - } - - /// Returns the number of remaining entries to read. - pub fn remaining(&self) -> u64 { - self.remaining - } - - /// Returns true if the index is empty. - pub fn is_empty(&self) -> bool { - self.remaining == 0 - } - - /// Read the next entry from the index. - pub fn read_next(&mut self) -> io::Result>> { - if self.remaining == 0 { - return Ok(None); - } - - // Read sort key - let mut len_buf = [0u8; 4]; - self.reader.read_exact(&mut len_buf)?; - let sk_len = u32::from_le_bytes(len_buf) as usize; - - let mut sk_bytes = vec![0u8; sk_len]; - self.reader.read_exact(&mut sk_bytes)?; - let sort_key: SortKey = binary_deserialize(&sk_bytes)?; - - // Read primary key - self.reader.read_exact(&mut len_buf)?; - let pk_len = u32::from_le_bytes(len_buf) as usize; - - let mut pk_bytes = vec![0u8; pk_len]; - self.reader.read_exact(&mut pk_bytes)?; - let primary_key: K = binary_deserialize(&pk_bytes)?; - - // Read value location - let location = ValueLocation::from_reader(&mut self.reader)?; - - self.remaining -= 1; - Ok(Some(IndexEntry { - sort_key, - primary_key, - location, - })) - } -} - -impl Iterator for SortedIndexReader -where - SortKey: for<'de> Deserialize<'de>, - K: for<'de> Deserialize<'de>, -{ - type Item = io::Result>; - - fn next(&mut self) -> Option { - match self.read_next() { - Ok(Some(entry)) => Some(Ok(entry)), - Ok(None) => None, - Err(e) => Some(Err(e)), - } - } -} - -/// Owned iterator that combines a sorted index with a primary `BPlusTree` file. -/// -/// Iterates through the index in order and reads values directly from the -/// tree file using stored offsets - no tree traversal needed (O(1) per item). -pub struct BPlusTreeSortedIteratorOwned { - index_reader: SortedIndexReader, - tree_file: Option>, - mmap: Option, - filepath: PathBuf, - block_cache: IndexMap>, - live_key_query: Option>, - _marker: PhantomData, -} - -const CACHE_CAPACITY: usize = 8; - -impl BPlusTreeSortedIteratorOwned -where - K: for<'de> Deserialize<'de>, - V: for<'de> Deserialize<'de>, - SortKey: for<'de> Deserialize<'de>, -{ - pub(crate) fn from_index_reader( - index_reader: SortedIndexReader, - filepath: PathBuf, - tree_file: Option>, - mmap: Option, - ) -> Self { - Self { - index_reader, - tree_file, - mmap, - filepath, - block_cache: IndexMap::with_capacity(CACHE_CAPACITY), - live_key_query: None, - _marker: PhantomData, - } - } - - /// Create a new owned sorted iterator. - /// - /// The index path is automatically derived from the tree filepath - /// by changing the extension to `.idx`. - pub fn new(filepath: PathBuf, tree_file: Option>) -> io::Result { - let index_path = get_file_path_for_db_index(&filepath); - let index_reader = SortedIndexReader::open(&index_path)?; - Ok(Self { - index_reader, - tree_file, - mmap: None, - filepath, - block_cache: IndexMap::with_capacity(CACHE_CAPACITY), - live_key_query: None, - _marker: PhantomData, - }) - } - - pub fn new_hybrid( - filepath: PathBuf, - tree_file: Option>, - mmap: Option, - ) -> io::Result { - let index_path = get_file_path_for_db_index(&filepath); - let index_reader = SortedIndexReader::open(&index_path)?; - Ok(Self { - index_reader, - tree_file, - mmap, - filepath, - block_cache: IndexMap::with_capacity(CACHE_CAPACITY), - live_key_query: None, - _marker: PhantomData, - }) - } - - /// Create from an explicit index path. - pub fn with_index_path( - filepath: PathBuf, - tree_file: Option>, - index_path: &Path, - ) -> io::Result { - let index_reader = SortedIndexReader::open(index_path)?; - Ok(Self { - index_reader, - tree_file, - mmap: None, - filepath, - block_cache: IndexMap::with_capacity(CACHE_CAPACITY), - live_key_query: None, - _marker: PhantomData, - }) - } - - pub fn with_index_path_hybrid( - filepath: PathBuf, - tree_file: Option>, - mmap: Option, - index_path: &Path, - ) -> io::Result { - let index_reader = SortedIndexReader::open(index_path)?; - Ok(Self { - index_reader, - tree_file, - mmap, - filepath, - block_cache: IndexMap::with_capacity(CACHE_CAPACITY), - live_key_query: None, - _marker: PhantomData, - }) - } - - /// Returns the number of remaining items. - pub fn remaining(&self) -> u64 { - self.index_reader.remaining - } - - /// Returns the path to the tree file. - pub fn filepath(&self) -> &Path { - &self.filepath - } - - /// Read a value from the tree file using the given `ValueLocation`. - fn read_value(&mut self, location: ValueLocation) -> io::Result { - match location { - ValueLocation::Single { offset, length } => { - self.read_value_single(offset, length) - } - ValueLocation::Packed { block_offset, index, length: _ } => { - self.read_value_packed(block_offset, index) - } - } - } - - /// Read a single value directly from the tree file at the given offset. - fn read_value_single(&mut self, offset: u64, length: u32) -> io::Result { - if let Some(mmap) = &self.mmap { - let mut cursor = io::Cursor::new(mmap.as_ref()); - cursor.set_position(offset); - - let mut flag = [0u8; 1]; - cursor.read_exact(&mut flag)?; - - let data = if flag[0] == COMPRESSION_FLAG_LZ4 { - let compressed_len = length as usize - 1; - let mut compressed = vec![0u8; compressed_len]; - cursor.read_exact(&mut compressed)?; - lz4_flex::decompress_size_prepended(&compressed).map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("LZ4 decompression failed: {e}")) - })? - } else { - let payload_len = length as usize - 1; - let mut data = vec![0u8; payload_len]; - cursor.read_exact(&mut data)?; - data - }; - crate::utils::binary_deserialize(&data) - } else if let Some(tree_file) = &mut self.tree_file { - tree_file.seek(SeekFrom::Start(offset))?; - - // Read compression flag - let mut flag = [0u8; 1]; - tree_file.read_exact(&mut flag)?; - - let data = if flag[0] == COMPRESSION_FLAG_LZ4 { - // Compressed: [flag:1][lz4_payload_with_prepended_size] - let compressed_len = length as usize - 1; - let mut compressed = vec![0u8; compressed_len]; - tree_file.read_exact(&mut compressed)?; - - lz4_flex::decompress_size_prepended(&compressed).map_err(|e| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("LZ4 decompression failed: {e}"), - ) - })? - } else { - // Uncompressed: [flag:1][payload] - let payload_len = length as usize - 1; - let mut data = vec![0u8; payload_len]; - tree_file.read_exact(&mut data)?; - data - }; - - crate::utils::binary_deserialize(&data) - } else { - Err(io::Error::other("No data source available")) - } - } - - /// Read a value from a packed block at the given index. - fn read_value_packed(&mut self, block_offset: u64, value_index: u16) -> io::Result { - if let Some(mmap) = &self.mmap { - let start = usize::try_from(block_offset).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - if start >= mmap.len() { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block offset out of range")); - } - let end = (start + PAGE_SIZE_USIZE).min(mmap.len()); - let block_buffer = &mmap[start..end]; - if block_buffer.len() < 4 { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block truncated")); - } - // Read count (first 4 bytes) - let count = u32::from_le_bytes(block_buffer[0..4].try_into().map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("Invalid count: {e}")) - })?); - - if u32::from(value_index) >= count { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Value index {value_index} out of bounds (count: {count})"), - )); - } - - let mut pos = 4; - // Skip to target value - for i in 0..=value_index { - if pos + 4 > block_buffer.len() { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block corrupted")); - } - let len = u32::from_le_bytes(block_buffer[pos..pos + 4].try_into().map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("Invalid length: {e}")) - })?) as usize; - pos += 4; - - if i == value_index { - if pos + len > block_buffer.len() { - return Err(io::Error::new(io::ErrorKind::InvalidData, "Packed block corrupted")); - } - let value_data = &block_buffer[pos..pos + len]; - return crate::utils::binary_deserialize(value_data); - } - pos += len; - } - return Err(io::Error::other("Value not found")); - } - - // Try cache first - if !self.block_cache.contains_key(&block_offset) { - // Miss - read from disk - if let Some(tree_file) = &mut self.tree_file { - tree_file.seek(SeekFrom::Start(block_offset))?; - let mut buf = vec![0u8; PAGE_SIZE_USIZE]; - tree_file.read_exact(&mut buf)?; - - // Update cache - if self.block_cache.len() >= CACHE_CAPACITY { - self.block_cache.shift_remove_index(0); // FIFO eviction - } - self.block_cache.insert(block_offset, buf); - } else { - return Err(io::Error::other("No data source available")); - } - } - - let block_buffer = self - .block_cache - .get(&block_offset) - .ok_or_else(|| io::Error::other("Packed block missing from cache"))?; - - // Read count (first 4 bytes) - let count = u32::from_le_bytes(block_buffer[0..4].try_into().map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("Invalid count: {e}")) - })?); - - if u32::from(value_index) >= count { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Value index {value_index} out of bounds (count: {count})"), - )); - } - - let mut pos = 4; - - // Skip to target value - for i in 0..=value_index { - if pos + 4 > PAGE_SIZE_USIZE { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Packed block corrupted: position {pos} exceeds block size"), - )); - } - - let len = u32::from_le_bytes(block_buffer[pos..pos + 4].try_into().map_err(|e| { - io::Error::new(io::ErrorKind::InvalidData, format!("Invalid length: {e}")) - })?) as usize; - pos += 4; - - if i == value_index { - // Found target value - if pos + len > PAGE_SIZE_USIZE { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Packed value corrupted: length {len} at position {pos} exceeds block size"), - )); - } - let value_data = &block_buffer[pos..pos + len]; - return crate::utils::binary_deserialize(value_data); - } - - pos += len; - } - - Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Value index {value_index} not found in packed block (count: {count})"), - )) - } - -} - -impl BPlusTreeSortedIteratorOwned -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - SortKey: for<'de> Deserialize<'de>, -{ - fn ensure_live_key_query(&mut self) -> io::Result<()> { - if self.live_key_query.is_some() || self.filepath.as_os_str().is_empty() { - return Ok(()); - } - - let query = BPlusTreeQuery::::try_new(&self.filepath)?; - self.live_key_query = Some(query); - Ok(()) - } - - fn is_live_key(&mut self, key: &K) -> io::Result { - self.ensure_live_key_query()?; - if let Some(query) = &mut self.live_key_query { - if !query.has_tombstones() { - return Ok(true); - } - return query - .contains_live_key(key) - .map_err(|err| io::Error::other(format!("sorted iterator live-key check failed: {err}"))); - } - - Ok(true) - } -} - -impl Iterator for BPlusTreeSortedIteratorOwned -where - K: Ord + Serialize + for<'de> Deserialize<'de> + Clone, - V: Serialize + for<'de> Deserialize<'de> + Clone, - SortKey: for<'de> Deserialize<'de>, -{ - type Item = io::Result<(K, V)>; - - fn next(&mut self) -> Option { - loop { - let entry = match self.index_reader.read_next() { - Ok(Some(entry)) => entry, - Ok(None) => return None, - Err(e) => return Some(Err(e)), - }; - - match self.is_live_key(&entry.primary_key) { - Ok(true) => {} - Ok(false) => continue, - Err(e) => return Some(Err(e)), - } - - match self.read_value(entry.location) { - Ok(value) => return Some(Ok((entry.primary_key, value))), - Err(e) => return Some(Err(e)), - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use tempfile::tempdir; - - #[test] - fn test_sorted_index_write_read() { - let dir = tempdir().unwrap(); - let path = dir.path().join("test.idx"); - - // Write entries with value locations - let mut writer = SortedIndexWriter::::new(&path).unwrap(); - writer.push( - &"apple".to_string(), - &1u32, - ValueLocation::Single { offset: 100, length: 50 }, - ).unwrap(); - writer.push( - &"banana".to_string(), - &2u32, - ValueLocation::Packed { block_offset: 200, index: 3, length: 75 }, - ).unwrap(); - writer.push( - &"cherry".to_string(), - &3u32, - ValueLocation::Single { offset: 300, length: 100 }, - ).unwrap(); - let count = writer.finish().unwrap(); - assert_eq!(count, 3); - - // Read entries - let reader = SortedIndexReader::::open(&path).unwrap(); - let entries: Vec<_> = reader.map(|r| r.unwrap()).collect(); - - assert_eq!(entries.len(), 3); - assert_eq!(entries[0].sort_key, "apple".to_string()); - assert_eq!(entries[0].primary_key, 1u32); - match entries[0].location { - ValueLocation::Single { offset, length } => { - assert_eq!(offset, 100); - assert_eq!(length, 50); - } - ValueLocation::Packed { .. } => panic!("Expected Single location"), - } - - assert_eq!(entries[1].sort_key, "banana".to_string()); - assert_eq!(entries[1].primary_key, 2u32); - match entries[1].location { - ValueLocation::Packed { block_offset, index, length } => { - assert_eq!(block_offset, 200); - assert_eq!(index, 3); - assert_eq!(length, 75); - } - ValueLocation::Single { .. } => panic!("Expected Packed location"), - } - - assert_eq!(entries[2].sort_key, "cherry".to_string()); - assert_eq!(entries[2].primary_key, 3u32); - match entries[2].location { - ValueLocation::Single { offset, length } => { - assert_eq!(offset, 300); - assert_eq!(length, 100); - } - ValueLocation::Packed { .. } => panic!("Expected Single location"), - } - } - - #[test] - fn test_empty_index() { - let dir = tempdir().unwrap(); - let path = dir.path().join("empty.idx"); - - let writer = SortedIndexWriter::::new(&path).unwrap(); - let count = writer.finish().unwrap(); - assert_eq!(count, 0); - - let reader = SortedIndexReader::::open(&path).unwrap(); - assert!(reader.is_empty()); - // Note: collect() consumes the reader, so we check is_empty() and remaining first - let remaining = reader.remaining; - assert_eq!(remaining, 0); - - let entries: Vec<_> = reader.collect(); - assert!(entries.is_empty()); - } - - #[test] - fn test_index_path_derivation() { - let tree_path = Path::new("/data/my_tree.bin"); - let idx_path = get_file_path_for_db_index(tree_path); - assert_eq!(idx_path, PathBuf::from("/data/my_tree.idx")); - - let tree_path2 = Path::new("/data/items"); - let idx_path2 = get_file_path_for_db_index(tree_path2); - assert_eq!(idx_path2, PathBuf::from("/data/items.idx")); - } -} diff --git a/backend/src/repository/target_id_mapping.rs b/backend/src/repository/target_id_mapping.rs index 9de8e2909..ce844cc99 100644 --- a/backend/src/repository/target_id_mapping.rs +++ b/backend/src/repository/target_id_mapping.rs @@ -119,15 +119,11 @@ impl TargetIdMapping { let mut mem_by_uuid = HashMap::new(); let mut mem_by_virtual_id = HashMap::new(); - tree.traverse(|keys, values| { - if let Some(max_key) = keys.iter().max() { - virtual_id_counter = max(virtual_id_counter, *max_key); - } - for v in values { - mem_by_uuid.insert(v.uuid, v.virtual_id); - mem_by_virtual_id.insert(v.virtual_id, v.clone()); - } - }); + for (key, value) in &tree { + virtual_id_counter = max(virtual_id_counter, *key); + mem_by_uuid.insert(value.uuid, value.virtual_id); + mem_by_virtual_id.insert(value.virtual_id, value.clone()); + } Ok(Self { virtual_id_counter, diff --git a/backend/src/repository/user_repository.rs b/backend/src/repository/user_repository.rs index 6a64347f1..9ecf3b59e 100644 --- a/backend/src/repository/user_repository.rs +++ b/backend/src/repository/user_repository.rs @@ -17,7 +17,7 @@ use std::path::{Path, PathBuf}; use tokio::task; // V6 (current): added network_access. V1-V5 are migrated to V6 at startup -// by `bplustree_migration::run_all_startup_migrations`. +// by `bplustree::run_all_startup_migrations`. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] struct StoredProxyUserCredentials { pub target: String, diff --git a/backend/src/repository/xtream_playlist_iterator.rs b/backend/src/repository/xtream_playlist_iterator.rs index 9110f75a8..504aad021 100644 --- a/backend/src/repository/xtream_playlist_iterator.rs +++ b/backend/src/repository/xtream_playlist_iterator.rs @@ -17,7 +17,7 @@ use tokio::sync::mpsc; use tokio::task; pub struct XtreamPlaylistIterator { - inner: LockedReceiverStream<(XtreamPlaylistItem, bool)>, + inner: LockedReceiverStream>, } fn is_cluster_allowed_for_user(user: &ProxyUserCredentials, cluster: XtreamCluster) -> bool { @@ -28,7 +28,7 @@ impl XtreamPlaylistIterator { fn empty() -> Self { // Ghost Channel pattern // When you immediately drop the sender with `_`, the channel is closed and receiver gets None. - let (_tx, rx) = mpsc::channel::<(XtreamPlaylistItem, bool)>(1); + let (_tx, rx) = mpsc::channel::>(1); Self { inner: LockedReceiverStream::new_empty(rx) } } @@ -73,9 +73,10 @@ impl XtreamPlaylistIterator { let xtream_path = xtream_path.clone(); let index_path = get_file_path_for_db_index(&xtream_path); - let (tx, rx) = mpsc::channel::<(XtreamPlaylistItem, bool)>(256); + let (tx, rx) = mpsc::channel::>(256); let xtream_path_for_log = xtream_path.clone(); + let join_error_tx = tx.clone(); let handle = task::spawn_blocking(move || { let _guard = bg_lock; let reader = match open_playlist_reader::( @@ -89,6 +90,7 @@ impl XtreamPlaylistIterator { "Failed to open Xtream playlist DB {} (cluster {cluster}): {err}", xtream_path.display() ); + let _ = tx.blocking_send(Err(err)); return; } }; @@ -98,7 +100,7 @@ impl XtreamPlaylistIterator { let item = match entry { Ok((_, item)) => item, Err(err) => { - error!("Error reading sorted index: {err}"); + error!("Skipping unreadable Xtream playlist entry: {err}"); continue; } }; @@ -108,14 +110,14 @@ impl XtreamPlaylistIterator { } if let Some(prev) = pending.replace(item) { - if tx.blocking_send((prev, true)).is_err() { + if tx.blocking_send(Ok((prev, true))).is_err() { return; } } } if let Some(last) = pending { - let _ = tx.blocking_send((last, false)); + let _ = tx.blocking_send(Ok((last, false))); } }); tokio::spawn(async move { @@ -124,6 +126,12 @@ impl XtreamPlaylistIterator { "Xtream playlist iterator task failed for {} (cluster {cluster}): {err}", xtream_path_for_log.display() ); + let _ = join_error_tx + .send(Err(TuliproxError::RepositoryXtream(format!( + "Xtream playlist iterator task failed for {}: {err}", + xtream_path_for_log.display() + )))) + .await; } }); @@ -153,7 +161,7 @@ impl XtreamPlaylistIterator { } impl Stream for XtreamPlaylistIterator { - type Item = (XtreamPlaylistItem, bool); + type Item = Result<(XtreamPlaylistItem, bool), TuliproxError>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_next(cx) } @@ -198,19 +206,18 @@ impl XtreamPlaylistJsonIterator { impl Stream for XtreamPlaylistJsonIterator { - type Item = (String, bool); + type Item = Result<(String, bool), TuliproxError>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { match Pin::new(&mut self.inner).poll_next(cx) { - Poll::Ready(Some((pli, has_next))) => { + Poll::Ready(Some(Ok((pli, has_next)))) => { let Some(options) = self.options.as_ref() else { return Poll::Ready(None); }; - let json = serde_json::to_string(&pli.to_document(options)).unwrap_or_else(|err| { - error!("Failed to serialize playlist item {}: {err}", pli.virtual_id); - "{}".to_string() - }); - Poll::Ready(Some((json, has_next))) + let json = serde_json::to_string(&pli.to_document(options)) + .map_err(|error| TuliproxError::RepositoryXtream(error.to_string())); + Poll::Ready(Some(json.map(|json| (json, has_next)))) } + Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))), Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, } @@ -222,8 +229,13 @@ impl Stream for XtreamPlaylistJsonIterator { mod tests { use super::{is_cluster_allowed_for_user, XtreamPlaylistIterator}; use crate::model::ProxyUserCredentials; - use shared::model::{ClusterFlags, XtreamCluster}; + use crate::repository::LockedReceiverStream; + use shared::{ + model::{ClusterFlags, PlaylistItemType, XtreamCluster, XtreamPlaylistItem}, + utils::Internable, + }; use futures::StreamExt; + use tokio::sync::mpsc; #[test] fn cluster_guard_respects_user_cluster_flags() { @@ -240,4 +252,38 @@ mod tests { let mut iter = XtreamPlaylistIterator::empty(); assert!(iter.next().await.is_none()); } + + #[tokio::test] + async fn iterator_forwards_one_storage_error_then_ends() { + let (tx, rx) = mpsc::channel(2); + let item = XtreamPlaylistItem { + virtual_id: 1, + provider_id: 1, + name: "name".intern(), + logo: "".intern(), + logo_small: "".intern(), + group: "group".intern(), + title: "title".intern(), + parent_code: "".intern(), + rec: "".intern(), + url: "http://example.test/live.ts".intern(), + epg_channel_id: None, + xtream_cluster: XtreamCluster::Live, + additional_properties: None, + item_type: PlaylistItemType::Live, + category_id: 1, + input_name: "input".intern(), + channel_no: 0, + source_ordinal: 0, + input_stream_id: "1".intern(), + }; + assert!(tx.send(Ok((item, true))).await.is_ok()); + assert!(tx.send(Err(shared::error::TuliproxError::RepositoryXtream("corrupt page".into()))).await.is_ok()); + drop(tx); + + let mut iterator = XtreamPlaylistIterator { inner: LockedReceiverStream::new_empty(rx) }; + assert!(iterator.next().await.is_some_and(|entry| entry.is_ok())); + assert!(iterator.next().await.is_some_and(|entry| entry.is_err())); + assert!(iterator.next().await.is_none()); + } } diff --git a/backend/src/repository/xtream_repository.rs b/backend/src/repository/xtream_repository.rs index 650c20dec..1189883ab 100644 --- a/backend/src/repository/xtream_repository.rs +++ b/backend/src/repository/xtream_repository.rs @@ -660,14 +660,14 @@ pub async fn xtream_load_rewrite_playlist( XtreamPlaylistJsonIterator::new(cluster, app_state, target, category_id, user).await } -pub async fn iter_raw_xtream_target_playlist(app_config: &AppConfig, target: &ConfigTarget, cluster: XtreamCluster) -> Option + Send + Unpin>> { +pub async fn iter_raw_xtream_target_playlist(app_config: &AppConfig, target: &ConfigTarget, cluster: XtreamCluster) -> Option> + Send + Unpin>> { let config = app_config.config.load(); let storage_path = xtream_get_storage_path(&config, target.name.as_str())?; let xtream_path = xtream_get_file_path(&storage_path, cluster); iter_raw_xtream_playlist(app_config, &xtream_path).await } -pub async fn iter_raw_xtream_input_playlist(app_config: &AppConfig, input: &ConfigInput, cluster: XtreamCluster) -> Option + Send + Unpin>> { +pub async fn iter_raw_xtream_input_playlist(app_config: &AppConfig, input: &ConfigInput, cluster: XtreamCluster) -> Option> + Send + Unpin>> { let config = app_config.config.load(); let storage_dir = &config.storage_dir; let storage_path = get_input_storage_path(&input.name, storage_dir).await.ok()?; @@ -676,7 +676,7 @@ pub async fn iter_raw_xtream_input_playlist(app_config: &AppConfig, input: &Conf iter_raw_xtream_playlist(app_config, &xtream_path).await } -async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> Option + Send + Unpin>> { +async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> Option> + Send + Unpin>> { if !file_exists_async(xtream_path).await { return None; } @@ -684,10 +684,11 @@ async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> let xtream_path = xtream_path.to_path_buf(); let index_path = get_file_path_for_db_index(&xtream_path); - let (tx, rx) = mpsc::channel::(256); + let (tx, rx) = mpsc::channel::>(256); let xtream_path_for_log = xtream_path.clone(); let index_path_for_log = index_path.clone(); + let join_error_tx = tx.clone(); let handle = tokio::task::spawn_blocking(move || { let _guard = bg_lock; let reader = match open_playlist_reader::( @@ -702,7 +703,7 @@ async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> xtream_path.display(), index_path.display() ); - drop(tx); + let _ = tx.blocking_send(Err(err)); return; } }; @@ -711,11 +712,11 @@ async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> let item = match entry { Ok((_, item)) => item, Err(err) => { - error!("Xtream playlist reader error: {err}"); + error!("Skipping unreadable Xtream playlist entry: {err}"); continue; } }; - if tx.blocking_send(item).is_err() { + if tx.blocking_send(Ok(item)).is_err() { break; } } @@ -727,10 +728,16 @@ async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> xtream_path_for_log.display(), index_path_for_log.display() ); + let _ = join_error_tx + .send(Err(TuliproxError::RepositoryXtream(format!( + "Xtream playlist producer task failed for {}: {err}", + xtream_path_for_log.display() + )))) + .await; } }); - let stream: Box + Send + Unpin> = + let stream: Box> + Send + Unpin> = Box::new(ReceiverStream::new(rx)); Some(stream) } @@ -785,7 +792,8 @@ fn preserve_details_input_xtream_playlist_cluster_to_disk( }; let mut updates: Vec<(u32, XtreamPlaylistItem)> = Vec::with_capacity(BATCH_SIZE); - for (_, old_item) in old_tree.iter() { + for entry in old_tree.iter() { + let (_, old_item) = entry.map_err(|error| TuliproxError::RepositoryXtream(error.to_string()))?; if let Some(old_props) = old_item.additional_properties.as_ref() { if old_props.has_details() { if let Ok(Some(mut new_item)) = new_tree.query(&old_item.provider_id) { @@ -969,7 +977,7 @@ pub async fn persist_input_xtream_playlist_cluster_to_disk( 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)?; + tree_update.compact()?; } Ok(()) }) @@ -1075,16 +1083,20 @@ pub async fn persist_input_xtream_playlist(app_config: &Arc, storage_ let stored_entries = match tokio::task::spawn_blocking(move || { let _guard = file_lock; let mut entries = IndexMap::new(); - if let Ok(mut query) = BPlusTreeQuery::::try_new(&xtream_path) { - for (_, doc) in query.iter() { - entries.insert(doc.provider_id, doc); - } + let mut query = BPlusTreeQuery::::try_new(&xtream_path)?; + for entry in query.iter() { + let (_, doc) = entry?; + entries.insert(doc.provider_id, doc); } - entries + Ok::<_, std::io::Error>(entries) }) .await { - Ok(entries) => Some(entries), + Ok(Ok(entries)) => Some(entries), + Ok(Err(err)) => { + errors.push(format!("Failed to read stored xtream playlist entries for {cluster}: {err}")); + None + } Err(err) => { errors.push(format!( "Failed to load stored xtream playlist entries for {cluster}: {err}" @@ -1099,6 +1111,10 @@ pub async fn persist_input_xtream_playlist(app_config: &Arc, storage_ } } + if !errors.is_empty() { + return (playlist, Some(TuliproxError::RepositoryXtream(errors.join("\n")))); + } + let mut groups = IndexMap::new(); for mut plg in playlist { @@ -1476,10 +1492,11 @@ pub async fn load_input_xtream_playlist(app_config: &Arc, storage_pat let items = tokio::task::spawn_blocking(move || -> Result, TuliproxError> { let _guard = file_lock; let mut items = Vec::new(); - if let Ok(mut query) = BPlusTreeQuery::::try_new(&xtream_path) { - for (_, item) in query.iter() { - items.push(item); - } + let mut query = BPlusTreeQuery::::try_new(&xtream_path) + .map_err(|error| TuliproxError::RepositoryXtream(error.to_string()))?; + for entry in query.iter() { + let (_, item) = entry.map_err(|error| TuliproxError::RepositoryXtream(error.to_string()))?; + items.push(item); } Ok(items) }) diff --git a/backend/src/utils/db_viewer.rs b/backend/src/utils/db_viewer.rs index 353c02116..193c47290 100644 --- a/backend/src/utils/db_viewer.rs +++ b/backend/src/utils/db_viewer.rs @@ -188,7 +188,15 @@ where println!("["); let mut first = true; - for (_, entry) in iterator { + for item in iterator { + let (_, entry) = match item { + Ok(entry) => entry, + Err(err) => { + error!("Failed to iterate database: {err}"); + error_count += 1; + break; + } + }; match to_human_readable_json_value(&entry).and_then(|value| serde_json::to_string(&value)) { Ok(json) => { if !first { diff --git a/backend/src/utils/network/content_coding.rs b/backend/src/utils/network/content_coding.rs index 394248e70..fe8093b4b 100644 --- a/backend/src/utils/network/content_coding.rs +++ b/backend/src/utils/network/content_coding.rs @@ -891,7 +891,9 @@ mod tests { let mut response = format!("HTTP/1.1 {} Test\r\nConnection: close\r\n", status.as_u16()); if !has_content_length { - response.push_str(&format!("Content-Length: {}\r\n", body.len())); + response.push_str("Content-Length: "); + response.push_str(&body.len().to_string()); + response.push_str("\r\n"); } for (name, value) in owned_headers { response.push_str(&name); diff --git a/backend/src/utils/network/request.rs b/backend/src/utils/network/request.rs index e440c4974..e0f76dfe4 100644 --- a/backend/src/utils/network/request.rs +++ b/backend/src/utils/network/request.rs @@ -3982,9 +3982,8 @@ mod tests { panic!("provider-only request lost the default idle-timeout guard"); } - let error = match request.await.expect("request task should join") { - Ok(_) => panic!("hanging request must time out"), - Err(error) => error, + let Err(error) = request.await.expect("request task should join") else { + panic!("hanging request must time out"); }; assert_eq!(error.kind(), ErrorKind::TimedOut); server.abort();