From 276a1e175e2dc61f343df296a5ae1436b2fd51e2 Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Sat, 30 Mar 2024 17:03:36 +0000 Subject: [PATCH] Add a method to update_only_files to chunk_tracker --- crates/librqbit/src/chunk_tracker.rs | 219 ++++++++++++++++++++++++++- 1 file changed, 216 insertions(+), 3 deletions(-) diff --git a/crates/librqbit/src/chunk_tracker.rs b/crates/librqbit/src/chunk_tracker.rs index b0bd39f..19b1642 100644 --- a/crates/librqbit/src/chunk_tracker.rs +++ b/crates/librqbit/src/chunk_tracker.rs @@ -1,5 +1,11 @@ +use std::collections::HashSet; + use anyhow::Context; -use librqbit_core::lengths::{ChunkInfo, Lengths, ValidPieceIndex}; +use buffers::ByteBufOwned; +use librqbit_core::{ + lengths::{ChunkInfo, Lengths, ValidPieceIndex}, + torrent_metainfo::{TorrentMetaV1Info, TorrentMetaV1Owned}, +}; use peer_binary_protocol::Piece; use tracing::{debug, trace}; @@ -29,6 +35,12 @@ pub struct ChunkTracker { total_selected_bytes: u64, } +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub struct HaveNeeded { + pub have_bytes: u64, + pub needed_bytes: u64, +} + // Comput the have-status of chunks. // // Save as "have_pieces", but there's one bit per chunk (not per piece). @@ -227,15 +239,96 @@ impl ChunkTracker { } Some(ChunkMarkingResult::NotCompleted) } + + // NOTE: this doesn't validate new_only_files. + // E.g. if there are indices there that don't make + // sense, they will be ignored. + pub fn update_only_files( + &mut self, + file_lengths_iterator: impl IntoIterator, + // TODO: maybe make this a BF + new_only_files: &HashSet, + ) -> anyhow::Result { + let mut piece_it = self.lengths.iter_piece_infos(); + let mut current_piece = piece_it + .next() + .context("bug: iter_piece_infos() returned empty iterator")?; + let mut current_piece_needed = false; + let mut current_piece_remaining = current_piece.len; + let mut have_bytes = 0u64; + let mut needed_bytes = 0u64; + + for (idx, len) in file_lengths_iterator.into_iter().enumerate() { + let file_required = new_only_files.contains(&idx); + + let mut remaining_file_len = len; + + while remaining_file_len > 0 { + current_piece_needed |= len > 0 && file_required; + let shift = std::cmp::min(current_piece_remaining as u64, remaining_file_len); + assert!(shift > 0); + remaining_file_len -= shift; + current_piece_remaining -= shift as u32; + + dbg!( + idx, + shift, + remaining_file_len, + current_piece_remaining, + current_piece_needed, + file_required, + current_piece + ); + + if current_piece_remaining == 0 { + let current_piece_have = self.have[current_piece.piece_index.get() as usize]; + if current_piece_have { + have_bytes += current_piece.len as u64; + } + if current_piece_needed { + needed_bytes += current_piece.len as u64; + } + match (current_piece_needed, current_piece_have) { + (true, true) => {} + (true, false) => { + dbg!(self.mark_piece_broken_if_not_have(current_piece.piece_index)) + } + (false, true) => {} + (false, false) => { + // don't need the piece, and don't have it - cancel downloading it + dbg!(self + .needed_pieces + .set(current_piece.piece_index.get() as usize, false)); + } + } + + if current_piece.piece_index != self.lengths.last_piece_id() { + current_piece = piece_it.next().context( + "bug: iter_piece_infos() pieces ended earlier than expected", + )?; + current_piece_needed = false; + current_piece_remaining = current_piece.len; + } + } + } + } + + Ok(HaveNeeded { + have_bytes, + needed_bytes, + }) + } } #[cfg(test)] mod tests { + use std::collections::HashSet; + use librqbit_core::{constants::CHUNK_SIZE, lengths::Lengths}; - use crate::type_aliases::BF; + use crate::{chunk_tracker::HaveNeeded, type_aliases::BF}; - use super::compute_chunk_have_status; + use super::{compute_chunk_have_status, ChunkTracker}; #[test] fn test_compute_chunk_status() { @@ -334,4 +427,124 @@ mod tests { } } } + + #[test] + fn test_update_only_files() { + let piece_len = CHUNK_SIZE * 2 + 1; + let total_len = piece_len as u64 * 2 + 1; + let l = Lengths::new(total_len, piece_len).unwrap(); + assert_eq!(l.total_pieces(), 3); + assert_eq!(l.total_chunks(), 7); + + let all_files = [ + piece_len as u64, // piece 0 and boundary + 1, // piece 1 + 0, // piece 1 (or none) + piece_len as u64, // piece 1 and 2 + ]; + + let bf_len = l.piece_bitfield_bytes(); + let initial_have = BF::from_boxed_slice(vec![0u8; bf_len].into_boxed_slice()); + let initial_needed = BF::from_boxed_slice(vec![u8::MAX; bf_len].into_boxed_slice()); + + // Initially, we need all files and all pieces. + let mut ct = ChunkTracker::new( + initial_needed.clone(), + initial_have.clone(), + l, + l.total_length(), + ) + .unwrap(); + + // Select all file, no changes. + assert_eq!( + ct.update_only_files(all_files.into_iter(), &HashSet::from_iter([0, 1, 2, 3])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: total_len + } + ); + assert_eq!(ct.have, initial_have); + assert_eq!(ct.needed_pieces, initial_needed); + + // Select only the first file. + println!("Select only the first file."); + assert_eq!( + ct.update_only_files(all_files, &HashSet::from_iter([0])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: all_files[0], + } + ); + assert_eq!(ct.needed_pieces[0], true); + assert_eq!(ct.needed_pieces[1], false); + assert_eq!(ct.needed_pieces[2], false); + + // Select only the second file. + assert_eq!( + ct.update_only_files(all_files, &HashSet::from_iter([1])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: piece_len as u64, + } + ); + assert_eq!(ct.needed_pieces[0], false); + assert_eq!(ct.needed_pieces[1], true); + assert_eq!(ct.needed_pieces[2], false); + + // Select only the third file (zero sized one!). + assert_eq!( + ct.update_only_files(all_files, &HashSet::from_iter([2])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: 0, + } + ); + assert_eq!(ct.needed_pieces[0], false); + assert_eq!(ct.needed_pieces[1], false); + assert_eq!(ct.needed_pieces[2], false); + + // Select only the fourth file. + assert_eq!( + ct.update_only_files(all_files, &HashSet::from_iter([3])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: (piece_len + 1) as u64, + } + ); + assert_eq!(ct.needed_pieces[0], false); + assert_eq!(ct.needed_pieces[1], true); + assert_eq!(ct.needed_pieces[2], true); + + // Select first and last file + assert_eq!( + ct.update_only_files(all_files.clone(), &HashSet::from_iter([0, 3])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: all_files[0] + all_files[3] + 1, + } + ); + assert_eq!(ct.needed_pieces[0], true); + assert_eq!(ct.needed_pieces[1], true); + assert_eq!(ct.needed_pieces[2], true); + + // Select all files + assert_eq!( + ct.update_only_files(all_files.clone(), &HashSet::from_iter([0, 1, 2, 3])) + .unwrap(), + HaveNeeded { + have_bytes: 0, + needed_bytes: total_len, + } + ); + assert_eq!(ct.needed_pieces[0], true); + assert_eq!(ct.needed_pieces[1], true); + assert_eq!(ct.needed_pieces[2], true); + } }