diff --git a/crates/librqbit/src/chunk_tracker.rs b/crates/librqbit/src/chunk_tracker.rs index 8db129b6e..071c9ce9f 100644 --- a/crates/librqbit/src/chunk_tracker.rs +++ b/crates/librqbit/src/chunk_tracker.rs @@ -1,4 +1,4 @@ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use anyhow::Context; use buffers::ByteBuf; @@ -9,7 +9,7 @@ use tracing::{debug, trace}; use crate::{ bitv::{BitV, BoxBitV}, file_info::FileInfo, - type_aliases::{BF, BS, FileInfos, FilePriorities}, + type_aliases::{FileInfos, FilePriorities, BF, BS}, }; pub struct ChunkTracker { @@ -42,6 +42,22 @@ pub struct ChunkTracker { // Quick to retrieve stats, that MUST be in sync with the BFs // above (have/selected). hns: HaveNeededSelected, + + // Track current streaming window per file (file_id -> (start_piece, end_piece)) + streaming_windows: HashMap>, +} + +/// Result of updating the streaming window +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StreamingWindowUpdate { + /// Number of pieces added to the download queue + pub pieces_added: usize, + /// Number of pieces removed from the download queue + pub pieces_removed: usize, + /// First piece index in the active window + pub window_start_piece: u32, + /// Last piece index (exclusive) in the active window + pub window_end_piece: u32, } #[derive(Default, Debug, PartialEq, Eq, Clone, Copy)] @@ -178,6 +194,7 @@ impl ChunkTracker { have: have_pieces, hns: HaveNeededSelected::default(), per_file_bytes: vec![0; file_infos.len()], + streaming_windows: HashMap::new(), }; ct.recalculate_per_file_bytes(file_infos); ct.hns = ct.calc_hns(); @@ -248,7 +265,7 @@ impl ChunkTracker { .filter_map(|id| self.lengths.validate_piece_index(id)) } - pub(crate) fn is_piece_have(&self, id: ValidPieceIndex) -> bool { + pub fn is_piece_have(&self, id: ValidPieceIndex) -> bool { self.have.as_slice()[id.get() as usize] } @@ -311,7 +328,9 @@ impl ChunkTracker { chunk_range.set(chunk_info.chunk_index as usize, true); trace!( "piece={}, chunk_info={:?}, bits={:?}", - piece.index, chunk_info, chunk_range, + piece.index, + chunk_info, + chunk_range, ); if chunk_range.all() { @@ -386,6 +405,87 @@ impl ChunkTracker { &self.per_file_bytes } + /// Update the download queue to only include pieces within a streaming window. + /// + /// This enables "stream-only" downloading where we only download pieces + /// from the current playback position forward, plus a small backward buffer. + /// + /// # Arguments + /// * `file_info` - The file being streamed + /// * `current_position` - Current byte position within the file (0-indexed) + /// * `backward_bytes` - How many bytes behind current position to keep in queue + /// * `forward_bytes` - How many bytes ahead of current position to keep in queue + /// + /// # Returns + /// The number of pieces added and removed from the queue + pub fn update_streaming_window( + &mut self, + file_id: usize, + file_info: &FileInfo, + current_position: u64, + backward_bytes: u64, + forward_bytes: u64, + ) -> StreamingWindowUpdate { + let file_start = file_info.offset_in_torrent; + let file_end = file_start + file_info.len; + + // Calculate absolute byte positions for the window + let abs_position = file_start + current_position.min(file_info.len); + let window_start = abs_position.saturating_sub(backward_bytes); + let window_end = (abs_position.saturating_add(forward_bytes)).min(file_end); + + // Convert to piece indices + let piece_len = self.lengths.default_piece_length() as u64; + let start_piece = (window_start / piece_len) as u32; + let end_piece = window_end.div_ceil(piece_len) as u32; + + let mut added = 0usize; + let mut removed = 0usize; + + // Iterate over all pieces in the file's range + for piece_idx in file_info.piece_range.clone() { + let idx = piece_idx as usize; + + // Check if this piece is within our streaming window + let in_window = piece_idx >= start_piece && piece_idx < end_piece; + let is_selected = self.selected.get(idx).map(|b| *b).unwrap_or(false); + let already_have = self.have.as_slice().get(idx).map(|b| *b).unwrap_or(false); + + // Piece should be queued if: in window AND selected AND not already downloaded + let should_be_queued = in_window && is_selected && !already_have; + let currently_queued = self.queue_pieces.get(idx).map(|b| *b).unwrap_or(false); + + if should_be_queued && !currently_queued { + self.queue_pieces.set(idx, true); + added += 1; + } else if !should_be_queued && currently_queued { + // Only remove from queue if it's part of this file + // (don't accidentally remove pieces from other files) + self.queue_pieces.set(idx, false); + removed += 1; + } + } + + // Recalculate stats + self.hns = self.calc_hns(); + + // Store the streaming window for this file + self.streaming_windows + .insert(file_id, start_piece..end_piece); + + StreamingWindowUpdate { + pieces_added: added, + pieces_removed: removed, + window_start_piece: start_piece, + window_end_piece: end_piece, + } + } + + /// Get the current streaming window for a file + pub fn get_streaming_window(&self, file_id: usize) -> Option> { + self.streaming_windows.get(&file_id).cloned() + } + // Returns remaining bytes pub fn update_file_have_on_piece_completed( &mut self, @@ -412,7 +512,7 @@ mod tests { bitv::BitV, chunk_tracker::HaveNeededSelected, file_info::FileInfo, type_aliases::BF, }; - use super::{ChunkTracker, compute_chunk_have_status}; + use super::{compute_chunk_have_status, ChunkTracker}; #[test] fn test_compute_chunk_status() { @@ -666,4 +766,222 @@ mod tests { assert!(ct.queue_pieces[1]); assert!(ct.queue_pieces[2]); } + + #[test] + fn test_streaming_window_basic() { + // Create a file with 10 pieces + let piece_len = CHUNK_SIZE * 2; + let total_len = piece_len as u64 * 10; + let l = Lengths::new(total_len, piece_len).unwrap(); + assert_eq!(l.total_pieces(), 10); + + let file_info = FileInfo { + relative_filename: "test.mp4".into(), + offset_in_torrent: 0, + piece_range: 0..10, + len: total_len, + attrs: Default::default(), + }; + + // Start with no pieces downloaded, all selected + let have_pieces = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + let selected_pieces = { + let mut bf = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + bf.get_mut(0..10).unwrap().fill(true); + bf + }; + + let mut ct = ChunkTracker::new( + have_pieces.into_dyn(), + selected_pieces, + l, + &vec![file_info.clone()], + ) + .unwrap(); + + // Initially all pieces should be queued + for i in 0..10 { + assert!(ct.queue_pieces[i], "piece {} should be queued initially", i); + } + + // Update streaming window: position at 50%, forward 30%, backward minimal + // Position: piece 5, forward window should include pieces 5-8 + let position = (piece_len as u64) * 5; // Start of piece 5 + let backward = piece_len as u64; // ~1 piece backward + let forward = (piece_len as u64) * 3; // ~3 pieces forward + + let result = ct.update_streaming_window(0, &file_info, position, backward, forward); + + // Window should be roughly pieces 4-8 (backward: 4, current+forward: 5,6,7,8) + assert!(result.pieces_removed > 0, "should have removed some pieces"); + + // Pieces before the window should NOT be queued + assert!(!ct.queue_pieces[0], "piece 0 should not be queued"); + assert!(!ct.queue_pieces[1], "piece 1 should not be queued"); + assert!(!ct.queue_pieces[2], "piece 2 should not be queued"); + assert!(!ct.queue_pieces[3], "piece 3 should not be queued"); + + // Pieces in the window SHOULD be queued + assert!( + ct.queue_pieces[4], + "piece 4 should be queued (backward buffer)" + ); + assert!(ct.queue_pieces[5], "piece 5 should be queued (current)"); + assert!(ct.queue_pieces[6], "piece 6 should be queued (forward)"); + assert!(ct.queue_pieces[7], "piece 7 should be queued (forward)"); + + // Pieces after the window should NOT be queued + assert!(!ct.queue_pieces[9], "piece 9 should not be queued"); + } + + #[test] + fn test_streaming_window_seek_forward() { + // Test seeking forward - old pieces should be removed from queue + let piece_len = CHUNK_SIZE * 2; + let total_len = piece_len as u64 * 10; + let l = Lengths::new(total_len, piece_len).unwrap(); + + let file_info = FileInfo { + relative_filename: "test.mp4".into(), + offset_in_torrent: 0, + piece_range: 0..10, + len: total_len, + attrs: Default::default(), + }; + + let have_pieces = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + let selected_pieces = { + let mut bf = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + bf.get_mut(0..10).unwrap().fill(true); + bf + }; + + let mut ct = ChunkTracker::new( + have_pieces.into_dyn(), + selected_pieces, + l, + &vec![file_info.clone()], + ) + .unwrap(); + + // First, set window at position 20% + let position1 = (piece_len as u64) * 2; + let forward = (piece_len as u64) * 3; + ct.update_streaming_window(0, &file_info, position1, piece_len as u64, forward); + + // Verify pieces 1-5 are in queue + assert!(ct.queue_pieces[2], "piece 2 should be queued"); + assert!(ct.queue_pieces[3], "piece 3 should be queued"); + assert!(ct.queue_pieces[4], "piece 4 should be queued"); + + // Now seek forward to 70% + let position2 = (piece_len as u64) * 7; + let result = + ct.update_streaming_window(0, &file_info, position2, piece_len as u64, forward); + + // Old pieces should be removed + assert!( + !ct.queue_pieces[2], + "piece 2 should be removed after seek forward" + ); + assert!( + !ct.queue_pieces[3], + "piece 3 should be removed after seek forward" + ); + assert!( + !ct.queue_pieces[4], + "piece 4 should be removed after seek forward" + ); + + // New window pieces should be queued + assert!( + ct.queue_pieces[6], + "piece 6 should be queued (backward buffer)" + ); + assert!(ct.queue_pieces[7], "piece 7 should be queued (current)"); + assert!(ct.queue_pieces[8], "piece 8 should be queued (forward)"); + assert!(ct.queue_pieces[9], "piece 9 should be queued (forward)"); + } + + #[test] + fn test_streaming_window_respects_have() { + // Pieces that are already downloaded should not be queued + let piece_len = CHUNK_SIZE * 2; + let total_len = piece_len as u64 * 10; + let l = Lengths::new(total_len, piece_len).unwrap(); + + let file_info = FileInfo { + relative_filename: "test.mp4".into(), + offset_in_torrent: 0, + piece_range: 0..10, + len: total_len, + attrs: Default::default(), + }; + + // Mark pieces 5 and 6 as already downloaded + let mut have_pieces = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + have_pieces.set(5, true); + have_pieces.set(6, true); + + let selected_pieces = { + let mut bf = + BF::from_boxed_slice(vec![0u8; l.piece_bitfield_bytes()].into_boxed_slice()); + bf.get_mut(0..10).unwrap().fill(true); + bf + }; + + let mut ct = ChunkTracker::new( + have_pieces.into_dyn(), + selected_pieces, + l, + &vec![file_info.clone()], + ) + .unwrap(); + + // Initially: pieces 5, 6 are NOT in queue (already have), others ARE in queue + assert!( + !ct.queue_pieces[5], + "piece 5 is already have, should not be queued initially" + ); + assert!( + !ct.queue_pieces[6], + "piece 6 is already have, should not be queued initially" + ); + assert!(ct.queue_pieces[7], "piece 7 should be queued initially"); + assert!(ct.queue_pieces[8], "piece 8 should be queued initially"); + + // Update window centered around pieces 5-8 + let position = (piece_len as u64) * 5; + let forward = (piece_len as u64) * 4; // Forward to piece 9 + ct.update_streaming_window(0, &file_info, position, piece_len as u64, forward); + + // Pieces 5 and 6 should STILL not be queued (already have) + assert!( + !ct.queue_pieces[5], + "piece 5 is already have, should not be queued" + ); + assert!( + !ct.queue_pieces[6], + "piece 6 is already have, should not be queued" + ); + + // Pieces 7 and 8 are in window and not downloaded - SHOULD be queued + assert!(ct.queue_pieces[7], "piece 7 should be queued"); + assert!(ct.queue_pieces[8], "piece 8 should be queued"); + + // Pieces outside window should NOT be queued + assert!( + !ct.queue_pieces[0], + "piece 0 should not be queued (outside window)" + ); + assert!( + !ct.queue_pieces[1], + "piece 1 should not be queued (outside window)" + ); + } } diff --git a/crates/librqbit/src/lib.rs b/crates/librqbit/src/lib.rs index e1c37beaf..392972537 100644 --- a/crates/librqbit/src/lib.rs +++ b/crates/librqbit/src/lib.rs @@ -82,6 +82,7 @@ mod vectored_traits; #[cfg(feature = "watch")] pub mod watch; +pub use chunk_tracker::StreamingWindowUpdate; pub use error::{Error, Result}; pub use api::Api; diff --git a/crates/librqbit/src/piece_tracker.rs b/crates/librqbit/src/piece_tracker.rs index 7d872877d..6fc0f7ce8 100644 --- a/crates/librqbit/src/piece_tracker.rs +++ b/crates/librqbit/src/piece_tracker.rs @@ -97,6 +97,11 @@ impl PieceTracker { &self.chunks } + /// Mutable access to the underlying ChunkTracker. + pub fn chunks_mut(&mut self) -> &mut ChunkTracker { + &mut self.chunks + } + /// Consume the PieceTracker, requeuing any in-flight pieces. /// /// This is used when pausing a torrent - any pieces that were being downloaded diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs index d50ff8ca4..625591557 100644 --- a/crates/librqbit/src/torrent_state/live/mod.rs +++ b/crates/librqbit/src/torrent_state/live/mod.rs @@ -152,6 +152,13 @@ impl TorrentStateLocked { self.pieces.as_mut().ok_or(Error::ChunkTrackerEmpty) } + pub(crate) fn get_chunks_mut(&mut self) -> crate::Result<&mut ChunkTracker> { + self.pieces + .as_mut() + .map(|p| p.chunks_mut()) + .ok_or(Error::ChunkTrackerEmpty) + } + fn try_flush_bitv(&mut self, shared: &ManagedTorrentShared, flush_async: bool) { if self.unflushed_bitv_bytes == 0 { return; diff --git a/crates/librqbit/src/torrent_state/mod.rs b/crates/librqbit/src/torrent_state/mod.rs index fe9e7e211..d7a09706d 100644 --- a/crates/librqbit/src/torrent_state/mod.rs +++ b/crates/librqbit/src/torrent_state/mod.rs @@ -248,7 +248,7 @@ impl ManagedTorrent { f(&mut self.locked.write().state) } - pub(crate) fn with_chunk_tracker( + pub fn with_chunk_tracker( &self, f: impl FnOnce(&ChunkTracker) -> R, ) -> anyhow::Result { @@ -263,6 +263,90 @@ impl ManagedTorrent { } } + pub(crate) fn with_chunk_tracker_mut( + &self, + f: impl FnOnce(&mut ChunkTracker) -> R, + ) -> anyhow::Result { + let g = self.locked.read(); + match &g.state { + ManagedTorrentState::Paused(_) => { + // Paused torrents can't update streaming window - they need to be live + bail!("cannot modify chunk tracker while paused") + } + ManagedTorrentState::Live(l) => Ok(f(l + .lock_write("chunk_tracker_mut") + .get_chunks_mut() + .context("error getting chunks")?)), + _ => bail!("no chunk tracker, torrent neither paused nor live"), + } + } + + /// Update the streaming download window + /// + /// This restricts downloading to only pieces within the specified window + /// around the current playback position. + pub fn update_streaming_window( + &self, + file_id: usize, + current_position: u64, + backward_bytes: u64, + forward_bytes: u64, + ) -> anyhow::Result { + let metadata = self.metadata.load_full().context("metadata not available")?; + let file_info = metadata + .file_infos + .get(file_id) + .context("invalid file_id")?; + + self.with_chunk_tracker_mut(|ct| { + ct.update_streaming_window(file_id, file_info, current_position, backward_bytes, forward_bytes) + }) + } + + /// Get the total number of pieces in the torrent + pub fn total_pieces(&self) -> usize { + self.with_metadata(|m| m.lengths().total_pieces() as usize) + .unwrap_or(0) + } + + /// Get a bitfield of downloaded pieces as Vec + pub fn get_piece_bitfield(&self) -> Vec { + self.with_chunk_tracker(|ct| { + let have_pieces = ct.get_have_pieces(); + have_pieces.as_slice() + .iter() + .map(|b| *b) + .collect() + }) + .unwrap_or_default() + } + + /// Check if a specific piece is downloaded + pub fn is_piece_downloaded(&self, piece_index: usize) -> bool { + self.with_chunk_tracker(|ct| { + let lengths = ct.get_lengths(); + lengths.validate_piece_index(piece_index as u32) + .map(|idx| ct.is_piece_have(idx)) + .unwrap_or(false) + }) + .unwrap_or(false) + } + + /// Get the number of downloaded pieces + pub fn downloaded_pieces_count(&self) -> usize { + self.with_chunk_tracker(|ct| { + ct.get_have_pieces().as_slice().count_ones() + }) + .unwrap_or(0) + } + + /// Get the current streaming window for a file + pub fn get_streaming_window(&self, file_id: usize) -> Option> { + self.with_chunk_tracker(|ct| ct.get_streaming_window(file_id)) + .ok() + .flatten() + } + /// Get the live state if the torrent is live. pub fn live(&self) -> Option> { let g = self.locked.read();