leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify initial check code to only return have_pieces

Browse source
This commit is contained in:
Igor Katson 2024-08-20 20:16:11 +01:00
parent 8135b31a5d
commit a55dfc6e0e
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
3 changed files with 56 additions and 96 deletions
Download patch file Download diff file Expand all files Collapse all files

59
crates/librqbit/src/torrent_state/initializing.rs
View file

@ -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<ManagedTorrentInfo>,
@ -42,21 +65,37 @@ impl TorrentStateInitializing {
pub async fn check(&self) -> anyhow::Result<TorrentStatePaused> {
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()?,