mirror of
https://github.com/euzu/tuliprox.git
synced 2026-10-06 16:02:36 +02:00
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:
@@ -1,4 +1,5 @@
|
||||
.tools
|
||||
.worktrees
|
||||
dist
|
||||
/resources/*.ts
|
||||
/release
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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()));
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
+367
-155
@@ -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(¤t_roots, temp_a.path())?;
|
||||
assert!(!stats.skipped_by_marker);
|
||||
assert!(!stats.marker_was_current);
|
||||
assert_eq!(stats.migrated_files, 1);
|
||||
|
||||
Ok(())
|
||||
@@ -1099,7 +1309,8 @@ mod tests {
|
||||
let db_path = temp.path().join(storage_const::API_USER_DB_FILE);
|
||||
let merge_guard_path = user_db_merge_guard_path(temp.path());
|
||||
|
||||
let mut v2_tree: BPlusTree<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 {
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
+212
-4
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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"]);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
})
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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")
|
||||
));
|
||||
}
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user