From a55dfc6e0ede0e066740f5b1131df2302b22163c Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Tue, 20 Aug 2024 20:16:11 +0100 Subject: [PATCH] Simplify initial check code to only return have_pieces --- crates/librqbit/src/chunk_tracker.rs | 1 - crates/librqbit/src/file_ops.rs | 92 +++---------------- .../src/torrent_state/initializing.rs | 59 +++++++++--- 3 files changed, 56 insertions(+), 96 deletions(-) diff --git a/crates/librqbit/src/chunk_tracker.rs b/crates/librqbit/src/chunk_tracker.rs index de6d026..ffab7c2 100644 --- a/crates/librqbit/src/chunk_tracker.rs +++ b/crates/librqbit/src/chunk_tracker.rs @@ -1,7 +1,6 @@ use std::collections::HashSet; use anyhow::Context; -use bitvec::{order::Lsb0, slice::BitSlice}; use librqbit_core::lengths::{ChunkInfo, Lengths, ValidPieceIndex}; use peer_binary_protocol::Piece; use tracing::{debug, trace}; diff --git a/crates/librqbit/src/file_ops.rs b/crates/librqbit/src/file_ops.rs index e8a200d..237b08f 100644 --- a/crates/librqbit/src/file_ops.rs +++ b/crates/librqbit/src/file_ops.rs @@ -19,29 +19,6 @@ use crate::{ type_aliases::{FileInfos, PeerHandle, BF}, }; -pub(crate) struct InitialCheckResults { - // A piece as flags based on these dimensions: - // - if the asked for it or not (only_files) - // - if we have it downloaded and verified - // - if we need to queue it for downloading - // this one depends if we queued it already or not. - - // The pieces we have downloaded. - pub have_pieces: BF, - // The pieces that the user selected to download. - pub selected_pieces: BF, - - // How many bytes we have. This can be MORE than "total_selected_bytes", - // if we downloaded some pieces, and later the "only_files" was changed. - pub have_bytes: u64, - // How many bytes we need to download. - pub needed_bytes: u64, - - // How many bytes are in selected pieces. - // If all selected, this must be equal to total torrent length. - pub selected_bytes: u64, -} - pub fn update_hash_from_file( file_id: usize, mut pos: u64, @@ -88,26 +65,16 @@ impl<'a> FileOps<'a> { } } - pub fn initial_check( - &self, - only_files: Option<&[usize]>, - progress: &AtomicU64, - ) -> anyhow::Result { - let mut needed_pieces = + // Returns the bitvector with pieces we have. + pub fn initial_check(&self, progress: &AtomicU64) -> anyhow::Result { + let mut have_pieces = BF::from_boxed_slice(vec![0u8; self.lengths.piece_bitfield_bytes()].into()); - let mut have_pieces = needed_pieces.clone(); - let mut selected_pieces = needed_pieces.clone(); - - let mut have_bytes = 0u64; - let mut needed_bytes = 0u64; - let mut total_selected_bytes = 0u64; let mut piece_files = Vec::::new(); #[derive(Debug)] struct CurrentFile<'a> { index: usize, fi: &'a FileInfo, - full_file_required: bool, processed_bytes: u64, is_broken: bool, } @@ -119,20 +86,16 @@ impl<'a> FileOps<'a> { self.processed_bytes += bytes } } - let mut file_iterator = self.file_infos.iter().enumerate().map(|(idx, fi)| { - let full_file_required = if let Some(only_files) = only_files { - only_files.contains(&idx) - } else { - true - }; - CurrentFile { + let mut file_iterator = self + .file_infos + .iter() + .enumerate() + .map(|(idx, fi)| CurrentFile { index: idx, fi, - full_file_required, processed_bytes: 0, is_broken: false, - } - }); + }); let mut current_file = file_iterator .next() @@ -145,7 +108,6 @@ impl<'a> FileOps<'a> { let mut computed_hash = Sha1::new(); let mut piece_remaining = piece_info.len as usize; let mut some_files_broken = false; - let mut piece_selected = current_file.full_file_required; progress.fetch_add(piece_info.len as u64, Ordering::Relaxed); while piece_remaining > 0 { @@ -158,8 +120,6 @@ impl<'a> FileOps<'a> { .next() .ok_or_else(|| anyhow::anyhow!("broken torrent metadata"))?; - piece_selected |= current_file.full_file_required; - to_read_in_file = std::cmp::min(current_file.remaining(), piece_remaining as u64) .try_into()?; @@ -193,18 +153,11 @@ impl<'a> FileOps<'a> { } } - if piece_selected { - total_selected_bytes += piece_info.len as u64; - selected_pieces.set(piece_info.piece_index.get() as usize, true); - } - - if piece_selected && some_files_broken { + if some_files_broken { trace!( "piece {} had errors, marking as needed", piece_info.piece_index ); - - needed_bytes += piece_info.len as u64; continue; } @@ -213,34 +166,11 @@ impl<'a> FileOps<'a> { .compare_hash(piece_info.piece_index.get(), computed_hash.finish()) .context("bug: either torrent info broken or we have a bug - piece index invalid")? { - trace!( - "piece {} is fine, not marking as needed", - piece_info.piece_index - ); - have_bytes += piece_info.len as u64; have_pieces.set(piece_info.piece_index.get() as usize, true); - } else if piece_selected { - trace!( - "piece {} hash does not match, marking as needed", - piece_info.piece_index - ); - needed_bytes += piece_info.len as u64; - needed_pieces.set(piece_info.piece_index.get() as usize, true); - } else { - trace!( - "piece {} hash does not match, but it is not required by any of the requested files, ignoring", - piece_info.piece_index - ); } } - Ok(InitialCheckResults { - have_pieces, - selected_pieces, - have_bytes, - needed_bytes, - selected_bytes: total_selected_bytes, - }) + Ok(have_pieces) } pub fn check_piece( diff --git a/crates/librqbit/src/torrent_state/initializing.rs b/crates/librqbit/src/torrent_state/initializing.rs index 9e97c09..62b3622 100644 --- a/crates/librqbit/src/torrent_state/initializing.rs +++ b/crates/librqbit/src/torrent_state/initializing.rs @@ -5,11 +5,16 @@ use std::{ use anyhow::Context; +use librqbit_core::lengths::Lengths; use size_format::SizeFormatterBinary as SF; use tracing::{debug, info, warn}; use crate::{ - bitv::BitV, chunk_tracker::ChunkTracker, file_ops::FileOps, type_aliases::FileStorage, + bitv::BitV, + chunk_tracker::ChunkTracker, + file_ops::FileOps, + type_aliases::{FileStorage, BF}, + FileInfos, }; use super::{paused::TorrentStatePaused, ManagedTorrentInfo}; @@ -21,6 +26,24 @@ pub struct TorrentStateInitializing { pub(crate) checked_bytes: AtomicU64, } +fn compute_selected_pieces( + lengths: &Lengths, + only_files: Option<&[usize]>, + file_infos: &FileInfos, +) -> BF { + let mut bf = BF::from_boxed_slice(vec![0u8; lengths.piece_bitfield_bytes()].into_boxed_slice()); + for (_, fi) in file_infos + .iter() + .enumerate() + .filter(|(id, _)| only_files.map(|of| of.contains(id)).unwrap_or(false)) + { + if let Some(r) = bf.get_mut(fi.piece_range_usize()) { + r.fill(true); + } + } + bf +} + impl TorrentStateInitializing { pub fn new( meta: Arc, @@ -42,21 +65,37 @@ impl TorrentStateInitializing { pub async fn check(&self) -> anyhow::Result { info!("Doing initial checksum validation, this might take a while..."); - let initial_check_results = self.meta.spawner.spawn_block_in_place(|| { + let have_pieces = self.meta.spawner.spawn_block_in_place(|| { FileOps::new( &self.meta.info, &self.files, &self.meta.file_infos, &self.meta.lengths, ) - .initial_check(self.only_files.as_deref(), &self.checked_bytes) + .initial_check(&self.checked_bytes) })?; + let selected_pieces = compute_selected_pieces( + &self.meta.lengths, + self.only_files.as_deref(), + &self.meta.file_infos, + ); + + let chunk_tracker = ChunkTracker::new( + have_pieces.into_dyn(), + selected_pieces, + self.meta.lengths, + &self.meta.file_infos, + ) + .context("error creating chunk tracker")?; + + let hns = chunk_tracker.get_hns(); + info!( "Initial check results: have {}, needed {}, total selected {}", - SF::new(initial_check_results.have_bytes), - SF::new(initial_check_results.needed_bytes), - SF::new(initial_check_results.selected_bytes) + SF::new(hns.have_bytes), + SF::new(hns.needed_bytes), + SF::new(hns.selected_bytes) ); // Ensure file lenghts are correct, and reopen read-only. @@ -87,14 +126,6 @@ impl TorrentStateInitializing { Ok::<_, anyhow::Error>(()) })?; - let chunk_tracker = ChunkTracker::new( - initial_check_results.have_pieces.into_dyn(), - initial_check_results.selected_pieces, - self.meta.lengths, - &self.meta.file_infos, - ) - .context("error creating chunk tracker")?; - let paused = TorrentStatePaused { info: self.meta.clone(), files: self.files.take()?,