BPlusTree refactoring / optimization (#792)

* 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.
This commit is contained in:
euzu
2026-07-17 17:56:29 +02:00
committed by GitHub
parent e3956a30cd
commit 5b49ea58d8
45 changed files with 14312 additions and 9881 deletions
+1
View File
@@ -1,4 +1,5 @@
.tools
.worktrees
dist
/resources/*.ts
/release
+15
View File
@@ -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
+9 -10
View File
@@ -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<str>`) 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
+162 -9
View File
@@ -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<P, S, E>(
accept: Option<&str>,
data: S,
) -> axum::response::Response
where
P: serde::Serialize + Send + 'static,
S: Stream<Item = Result<P, E>> + 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<P, S, E>(stream: S) -> axum::response::Response
where
P: serde::Serialize + Send + 'static,
S: Stream<Item = Result<P, E>> + 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<P, S>(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<P, S, E>(stream: S) -> axum::response::Response
where
P: serde::Serialize + Send + 'static,
S: Stream<Item = Result<P, E>> + 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<S, E>(stream: S) -> impl Stream<Item = Result<Bytes, E>>
where
S: Stream<Item = Result<Bytes, E>> + 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<AppState>) -> 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::<Vec<u32>>(&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::<Result<Bytes, ()>>());
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,
+20 -20
View File
@@ -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()
+12 -9
View File
@@ -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<I>(
channels: Option<I>,
) -> impl Stream<Item = Result<Bytes, String>>
where
I: Stream<Item = (XtreamPlaylistItem, bool)> + Send + Unpin + 'static,
I: Stream<Item = Result<(XtreamPlaylistItem, bool), TuliproxError>> + 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<I>(channels: Option<I>) -> impl Stream<Item = Result<Bytes, String>>
where
I: Stream<Item = (M3uPlaylistItem, bool)> + Send + Unpin + 'static,
I: Stream<Item = Result<(M3uPlaylistItem, bool), TuliproxError>> + 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()
}
+7 -4
View File
@@ -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, String>(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()));
+61 -18
View File
@@ -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::<EpgChannel>(64);
let (tx, rx) = mpsc::channel::<Result<EpgChannel, String>>(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::<Arc<str>, EpgChannel>::try_new(&epg_path) else {
error!("Failed to open epg db for target {} {}", target_name, epg_path.display());
return;
let query = match BPlusTreeQuery::<Arc<str>, 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::<EpgChannel>(256);
let (channel_tx, mut channel_rx) = mpsc::channel::<Result<EpgChannel, String>>(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::<Arc<str>, EpgChannel>::try_new(&epg_path) else {
error!("Failed to open BPlusTreeQuery {}", epg_path.display());
return;
let mut query = match BPlusTreeQuery::<Arc<str>, 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.
+11 -9
View File
@@ -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<S>(xtream_iter: S) -> impl Stream<Item = Result<Bytes, String>>
where
S: Stream<Item = (String, bool)> + Send + Unpin + 'static,
S: Stream<Item = Result<(String, bool), TuliproxError>> + Send + Unpin + 'static,
{
let mapped = xtream_iter.map(move |(mut line, has_next)| {
if has_next {
line.push(',');
}
Ok::<Bytes, String>(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, String>(Bytes::from("[")) })
coalesce_byte_stream(stream::once(async { Ok::<Bytes, String>(Bytes::from("[")) })
.chain(mapped)
.chain(stream::once(async { Ok::<Bytes, String>(Bytes::from("]")) }))
.chain(stream::once(async { Ok::<Bytes, String>(Bytes::from("]")) })))
}
async fn xtream_player_api_get(
@@ -1865,6 +1865,7 @@ pub fn retry_after_delay_ms(headers: &HeaderMap) -> Option<u64> {
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),
}
}
+3 -3
View File
@@ -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);
@@ -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;
@@ -1118,7 +1118,7 @@ mod tests {
}
struct TestTransientCacheFixture {
_temp_dir: tempfile::TempDir,
temp_dir: tempfile::TempDir,
segment_cache: Arc<HlsSegmentCache>,
segment_repair: Arc<HlsSegmentRepairManager>,
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 { .. }));
@@ -496,7 +496,8 @@ fn load_metadata_retry_states_from_disk(path: &Path) -> io::Result<HashMap<TaskK
let mut result = HashMap::new();
let mut stale_keys: Vec<MetadataRetryDbKey> = Vec::new();
let mut query = BPlusTreeQuery::<MetadataRetryDbKey, MetadataRetryDbValue>::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 {
@@ -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::<String>();
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);
File diff suppressed because it is too large Load Diff
+103
View File
@@ -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<Mmap> {
// 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<io::Error> 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),
}
}
}
@@ -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<bool> {
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<FileMigrationOutcome> {
fn migrate_file_if_needed(path: &Path, roots: &[PathBuf]) -> io::Result<FileMigrationOutcome> {
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::<u8, u8>::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<bool> {
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::<MetadataRetryDbKey, MetadataRetryDbValue>(path)?,
TypedMigration::QosSnapshot => migrate_v2_typed::<String, QosSnapshotRecord>(path)?,
TypedMigration::QosCheckpoint => migrate_v2_typed::<u8, QosAggregationCheckpoint>(path)?,
TypedMigration::GeoIp => migrate_v2_typed::<u32, (u32, String)>(path)?,
TypedMigration::IdMapping => migrate_v2_typed::<u32, VirtualIdRecord>(path)?,
TypedMigration::UuidMapping => migrate_v2_typed::<UUIDType, u32>(path)?,
TypedMigration::TargetM3u => {
migrate_v2_typed_with_index::<u32, M3uPlaylistItem, u32, _>(path, |item| item.source_ordinal)?
}
TypedMigration::InputM3u => migrate_v2_typed::<Arc<str>, M3uPlaylistItem>(path)?,
TypedMigration::Library => migrate_v2_typed::<UUIDType, XtreamPlaylistItem>(path)?,
TypedMigration::TargetXtream => {
migrate_v2_typed_with_index::<u32, XtreamPlaylistItem, u32, _>(path, |item| item.source_ordinal)?
}
TypedMigration::InputXtream => migrate_v2_typed::<u32, XtreamPlaylistItem>(path)?,
TypedMigration::Epg => migrate_v2_typed::<Arc<str>, EpgChannel>(path)?,
};
Ok(())
}
fn migration_kind(path: &Path, roots: &[PathBuf]) -> Option<TypedMigration> {
let relative = roots.iter().find_map(|root| path.strip_prefix(root).ok())?;
let components = relative.iter().filter_map(OsStr::to_str).collect::<Vec<_>>();
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<SourceV, Map>(
db_path: &Path,
merge_guard_path: &Path,
map: Map,
) -> io::Result<bool>
where
SourceV: serde::Serialize + for<'de> serde::Deserialize<'de> + Clone,
Map: FnMut(SourceV) -> StoredApiUserV6,
{
if super::v3::migration::migrate_v2_typed_map::<String, SourceV, StoredApiUserV6, _>(db_path, map).is_err()
{
return Ok(false);
}
create_user_db_merge_guard(merge_guard_path)?;
Ok(true)
}
fn migrate_current_user_schema<SourceV, Map>(
db_path: &Path,
merge_guard_path: &Path,
map: Map,
) -> io::Result<bool>
where
SourceV: for<'de> serde::Deserialize<'de>,
Map: Fn(&SourceV) -> StoredApiUserV6,
{
let Ok(tree) = BPlusTree::<String, SourceV>::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::<String, StoredApiUserV5>::load(db_path) {
let mut v6_tree: BPlusTree<String, StoredApiUserV6> = 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::<String, StoredApiUserV6>(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::<StoredApiUserV5, _>(db_path, merge_guard_path, |user| {
StoredApiUserV6::from_v5(&user)
})? {
return Ok(true);
}
if migrate_legacy_user_schema::<StoredApiUserV4, _>(db_path, merge_guard_path, |user| {
StoredApiUserV6::from_v4(&user)
})? {
return Ok(true);
}
if migrate_legacy_user_schema::<StoredApiUserV3, _>(db_path, merge_guard_path, |user| {
StoredApiUserV6::from_v3(&user)
})? {
return Ok(true);
}
if migrate_legacy_user_schema::<StoredApiUserV2, _>(db_path, merge_guard_path, |user| {
StoredApiUserV6::from_v2(&user)
})? {
return Ok(true);
}
if migrate_legacy_user_schema::<StoredApiUserV1, _>(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::<String, StoredApiUserV6>::load(db_path).is_ok() {
return Ok(false);
}
if let Ok(tree) = BPlusTree::<String, StoredApiUserV4>::load(db_path) {
let mut v6_tree: BPlusTree<String, StoredApiUserV6> = 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::<StoredApiUserV5, _>(db_path, merge_guard_path, StoredApiUserV6::from_v5)? {
return Ok(true);
}
if let Ok(tree) = BPlusTree::<String, StoredApiUserV3>::load(db_path) {
let mut v6_tree: BPlusTree<String, StoredApiUserV6> = 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::<StoredApiUserV4, _>(db_path, merge_guard_path, StoredApiUserV6::from_v4)? {
return Ok(true);
}
if let Ok(tree) = BPlusTree::<String, StoredApiUserV2>::load(db_path) {
let mut v6_tree: BPlusTree<String, StoredApiUserV6> = 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::<StoredApiUserV3, _>(db_path, merge_guard_path, StoredApiUserV6::from_v3)? {
return Ok(true);
}
if let Ok(tree) = BPlusTree::<String, StoredApiUserV1>::load(db_path) {
let mut v6_tree: BPlusTree<String, StoredApiUserV6> = 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::<StoredApiUserV2, _>(db_path, merge_guard_path, StoredApiUserV6::from_v2)? {
return Ok(true);
}
if migrate_current_user_schema::<StoredApiUserV1, _>(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<AllStartupMigrationStats> {
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<PathBuf> {
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::<u32, (u32, String)>::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<str>, 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<str>, EpgChannel);
empty_v2!("target/xtream/epg.db", Arc<str>, 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(&current_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<String, StoredApiUserV2> = BPlusTree::new();
let mut v2_tree: super::super::v2::BPlusTree<String, StoredApiUserV2> =
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<String, StoredApiUserV6> = BPlusTree::new();
let mut v6_tree: super::super::v2::BPlusTree<String, StoredApiUserV6> =
super::super::v2::BPlusTree::new();
v6_tree.insert(
"erin".to_string(),
StoredApiUserV6 {
+11
View File
@@ -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::*;
@@ -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<String>) -> 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<u32> {
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<u64> {
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<u64> {
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<SortKey, K> {
output: BufWriter<File>,
database_id: [u8; 16],
generation: u64,
count: u64,
body: Vec<u8>,
_marker: PhantomData<(SortKey, K)>,
}
impl<SortKey, K> Writer<SortKey, K>
where
SortKey: Serialize,
K: Serialize,
{
pub(crate) fn new(path: &Path, database_id: [u8; 16], generation: u64) -> io::Result<Self> {
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<u64> {
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<SortKey, K> {
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<SortKey, K> {
input: BufReader<File>,
remaining_entries: u64,
remaining_bytes: u64,
failed: bool,
body: Vec<u8>,
_marker: PhantomData<(SortKey, K)>,
}
impl<SortKey, K> Reader<SortKey, K>
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<Self> {
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<Option<Entry<SortKey, K>>> {
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<Option<BorrowedEntry<'_, SortKey, K>>> {
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<Option<BorrowedEntry<'_, SortKey, K>>> {
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<K, V, SortKey> {
reader: Reader<SortKey, K>,
query: BPlusTreeQuery<K, V>,
previous_sort_key: Option<SortKey>,
finished: bool,
}
impl<K, V, SortKey> OwnedIterator<K, V, SortKey>
where
K: Ord + for<'de> Deserialize<'de>,
V: for<'de> Deserialize<'de>,
SortKey: Ord + for<'de> Deserialize<'de>,
{
pub(crate) fn open(query: BPlusTreeQuery<K, V>, index_path: &Path) -> io::Result<Self> {
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<K, V, SortKey> Iterator for OwnedIterator<K, V, SortKey>
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<Self::Item> {
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::<String, u32>::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::<String, u32>::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::<String, u32>::open(&path, [0; 16], 9).is_err());
assert!(v4::Reader::<String, u32>::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::<u32, u32>::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::<u32, u32>::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::<u32, u32>::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::<u32, u32>::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::<u32, String>::try_new(&database)?;
let entries = query.collect_with_locators()?;
let (database_id, generation) = query.snapshot_identity();
drop(query);
let mut writer = v4::Writer::<u32, u32>::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::<u32, String>::try_new(&database)?;
let mut iterator = v4::OwnedIterator::<u32, String, u32>::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::<u32, String>::try_new(&database)?;
let entries = query.collect_with_locators()?;
let (database_id, generation) = query.snapshot_identity();
drop(query);
let mut writer = v4::Writer::<u32, u32>::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::<u32, String>::try_new(&database)?;
let mut iterator = v4::OwnedIterator::<u32, String, u32>::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(())
}
}
@@ -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::<u32, String>::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<Vec<Duration>> {
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::<Vec<_>>().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::<Vec<_>>();
let v2_point = comparison_runs("v2 point lookup", || {
let mut query = super::v2::BPlusTreeQuery::<String, ComparisonValue>::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::<String, ComparisonValue>::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::<String, ComparisonValue>::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::<String, ComparisonValue>::try_new(&v2_path)?;
let count = query
.range_iter(Bound::Unbounded, Bound::Unbounded)
.collect::<Result<Vec<_>, _>>()
.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::<String, ComparisonValue>::try_new(&v3_path)?;
black_box(query.iter().collect::<std::io::Result<Vec<_>>>()?.len());
Ok(())
})?;
comparison_runs("v2 sorted full scan", || {
let mut query = super::v2::BPlusTreeQuery::<String, ComparisonValue>::try_new(&v2_path)?;
let mut values = query
.range_iter(Bound::Unbounded, Bound::Unbounded)
.collect::<Result<Vec<_>, _>>()
.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::<String, ComparisonValue>::try_new(&v3_path)?;
let iterator = super::sorted_index::v4::OwnedIterator::<String, ComparisonValue, u32>::open(query, &index_path)?;
black_box(iterator.collect::<std::io::Result<Vec<_>>>()?.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::<Vec<_>>();
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::<Vec<_>>();
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::<String, ComparisonValue>::try_new(&path)?;
let references = updates.iter().map(|(key, value)| (key, value)).collect::<Vec<_>>();
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::<String, ComparisonValue>::try_new(&compact_path)?;
let growing_refs = growing.iter().map(|(key, value)| (key, value)).collect::<Vec<_>>();
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(())
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -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<Option<u32>> {
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<Option<PathBuf>> {
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<BPlusTreeMetadata> {
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<K, V>(source: &Path) -> io::Result<MigrationValidation>
where
K: Ord + Serialize + for<'de> Deserialize<'de> + Clone,
V: Serialize + for<'de> Deserialize<'de> + Clone,
{
migrate_v2_typed_inner::<K, V, V, _, _, _>(
source,
std::convert::identity,
BPlusTree::store_verified,
|_, _| Ok(()),
false,
)
}
pub(crate) fn migrate_v2_typed_with_index<K, V, SortKey, F>(
source: &Path,
sort_key: F,
) -> io::Result<MigrationValidation>
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::<K, V, V, _, _, _>(
source,
std::convert::identity,
|tree, destination| tree.store_with_index_verified(destination, sort_key),
|destination, entries| {
let query = super::BPlusTreeQuery::<K, V>::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::<K, V, SortKey>::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<K, SourceV, DestinationV, Map>(
source: &Path,
map: Map,
) -> io::Result<MigrationValidation>
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::<K, SourceV, DestinationV, _, _, _>(
source,
map,
BPlusTree::store_verified,
|_, _| Ok(()),
false,
)
}
fn migrate_v2_typed_inner<K, SourceV, DestinationV, Map, Store, Validate>(
source: &Path,
mut map: Map,
store: Store,
validate: Validate,
indexed: bool,
) -> io::Result<MigrationValidation>
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<K, DestinationV>, &Path) -> io::Result<super::tree::VerificationReport>,
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::<K, SourceV>::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::<K, DestinationV>::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<PathBuf> {
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::<u32, String>(&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::<u32, String>(&path)?;
assert_eq!(validation.entries, 2);
assert_eq!(validation.generation, 1);
assert_eq!(storage_version(&path)?, Some(3));
let mut query = super::super::BPlusTreeQuery::<u32, String>::try_new(&path)?;
assert_eq!(query.iter().collect::<io::Result<Vec<_>>>()?, 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::<u32, String>(&path)?;
assert_eq!(validation.entries, 80);
let mut query = super::super::BPlusTreeQuery::<u32, String>::try_new(&path)?;
assert_eq!(query.iter().collect::<io::Result<Vec<_>>>()?.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::<u32, String>(&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::<u32, String>(&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::<u32, String>(&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::<u32, String, usize, _>(&path, String::len)?;
assert_eq!(validation.entries, 2);
let query = super::super::BPlusTreeQuery::<u32, String>::try_new(&path)?;
let index = crate::repository::storage::get_file_path_for_db_index(&path);
let values = crate::repository::bplustree::sorted_index::v4::OwnedIterator::<u32, String, usize>::open(query, &index)
?.collect::<io::Result<Vec<_>>>()?;
assert_eq!(values, vec![(2, String::from("a")), (1, String::from("bbb"))]);
Ok(())
}
}
+140
View File
@@ -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<u8> {
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::<io::Result<Vec<_>>>()?.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::<Vec<_>>(), 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::<io::Result<Vec<_>>>()?, expected);
let mut query = Query::try_new(&path)?;
assert_eq!(
query
.range_iter(Bound::Included(&20), Bound::Included(&30))
.collect::<io::Result<Vec<_>>>()?,
expected[1..].to_vec()
);
Ok(())
}
};
}
storage_contract!(
v3_storage_contract,
BPlusTree<u32, Vec<u8>>,
BPlusTreeQuery<u32, Vec<u8>>,
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::<u32, Vec<u8>>::try_new(&path)?;
let actual = query
.range_iter(Bound::Unbounded, Bound::Unbounded)
.collect::<Result<Vec<_>, _>>()
.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::<String, u32>::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::<String, u32>::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::<u32, Vec<u8>>::try_new(&path)?;
assert_eq!(verify_full(&mut query)?.live_entries, 1);
Ok(())
}
}
+689
View File
@@ -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<usize> = 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<usize> {
offset.checked_add(length).ok_or_else(|| io::Error::new(kind, "page offset overflow"))
}
fn slot_base(page_type: PageType) -> io::Result<usize> {
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<B> {
bytes: B,
header: PageHeader,
page_id: u64,
next_page_id: u64,
overflow_payload_length: Option<u16>,
}
#[derive(Clone, Copy)]
pub(crate) struct PageValidation {
header: PageHeader,
page_id: u64,
next_page_id: u64,
overflow_payload_length: Option<u16>,
}
impl<B: AsRef<[u8]>> SlottedPage<B> {
pub(crate) fn open(bytes: B, page_id: u64, next_page_id: u64) -> io::Result<Self> {
#[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<Self> {
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<Slot> {
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<std::ops::Range<usize>> {
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<B: AsRef<[u8]>>(page: &SlottedPage<B>) -> 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<B>,
next: usize,
}
impl<'a, B: AsRef<[u8]>> Iterator for Cells<'a, B> {
type Item = io::Result<&'a [u8]>;
fn next(&mut self) -> Option<Self::Item> {
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<usize>) {
let remaining = self.len();
(remaining, Some(remaining))
}
}
impl<B: AsRef<[u8]>> ExactSizeIterator for Cells<'_, B> {
fn len(&self) -> usize { usize::from(self.page.header.cell_count).saturating_sub(self.next) }
}
impl<B: AsRef<[u8]>> FusedIterator for Cells<'_, B> {}
impl<B: AsRef<[u8]> + AsMut<[u8]>> SlottedPage<B> {
pub(crate) fn rebuild_ordered<'a, I>(&mut self, cells: I) -> io::Result<()>
where
I: IntoIterator<Item = &'a [u8]>,
{
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<T>(result: io::Result<T>) -> 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<T>(result: io::Result<T>) -> 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::<io::Result<Vec<_>>>()?, [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::<io::Result<Vec<_>>>()?, [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::<io::Result<Vec<_>>>()?, [b"a".as_slice(), b"bc".as_slice()]);
Ok(())
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+7 -1
View File
@@ -468,7 +468,13 @@ mod tests {
.expect("case-preserving EPG should be written");
let mut query = BPlusTreeQuery::<Arc<str>, EpgChannel>::try_new(&path).expect("EPG DB should open");
let stored_ids = query.iter().map(|(_, channel)| channel.id).collect::<Vec<_>>();
let stored_ids = query
.iter()
.collect::<std::io::Result<Vec<_>>>()
.expect("EPG entries should be readable")
.into_iter()
.map(|(_, channel)| channel.id)
.collect::<Vec<_>>();
assert_eq!(stored_ids.iter().map(AsRef::as_ref).collect::<Vec<_>>(), vec!["Z.Channel", "a.channel"]);
}
+33 -21
View File
@@ -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<UUIDType, XtreamPlaylistItem> = HashMap::new();
if library_path.exists() {
if let Ok(mut query) = BPlusTreeQuery::<UUIDType, XtreamPlaylistItem>::try_new(&library_path) {
for (uuid, item) in query.iter() {
existing_by_uuid.insert(uuid, item);
}
let mut query = BPlusTreeQuery::<UUIDType, XtreamPlaylistItem>::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<AppConfig>, lib_path: &Path) -> Result<Vec<PlaylistGroup>, 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<AppConfig>, lib_
let groups = task::spawn_blocking(move || -> Result<Vec<PlaylistGroup>, TuliproxError> {
let _guard = file_lock;
let mut groups: IndexMap<CategoryKey, PlaylistGroup> = IndexMap::new();
if let Ok(mut query) = BPlusTreeQuery::<UUIDType, XtreamPlaylistItem>::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::<UUIDType, XtreamPlaylistItem>::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,
+34 -12
View File
@@ -23,7 +23,7 @@ use tokio::task;
create_bitset!(u8, M3uPlaylistIteratorFlags, MaskRedirectUrl, IncludeTypeInUrl, RewriteResource);
pub struct M3uPlaylistIterator {
inner: LockedReceiverStream<(M3uPlaylistItem, bool)>,
inner: LockedReceiverStream<Result<(M3uPlaylistItem, bool), TuliproxError>>,
}
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::<Result<(M3uPlaylistItem, bool), TuliproxError>>(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::<u32, M3uPlaylistItem, u32>(
@@ -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<Option<Self::Item>> {
Pin::new(&mut self.inner).poll_next(cx)
@@ -363,19 +371,20 @@ impl M3uPlaylistM3uTextIterator {
}
impl Stream for M3uPlaylistM3uTextIterator {
type Item = String;
type Item = Result<String, String>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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::<Result<(M3uPlaylistItem, bool), shared::error::TuliproxError>>(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());
}
}
+22 -21
View File
@@ -384,7 +384,7 @@ async fn iter_raw_m3u_playlist<SortKey, ItemKey>(
) -> Option<Box<dyn Stream<Item=Result<M3uPlaylistItem, TuliproxError>> + 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<Vec<PlaylistGroup>, TuliproxError> {
let _guard = file_lock;
let mut groups: IndexMap<CategoryKey, PlaylistGroup> = IndexMap::new();
if let Ok(mut query) = BPlusTreeQuery::<Arc<str>, 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::<Arc<str>, 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())
})
-7
View File
@@ -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;
@@ -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?;
+110 -50
View File
@@ -457,8 +457,18 @@ pub struct XtreamDiskPlaylistSource {
series: Option<XtreamQueryHandle>,
}
fn log_and_skip_btree_error<T>(entry: std::io::Result<T>) -> Option<T> {
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<AppConfig>, storage_path: &Path) -> Self {
pub(crate) async fn new(app_config: &Arc<AppConfig>, storage_path: &Path) -> Result<Self, TuliproxError> {
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::<u32, XtreamPlaylistItem>(&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::<u32, XtreamPlaylistItem>(&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::<u32, XtreamPlaylistItem>(&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<Q>(query: &mut Option<(BPlusTreeQuery<u32, XtreamPlaylistItem>, Q)>, groups: &mut HashSet<Arc<str>>) {
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<dyn Iterator<Item=PlaylistItem> + 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<dyn Iterator<Item=Cow<'a, PlaylistItem>> + 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<dyn Iterator<Item=XtreamPlaylistItem> + 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<AppConfig>, file_path: &Path) -> Self {
pub(crate) async fn new(app_config: &Arc<AppConfig>, file_path: &Path) -> Result<Self, TuliproxError> {
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<XtreamCluster>) -> usize {
let mut groups = HashSet::<(XtreamCluster, Arc<str>)>::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<dyn Iterator<Item=PlaylistItem> + 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<dyn Iterator<Item=Cow<'a, PlaylistItem>> + 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<str>), 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<K, P>(app_config: &Arc<AppConfig>, file_path: &Path) -> Option<(BPlusTreeQuery<K, P>, FileReadGuard)>
async fn load_bplustree_query<K, P>(
app_config: &Arc<AppConfig>,
file_path: &Path,
) -> Result<Option<(BPlusTreeQuery<K, P>, 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::<u32, XtreamPlaylistItem>(&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::<u32, XtreamPlaylistItem>(
&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")
));
}
+115 -9
View File
@@ -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<K, V, SortKey> {
Sorted {
iterator: super::bplustree::sorted_index::v4::OwnedIterator<K, V, SortKey>,
fallback_path: std::path::PathBuf,
yielded: bool,
},
Unsorted(BPlusTreeDiskIteratorOwned<K, V>),
}
impl<K, V, SortKey> Iterator for PlaylistIteratorReader<K, V, SortKey>
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<Self::Item> {
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<T> {
@@ -51,7 +94,7 @@ impl<T> Stream for LockedReceiverStream<T> {
/// 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<K, V, SortKey>(
pub(crate) fn open_playlist_reader<K, V, SortKey>(
path: &Path,
index_path: &Path,
sorted_err_prefix: Option<&str>,
@@ -59,7 +102,7 @@ pub fn open_playlist_reader<K, V, SortKey>(
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::<K, V>::try_new(path)
.map_err(|err| TuliproxError::Config(format!(
@@ -68,11 +111,13 @@ where
)))?;
if index_path.exists() {
match SortedIndexReader::<SortKey, K>::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::<K, V>::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::<u32, String, usize>(&database, &index, None)
.map_err(io::Error::other)?;
assert_eq!(reader.collect::<io::Result<Vec<_>>>()?, 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::<u32, String>::try_new(&database)?;
let entries = query.collect_with_locators()?;
let (database_id, generation) = query.snapshot_identity();
drop(query);
let mut writer = v4::Writer::<u32, u32>::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::<u32, String, u32>(&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(())
}
}
@@ -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(())
}
-788
View File
@@ -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<u8> {
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<R: Read>(reader: &mut R) -> io::Result<Self> {
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<SortKey, K> {
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<SortKey, K> {
writer: BufWriter<File>,
count: u64,
_marker: PhantomData<(SortKey, K)>,
}
impl<SortKey, K> SortedIndexWriter<SortKey, K>
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<Self> {
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<u64> {
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<SortKey, K> {
reader: BufReader<File>,
remaining: u64,
_marker: PhantomData<(SortKey, K)>,
}
impl<SortKey, K> SortedIndexReader<SortKey, K>
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<Self> {
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<Option<IndexEntry<SortKey, K>>> {
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<SortKey, K> Iterator for SortedIndexReader<SortKey, K>
where
SortKey: for<'de> Deserialize<'de>,
K: for<'de> Deserialize<'de>,
{
type Item = io::Result<IndexEntry<SortKey, K>>;
fn next(&mut self) -> Option<Self::Item> {
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<K, V, SortKey> {
index_reader: SortedIndexReader<SortKey, K>,
tree_file: Option<BufReader<File>>,
mmap: Option<memmap2::Mmap>,
filepath: PathBuf,
block_cache: IndexMap<u64, Vec<u8>>,
live_key_query: Option<BPlusTreeQuery<K, V>>,
_marker: PhantomData<V>,
}
const CACHE_CAPACITY: usize = 8;
impl<K, V, SortKey> BPlusTreeSortedIteratorOwned<K, V, SortKey>
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<SortKey, K>,
filepath: PathBuf,
tree_file: Option<BufReader<File>>,
mmap: Option<memmap2::Mmap>,
) -> 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<BufReader<File>>) -> io::Result<Self> {
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<BufReader<File>>,
mmap: Option<memmap2::Mmap>,
) -> io::Result<Self> {
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<BufReader<File>>,
index_path: &Path,
) -> io::Result<Self> {
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<BufReader<File>>,
mmap: Option<memmap2::Mmap>,
index_path: &Path,
) -> io::Result<Self> {
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<V> {
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<V> {
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<V> {
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<K, V, SortKey> BPlusTreeSortedIteratorOwned<K, V, SortKey>
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::<K, V>::try_new(&self.filepath)?;
self.live_key_query = Some(query);
Ok(())
}
fn is_live_key(&mut self, key: &K) -> io::Result<bool> {
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<K, V, SortKey> Iterator for BPlusTreeSortedIteratorOwned<K, V, SortKey>
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<Self::Item> {
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::<String, u32>::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::<String, u32>::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::<String, u32>::new(&path).unwrap();
let count = writer.finish().unwrap();
assert_eq!(count, 0);
let reader = SortedIndexReader::<String, u32>::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"));
}
}
+5 -9
View File
@@ -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,
+1 -1
View File
@@ -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,
@@ -17,7 +17,7 @@ use tokio::sync::mpsc;
use tokio::task;
pub struct XtreamPlaylistIterator {
inner: LockedReceiverStream<(XtreamPlaylistItem, bool)>,
inner: LockedReceiverStream<Result<(XtreamPlaylistItem, bool), TuliproxError>>,
}
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::<Result<(XtreamPlaylistItem, bool), TuliproxError>>(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::<Result<(XtreamPlaylistItem, bool), TuliproxError>>(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::<u32, XtreamPlaylistItem, u32>(
@@ -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<Option<Self::Item>> {
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<Option<Self::Item>> {
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());
}
}
+37 -20
View File
@@ -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<Box<dyn Stream<Item=XtreamPlaylistItem> + Send + Unpin>> {
pub async fn iter_raw_xtream_target_playlist(app_config: &AppConfig, target: &ConfigTarget, cluster: XtreamCluster) -> Option<Box<dyn Stream<Item=Result<XtreamPlaylistItem, TuliproxError>> + 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<Box<dyn Stream<Item=XtreamPlaylistItem> + Send + Unpin>> {
pub async fn iter_raw_xtream_input_playlist(app_config: &AppConfig, input: &ConfigInput, cluster: XtreamCluster) -> Option<Box<dyn Stream<Item=Result<XtreamPlaylistItem, TuliproxError>> + 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<Box<dyn Stream<Item=XtreamPlaylistItem> + Send + Unpin>> {
async fn iter_raw_xtream_playlist(app_config: &AppConfig, xtream_path: &Path) -> Option<Box<dyn Stream<Item=Result<XtreamPlaylistItem, TuliproxError>> + 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::<XtreamPlaylistItem>(256);
let (tx, rx) = mpsc::channel::<Result<XtreamPlaylistItem, TuliproxError>>(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::<u32, XtreamPlaylistItem, u32>(
@@ -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<dyn Stream<Item=XtreamPlaylistItem> + Send + Unpin> =
let stream: Box<dyn Stream<Item=Result<XtreamPlaylistItem, TuliproxError>> + 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::<u32, XtreamPlaylistItem>::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<AppConfig>, 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::<u32, XtreamPlaylistItem>::try_new(&xtream_path) {
for (_, doc) in query.iter() {
entries.insert(doc.provider_id, doc);
}
let mut query = BPlusTreeQuery::<u32, XtreamPlaylistItem>::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<AppConfig>, 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<AppConfig>, storage_pat
let items = tokio::task::spawn_blocking(move || -> Result<Vec<XtreamPlaylistItem>, TuliproxError> {
let _guard = file_lock;
let mut items = Vec::new();
if let Ok(mut query) = BPlusTreeQuery::<u32, XtreamPlaylistItem>::try_new(&xtream_path) {
for (_, item) in query.iter() {
items.push(item);
}
let mut query = BPlusTreeQuery::<u32, XtreamPlaylistItem>::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)
})
+9 -1
View File
@@ -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 {
+3 -1
View File
@@ -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);
+2 -3
View File
@@ -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();