Files
aria2-rust-pro/crates/aria2-rust-pro-core/src/engine/session_persistence.rs
T

622 lines
21 KiB
Rust

use super::{
BtFileInfo, BtRuntimeState, ControlFileVersion, ControlMetadata, DownloadEngine, DownloadId,
DownloadRegistry, DownloadStatus, OptionKey, OptionValue, Path, PathBuf, RequestContext,
RequestGroup, SessionFile, SessionFileEntry, StoredDownloadFile, StoredPieceIndex,
StoredPieceState, infer_total_length, usize_to_u32,
};
use std::collections::BTreeMap;
use crate::{ResumeState, RetryAttempt};
impl DownloadEngine {
/// Builds a serializable session file from every persistable group in the registry.
pub(super) fn build_session_file(&self, session_path: &Path) -> SessionFile {
let entries = self
.registry
.groups
.values()
.filter(|group| should_persist_group(*group.status()))
.map(|group| build_session_entry(group, self.session.global_options(), session_path))
.collect();
SessionFile { entries }
}
/// Reconstructs the in-memory registry and waiting queue from a persisted session file.
pub(super) fn rebuild_registry_from_session_file(
&mut self,
session_path: &Path,
session_file: SessionFile,
) {
let mut registry = DownloadRegistry::new();
let mut reserved_queue = Vec::new();
let mut max_gid = 0_u64;
for entry in session_file.entries {
let mut context = RequestContext::new(entry.uri.clone());
if entry.uris.is_empty() {
context.replace_uris(vec![entry.uri.clone()]);
} else {
context.replace_uris(entry.uris.clone());
}
let mut group = if let Some(gid) = DownloadId::parse_hex(&entry.gid) {
max_gid = max_gid.max(gid.as_u64());
RequestGroup::with_context(gid, context)
} else {
let gid = registry.allocate_gid();
max_gid = max_gid.max(gid.as_u64());
RequestGroup::with_context(gid, context)
};
restore_group_metadata(&mut group, entry.metadata);
let control_path = entry
.metadata_path
.unwrap_or_else(|| control_path_for_session_entry(session_path, group.gid()));
if let Ok(control) = aria2_rust_pro_storage::read_aria2_control_file(&control_path) {
restore_control_metadata(&mut group, control);
}
if matches!(
group.status(),
DownloadStatus::Waiting | DownloadStatus::Paused
) {
reserved_queue.push(group.gid());
}
registry.insert(group);
}
registry.next_gid = max_gid.saturating_add(1).max(1);
self.registry = registry;
self.reserved_queue = reserved_queue;
}
}
/// Returns whether a group status should be persisted into session artifacts.
pub(super) fn should_persist_group(status: DownloadStatus) -> bool {
!matches!(status, DownloadStatus::Complete | DownloadStatus::Removed)
}
/// Builds one persisted session entry from a request group.
fn build_session_entry(
group: &RequestGroup,
global_options: &crate::session::GlobalOptions,
session_path: &Path,
) -> SessionFileEntry {
let target_path = resolve_target_path(group, global_options);
let metadata = Some(build_group_metadata(group));
SessionFileEntry {
gid: group.gid().to_string(),
uri: group.uri().to_owned(),
uris: group.uris().to_vec(),
target_path,
metadata_path: Some(control_path_for_session_entry(session_path, group.gid())),
metadata,
}
}
/// Serializes selected request-group runtime metadata into session-file string fields.
fn build_group_metadata(group: &RequestGroup) -> BTreeMap<String, String> {
let mut metadata = BTreeMap::from([
(
"status".to_owned(),
group.status().as_rpc_status().to_owned(),
),
(
"num_connections".to_owned(),
group.num_connections().to_string(),
),
(
"download_speed".to_owned(),
group.download_speed().to_string(),
),
(
"upload_length".to_owned(),
group.upload_length().to_string(),
),
(
"completed_length".to_owned(),
group.completed_length().to_string(),
),
("retry_count".to_owned(), group.retry_count().to_string()),
]);
if let Some(resume_state) = group.resume_state() {
metadata.insert(
"resume_state".to_owned(),
encode_resume_state_metadata(resume_state),
);
}
if !group.retry_attempts().is_empty() {
metadata.insert(
"retry_attempts".to_owned(),
encode_retry_attempts_metadata(group.retry_attempts()),
);
}
for (key, value) in group.options().entries() {
metadata.insert(format!("opt.{}", key.as_str()), option_value_text(value));
}
if let Some(bt) = group.bt() {
metadata.insert("bt.info_hash".to_owned(), bt.info_hash.clone());
metadata.insert("bt.metadata_only".to_owned(), bt.metadata_only.to_string());
if let Some(name) = &bt.name {
metadata.insert("bt.name".to_owned(), escape_metadata_field(name));
}
if let Some(magnet_uri) = &bt.magnet_uri {
metadata.insert(
"bt.magnet_uri".to_owned(),
escape_metadata_field(magnet_uri),
);
}
if let Some(creation_date) = &bt.creation_date {
metadata.insert(
"bt.creation_date".to_owned(),
escape_metadata_field(creation_date),
);
}
if let Some(comment) = &bt.comment {
metadata.insert("bt.comment".to_owned(), escape_metadata_field(comment));
}
metadata.insert("bt.files_count".to_owned(), bt.files.len().to_string());
for (index, file) in bt.files.iter().enumerate() {
metadata.insert(
format!("bt.file.{index}.path"),
escape_metadata_field(&file.path),
);
metadata.insert(format!("bt.file.{index}.length"), file.length.to_string());
metadata.insert(
format!("bt.file.{index}.selected"),
file.selected.to_string(),
);
if let Some(piece_offset) = file.piece_offset {
metadata.insert(
format!("bt.file.{index}.piece_offset"),
piece_offset.to_string(),
);
}
}
}
metadata
}
/// Resolves the output path that should be associated with a request group.
pub(super) fn resolve_target_path(
group: &RequestGroup,
global_options: &crate::session::GlobalOptions,
) -> PathBuf {
let dir = group
.options()
.get(&OptionKey::from("dir"))
.or_else(|| global_options.get(&OptionKey::from("dir")))
.and_then(OptionValue::as_text)
.map(PathBuf::from);
let file_name = group
.options()
.get(&OptionKey::from("out"))
.and_then(OptionValue::as_text)
.map(str::to_owned)
.or_else(|| uri_file_name(group.uri()))
.unwrap_or_else(|| group.gid().to_string());
match dir {
Some(dir) => dir.join(file_name),
None => PathBuf::from(file_name),
}
}
/// Extracts a best-effort file name from a URI path component.
fn uri_file_name(uri: &str) -> Option<String> {
let trimmed = uri
.split(['?', '#'])
.next()
.unwrap_or(uri)
.trim_end_matches('/');
let candidate = trimmed.rsplit('/').next()?;
if candidate.is_empty() {
None
} else {
Some(candidate.to_owned())
}
}
/// Resolves the per-download control-file path relative to a session file path.
pub(super) fn control_path_for_session_entry(session_path: &Path, gid: DownloadId) -> PathBuf {
let parent = session_path
.parent()
.map_or_else(|| PathBuf::from("."), Path::to_path_buf);
parent.join("control").join(format!("{gid}.aria2"))
}
/// Builds persisted control-file metadata for a request group.
pub(super) fn build_control_metadata(
group: &RequestGroup,
target_path: &Path,
piece_length: u64,
) -> ControlMetadata {
let resolved_piece_length = group.piece_length().max(piece_length);
let inferred_total_length = infer_total_length(group, resolved_piece_length);
ControlMetadata {
version: ControlFileVersion::CURRENT,
files: vec![StoredDownloadFile {
path: target_path.to_path_buf(),
length: inferred_total_length,
piece_length: resolved_piece_length,
}],
checksums: Vec::new(),
completed_length: group.completed_length(),
retry_count: group.retry_count(),
last_error: group
.retry_attempts()
.last()
.and_then(|attempt| attempt.error.clone()),
last_error_at_unix_ms: None,
last_retry_at_unix_ms: None,
next_retry_at_unix_ms: None,
consecutive_failure_count: Some(usize_to_u32(group.retry_attempts().len())),
active_segment_count: Some(group.num_connections()),
resume_verified_at_unix_ms: None,
resume_generation: group
.resume_state()
.and_then(|resume_state| resume_state.persisted.then_some(1)),
piece_states: group
.piece_map()
.iter()
.map(|(piece, state)| {
(
StoredPieceIndex(piece.0),
map_piece_state_to_storage(*state),
)
})
.collect(),
}
}
/// Maps in-memory piece states into storage-layer piece states.
fn map_piece_state_to_storage(state: crate::piece::PieceState) -> StoredPieceState {
match state {
crate::piece::PieceState::Verified => StoredPieceState::Verified,
crate::piece::PieceState::Queued | crate::piece::PieceState::Downloading => {
StoredPieceState::InFlight
}
crate::piece::PieceState::Pending
| crate::piece::PieceState::Missing
| crate::piece::PieceState::Skipped => StoredPieceState::Pending,
}
}
/// Maps storage-layer piece states back into in-memory piece states.
fn map_piece_state_from_storage(state: StoredPieceState) -> crate::piece::PieceState {
match state {
StoredPieceState::Pending => crate::piece::PieceState::Pending,
StoredPieceState::InFlight => crate::piece::PieceState::Downloading,
StoredPieceState::Verified => crate::piece::PieceState::Verified,
StoredPieceState::Failed => crate::piece::PieceState::Missing,
}
}
/// Serializes an option value into the session metadata text format.
fn option_value_text(value: &OptionValue) -> String {
match value {
OptionValue::Bool(value) => value.to_string(),
OptionValue::Int(value) => value.to_string(),
OptionValue::UInt(value) => value.to_string(),
OptionValue::Text(value) => value.clone(),
OptionValue::List(value) => value.join(","),
OptionValue::Map(value) => value
.iter()
.map(|(key, value)| format!("{key}={value}"))
.collect::<Vec<_>>()
.join(","),
OptionValue::Empty => String::default(),
}
}
/// Restores request-group runtime metadata from persisted session metadata fields.
fn restore_group_metadata(group: &mut RequestGroup, metadata: Option<BTreeMap<String, String>>) {
let Some(metadata) = metadata else {
return;
};
let mut bt = BtRuntimeState::default();
let mut bt_seen = false;
let mut bt_files: BTreeMap<usize, BtFileInfo> = BTreeMap::new();
for (key, value) in metadata {
if restore_group_metadata_field(group, &key, &value) {
continue;
}
if restore_bt_metadata_field(&mut bt, &mut bt_files, &key, &value) {
bt_seen = true;
}
}
if bt_seen {
bt.files = bt_files.into_values().collect();
group.set_bt(bt);
}
}
/// Restores one non-BitTorrent metadata field onto a request group.
fn restore_group_metadata_field(group: &mut RequestGroup, key: &str, value: &str) -> bool {
match key {
"status" => {
if let Some(status) = parse_status(value) {
group.set_status(status);
}
true
}
"num_connections" => {
if let Ok(parsed) = value.parse::<u32>() {
group.set_num_connections(parsed);
}
true
}
"download_speed" => {
if let Ok(parsed) = value.parse::<u64>() {
group.set_download_speed(parsed);
}
true
}
"upload_length" => {
if let Ok(parsed) = value.parse::<u64>() {
group.set_upload_length(parsed);
}
true
}
"completed_length" => {
if let Ok(parsed) = value.parse::<u64>() {
group.set_completed_length(parsed);
}
true
}
"retry_count" => {
if let Ok(parsed) = value.parse::<u32>() {
group.set_retry_count(parsed);
}
true
}
"resume_state" => {
if let Some(parsed) = decode_resume_state_metadata(value) {
group.set_resume_state(parsed);
}
true
}
"retry_attempts" => {
group.set_retry_attempts(decode_retry_attempts_metadata(value));
true
}
_ => key.strip_prefix("opt.").is_some_and(|option_key| {
group.set_option(option_key.to_owned(), value.to_owned());
true
}),
}
}
/// Restores one `BitTorrent` metadata field.
fn restore_bt_metadata_field(
bt: &mut BtRuntimeState,
bt_files: &mut BTreeMap<usize, BtFileInfo>,
key: &str,
value: &str,
) -> bool {
match key {
"bt.info_hash" => {
value.clone_into(&mut bt.info_hash);
true
}
"bt.metadata_only" => {
bt.metadata_only = value.parse::<bool>().unwrap_or(false);
true
}
"bt.name" => {
bt.name = Some(unescape_metadata_field(value));
true
}
"bt.magnet_uri" => {
bt.magnet_uri = Some(unescape_metadata_field(value));
true
}
"bt.creation_date" => {
bt.creation_date = Some(unescape_metadata_field(value));
true
}
"bt.comment" => {
bt.comment = Some(unescape_metadata_field(value));
true
}
_ => key
.strip_prefix("bt.file.")
.is_some_and(|rest| restore_bt_file_metadata_field(bt_files, rest, value)),
}
}
/// Restores one `BitTorrent` file metadata field.
fn restore_bt_file_metadata_field(
bt_files: &mut BTreeMap<usize, BtFileInfo>,
rest: &str,
value: &str,
) -> bool {
let mut parts = rest.split('.');
let Some(index_raw) = parts.next() else {
return false;
};
let Some(field) = parts.next() else {
return false;
};
if parts.next().is_some() {
return false;
}
let Ok(index) = index_raw.parse::<usize>() else {
return false;
};
let file = bt_files.entry(index).or_default();
match field {
"path" => file.path = unescape_metadata_field(value),
"length" => {
if let Ok(parsed) = value.parse::<u64>() {
file.length = parsed;
}
}
"selected" => {
file.selected = value.parse::<bool>().unwrap_or(false);
}
"piece_offset" => {
if let Ok(parsed) = value.parse::<u64>() {
file.piece_offset = Some(parsed);
}
}
_ => {}
}
true
}
/// Restores request-group progress and piece state from persisted control metadata.
fn restore_control_metadata(group: &mut RequestGroup, control: ControlMetadata) {
if let Some(file) = control.files.first() {
group.set_total_length(file.length);
group.set_piece_length(file.piece_length);
}
group.set_completed_length(control.completed_length);
group.set_retry_count(control.retry_count);
if control.retry_count > 0 && group.retry_attempts().is_empty() {
let mut attempt = RetryAttempt::new(control.retry_count, control.completed_length);
attempt.length = Some(control.completed_length);
attempt.error.clone_from(&control.last_error);
attempt.recoverable = true;
group.push_retry_attempt(attempt);
}
if control.completed_length > 0 || control.resume_generation.is_some() {
group.set_resume_state(ResumeState {
persisted: control.resume_generation.is_some(),
resume_offset: control.completed_length,
validated_length: Some(control.completed_length),
segment_cursor: None,
});
}
for (piece, state) in control.piece_states {
group.set_piece_state(
crate::piece::PieceId(piece.0),
map_piece_state_from_storage(state),
);
}
}
/// Encodes retry attempts into a compact session-metadata string.
fn encode_retry_attempts_metadata(attempts: &[RetryAttempt]) -> String {
attempts
.iter()
.map(|attempt| {
let length = attempt
.length
.map_or_else(|| "-".to_owned(), |value| value.to_string());
let error = attempt.error.as_deref().unwrap_or("-");
format!(
"{}:{}:{}:{}:{}",
attempt.attempt,
attempt.offset,
length,
u8::from(attempt.recoverable),
escape_metadata_field(error)
)
})
.collect::<Vec<_>>()
.join(",")
}
/// Decodes retry attempts from the compact session-metadata string.
fn decode_retry_attempts_metadata(raw: &str) -> Vec<RetryAttempt> {
raw.split(',')
.filter(|entry| !entry.trim().is_empty())
.filter_map(|entry| {
let mut parts = entry.splitn(5, ':');
let attempt = parts.next()?.parse().ok()?;
let offset = parts.next()?.parse().ok()?;
let length = match parts.next()? {
"-" => None,
value => value.parse().ok(),
};
let recoverable = matches!(parts.next()?, "1" | "true");
let error = match parts.next()? {
"-" => None,
value => Some(unescape_metadata_field(value)),
};
Some(RetryAttempt {
attempt,
offset,
length,
error,
recoverable,
})
})
.collect()
}
/// Encodes resume-state metadata into a compact session-metadata string.
fn encode_resume_state_metadata(resume_state: &ResumeState) -> String {
let validated_length = resume_state
.validated_length
.map_or_else(|| "-".to_owned(), |value| value.to_string());
let segment_cursor = resume_state
.segment_cursor
.map_or_else(|| "-".to_owned(), |piece| piece.0.to_string());
format!(
"{}:{}:{}:{}",
u8::from(resume_state.persisted),
resume_state.resume_offset,
validated_length,
segment_cursor
)
}
/// Decodes resume-state metadata from the compact session-metadata string.
fn decode_resume_state_metadata(raw: &str) -> Option<ResumeState> {
let mut parts = raw.splitn(4, ':');
let persisted = matches!(parts.next()?, "1" | "true");
let resume_offset = parts.next()?.parse().ok()?;
let validated_length = match parts.next()? {
"-" => None,
value => value.parse().ok(),
};
let segment_cursor = match parts.next()? {
"-" => None,
value => value.parse().ok().map(crate::piece::PieceId),
};
Some(ResumeState {
persisted,
resume_offset,
validated_length,
segment_cursor,
})
}
/// Escapes reserved delimiters used by compact metadata encodings.
fn escape_metadata_field(raw: &str) -> String {
raw.replace('\\', "\\\\")
.replace(',', "\\c")
.replace(':', "\\d")
}
/// Reverses `escape_metadata_field`.
fn unescape_metadata_field(raw: &str) -> String {
let mut out = String::new();
let mut chars = raw.chars();
while let Some(ch) = chars.next() {
if ch == '\\' {
match chars.next() {
Some('c') => out.push(','),
Some('d') => out.push(':'),
Some('\\') | None => out.push('\\'),
Some(other) => {
out.push('\\');
out.push(other);
}
}
} else {
out.push(ch);
}
}
out
}
/// Parses the persisted RPC-style download status string.
fn parse_status(value: &str) -> Option<DownloadStatus> {
match value {
"active" => Some(DownloadStatus::Active),
"waiting" => Some(DownloadStatus::Waiting),
"paused" => Some(DownloadStatus::Paused),
"error" => Some(DownloadStatus::Error),
"complete" => Some(DownloadStatus::Complete),
"removed" => Some(DownloadStatus::Removed),
_ => None,
}
}