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

Storages

Browse source
This commit is contained in:
Igor Katson 2024-04-29 21:44:21 +01:00
parent 609f9d92ae
commit 1b49257019
13 changed files with 499 additions and 230 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -10,8 +10,10 @@ use size_format::SizeFormatterBinary as SF;
use tracing::{debug, info, warn};
use crate::{
chunk_tracker::ChunkTracker, file_ops::FileOps, opened_file::OpenedFile,
type_aliases::OpenedFiles,
chunk_tracker::ChunkTracker,
file_ops::FileOps,
opened_file::OpenedFile,
storage::{FilesystemStorage, InMemoryGarbageCollectingStorage, TorrentStorage},
};
use super::{paused::TorrentStatePaused, ManagedTorrentInfo};
@ -37,7 +39,23 @@ impl TorrentStateInitializing {
}
pub async fn check(&self) -> anyhow::Result<TorrentStatePaused> {
let mut files = OpenedFiles::new();
// Return in-memory store
let store =
InMemoryGarbageCollectingStorage::new(self.meta.lengths, self.meta.file_infos.clone())?;
let ct = ChunkTracker::new_empty(self.meta.lengths, &self.meta.file_infos)?;
Ok(TorrentStatePaused {
info: self.meta.clone(),
files: Box::new(store),
chunk_tracker: ct,
streams: Arc::new(Default::default()),
})
// self.check_disk().await
}
pub async fn check_disk(&self) -> anyhow::Result<TorrentStatePaused> {
let mut files = Vec::<OpenedFile>::new();
for file_details in self.meta.info.iter_file_details(&self.meta.lengths)? {
let mut full_path = self.meta.out_dir.clone();
let relative_path = file_details
@ -64,26 +82,21 @@ impl TorrentStateInitializing {
.with_context(|| format!("error creating {:?}", &full_path))?;
OpenOptions::new().read(true).write(true).open(&full_path)?
};
files.push(OpenedFile::new(
file,
full_path,
0,
file_details.len,
file_details.offset,
file_details.pieces,
));
files.push(OpenedFile::new(file));
}
let files: Box<dyn TorrentStorage> = Box::new(FilesystemStorage::new(files));
debug!("computed lengths: {:?}", &self.meta.lengths);
info!("Doing initial checksum validation, this might take a while...");
let initial_check_results = self.meta.spawner.spawn_block_in_place(|| {
FileOps::new(&self.meta.info, &files, &self.meta.lengths).initial_check(
self.only_files.as_deref(),
FileOps::new(
&self.meta.info,
&files,
&self.meta.file_infos,
&self.meta.lengths,
&self.checked_bytes,
)
.initial_check(self.only_files.as_deref(), &self.checked_bytes)
})?;
info!(
@ -95,7 +108,7 @@ impl TorrentStateInitializing {
// Ensure file lenghts are correct, and reopen read-only.
self.meta.spawner.spawn_block_in_place(|| {
for (idx, file) in files.iter().enumerate() {
for (idx, fi) in self.meta.file_infos.iter().enumerate() {
if self
.only_files
.as_ref()
@ -103,16 +116,16 @@ impl TorrentStateInitializing {
.unwrap_or(true)
{
let now = Instant::now();
if let Err(err) = file.file.lock().set_len(file.len) {
if let Err(err) = files.ensure_file_length(idx, fi.len) {
warn!(
"Error setting length for file {:?} to {}: {:#?}",
file.filename, file.len, err
fi.filename, fi.len, err
);
} else {
debug!(
"Set length for file {:?} to {} in {:?}",
file.filename,
SF::new(file.len),
fi.filename,
SF::new(fi.len),
now.elapsed()
);
}
@ -125,6 +138,7 @@ impl TorrentStateInitializing {
initial_check_results.have_pieces,
initial_check_results.selected_pieces,
self.meta.lengths,
&self.meta.file_infos,
)
.context("error creating chunk tracker")?;

83
crates/librqbit/src/torrent_state/live/mod.rs
View file

@ -86,8 +86,9 @@ use crate::{
PeerConnection, PeerConnectionHandler, PeerConnectionOptions, WriterRequest,
},
session::CheckedIncomingConnection,
storage::TorrentStorage,
torrent_state::{peer::Peer, utils::atomic_inc},
type_aliases::{FilePriorities, OpenedFiles, PeerHandle, BF},
type_aliases::{FilePriorities, FileStorage, PeerHandle, BF},
};
use self::{
@ -141,7 +142,7 @@ impl TorrentStateLocked {
.context("chunk tracker empty, torrent was paused")
}
fn get_chunks_mut(&mut self) -> anyhow::Result<&mut ChunkTracker> {
pub(crate) fn get_chunks_mut(&mut self) -> anyhow::Result<&mut ChunkTracker> {
self.chunks
.as_mut()
.context("chunk tracker empty, torrent was paused")
@ -159,7 +160,7 @@ pub struct TorrentStateLive {
meta: Arc<ManagedTorrentInfo>,
locked: RwLock<TorrentStateLocked>,
pub(crate) files: OpenedFiles,
pub(crate) files: FileStorage,
stats: AtomicStats,
lengths: Lengths,
@ -195,9 +196,15 @@ impl TorrentStateLive {
// TODO: make it configurable
let file_priorities = {
let mut pri = (0..paused.files.len()).collect::<Vec<usize>>();
let mut pri = (0..paused.info.file_infos.len()).collect::<Vec<usize>>();
// sort by filename, cause many torrents have random sort order.
pri.sort_unstable_by_key(|id| paused.files.get(*id).map(|op| op.filename.as_path()));
pri.sort_unstable_by_key(|id| {
paused
.info
.file_infos
.get(*id)
.map(|fi| fi.filename.as_path())
});
pri
};
@ -482,7 +489,12 @@ impl TorrentStateLive {
self.meta.peer_id
}
pub(crate) fn file_ops(&self) -> FileOps<'_> {
FileOps::new(&self.meta.info, &self.files, &self.lengths)
FileOps::new(
&self.meta.info,
&self.files,
&self.meta().file_infos,
&self.lengths,
)
}
pub(crate) fn lock_read(
@ -632,12 +644,6 @@ impl TorrentStateLive {
// It should be impossible to make a fatal error after pausing.
g.fatal_errors_tx.take();
let files = self
.files
.iter()
.map(|f| f.take_clone())
.collect::<anyhow::Result<Vec<_>>>()?;
let mut chunk_tracker = g
.chunks
.take()
@ -649,7 +655,7 @@ impl TorrentStateLive {
// g.chunks;
Ok(TorrentStatePaused {
info: self.meta.clone(),
files,
files: self.files.take()?,
chunk_tracker,
streams: self.streams.clone(),
})
@ -671,7 +677,7 @@ impl TorrentStateLive {
pub(crate) fn update_only_files(&self, only_files: &HashSet<usize>) -> anyhow::Result<()> {
let mut g = self.lock_write("update_only_files");
let ct = g.get_chunks_mut()?;
let hns = ct.update_only_files(self.files.iter().map(|f| f.len), only_files)?;
let hns = ct.update_only_files(self.meta().file_infos.iter().map(|f| f.len), only_files)?;
if !hns.finished() {
self.reconnect_all_not_needed_peers();
}
@ -682,41 +688,49 @@ impl TorrentStateLive {
self.get_hns().map(|h| h.finished()).unwrap_or_default()
}
pub(crate) fn has_active_streams_unfinished_files(&self) -> bool {
fn has_active_streams_unfinished_files(&self, state: &TorrentStateLocked) -> bool {
let chunks = match state.get_chunks() {
Ok(c) => c,
Err(_) => return false,
};
self.streams
.streamed_file_ids()
.any(|file_id| !self.files[file_id].approx_is_finished())
.any(|file_id| !chunks.is_file_finished(&self.meta.file_infos[file_id]))
}
pub(crate) fn is_finished_and_dont_need_peers(&self) -> bool {
self.is_finished() && !self.has_active_streams_unfinished_files()
fn is_finished_and_dont_need_peers(&self) -> bool {
self.is_finished()
&& !self.has_active_streams_unfinished_files(
&self.lock_read("is_finished_and_dont_need_peers"),
)
}
fn on_piece_completed(&self, id: ValidPieceIndex) -> anyhow::Result<()> {
let mut g = self.lock_write("on_piece_completed");
let chunks = g.get_chunks_mut()?;
// if we have all the pieces of the file, reopen it read only
for (idx, opened_file) in self
.files
for (idx, file_info) in self
.meta()
.file_infos
.iter()
.enumerate()
.skip_while(|fd| !fd.1.piece_range.contains(&id.get()))
.take_while(|fd| fd.1.piece_range.contains(&id.get()))
.skip_while(|(_, fi)| !fi.piece_range.contains(&id.get()))
.take_while(|(_, fi)| fi.piece_range.contains(&id.get()))
{
let bytes = opened_file.update_have_on_piece_completed(id.get(), &self.lengths);
if bytes == 0 {
warn!(file_id=idx, piece_id=id.get(), "bug: update_have_on_piece_completed() returned 0, although this piece is present in the file");
}
let _remaining = chunks.update_file_have_on_piece_completed(id, idx, file_info);
}
self.streams
.wake_streams_on_piece_completed(id, &self.meta.lengths);
if self.is_finished() {
if self.lock_read("chunks").get_chunks()?.get_selected_pieces()[id.get_usize()] {
if chunks.is_finished() {
if chunks.get_selected_pieces()[id.get_usize()] {
info!("torrent finished downloading");
}
self.finished_notify.notify_waiters();
if !self.has_active_streams_unfinished_files() {
if !self.has_active_streams_unfinished_files(&g) {
// There is not poing being connected to peers that have all the torrent, when
// we don't need anything from them, and they don't need anything from us.
self.disconnect_all_peers_that_have_full_torrent();
@ -749,13 +763,6 @@ impl TorrentStateLive {
}
}
}
pub(crate) fn get_file_progress(&self) -> Vec<u64> {
self.files
.iter()
.map(|fd| fd.have.load(Ordering::Relaxed))
.collect()
}
}
struct PeerHandlerLocked {
@ -989,8 +996,8 @@ impl PeerHandler {
!chunk_tracker.is_piece_have(*pid)
&& !g.inflight_pieces.contains_key(pid)
});
let natural_order_pieces =
chunk_tracker.iter_queued_pieces(&g.file_priorities, &self.state.files);
let natural_order_pieces = chunk_tracker
.iter_queued_pieces(&g.file_priorities, &self.state.meta().file_infos);
for n in priority_streamed_pieces.chain(natural_order_pieces) {
if bf.get(n.get() as usize).map(|v| *v) == Some(true) {
n_opt = Some(n);

29
crates/librqbit/src/torrent_state/mod.rs
View file

@ -34,8 +34,10 @@ use tracing::error_span;
use tracing::warn;
use crate::chunk_tracker::ChunkTracker;
use crate::file_info::FileInfo;
use crate::spawn_utils::BlockingSpawner;
use crate::torrent_state::stats::LiveStats;
use crate::type_aliases::FileInfos;
use crate::type_aliases::PeerStream;
use initializing::TorrentStateInitializing;
@ -98,6 +100,7 @@ pub struct ManagedTorrentInfo {
pub trackers: HashSet<String>,
pub peer_id: Id20,
pub lengths: Lengths,
pub file_infos: FileInfos,
pub span: tracing::Span,
pub(crate) options: ManagedTorrentOptions,
}
@ -370,11 +373,7 @@ impl ManagedTorrent {
resp.total_bytes = hns.total();
resp.progress_bytes = hns.progress();
resp.finished = hns.finished();
resp.file_progress = p
.files
.iter()
.map(|f| f.have.load(Ordering::Relaxed))
.collect();
resp.file_progress = p.chunk_tracker.per_file_have_bytes().to_owned();
}
ManagedTorrentState::Live(l) => {
resp.state = S::Live;
@ -384,7 +383,12 @@ impl ManagedTorrent {
resp.progress_bytes = hns.progress();
resp.finished = hns.finished();
resp.uploaded_bytes = l.get_uploaded_bytes();
resp.file_progress = l.get_file_progress();
resp.file_progress = l
.lock_read("file_progress")
.get_chunks()
.ok()
.map(|c| c.per_file_have_bytes().to_owned())
.unwrap_or_default();
resp.live = Some(live_stats);
}
ManagedTorrentState::Error(e) => {
@ -534,8 +538,21 @@ impl ManagedTorrentBuilder {
pub(crate) fn build(self, span: tracing::Span) -> anyhow::Result<ManagedTorrentHandle> {
let lengths = Lengths::from_torrent(&self.info)?;
let file_infos = self
.info
.iter_file_details(&lengths)?
.map(|fd| {
Ok::<_, anyhow::Error>(FileInfo {
filename: self.output_folder.join(fd.filename.to_pathbuf()?),
offset_in_torrent: fd.offset,
piece_range: fd.pieces,
len: fd.len,
})
})
.collect::<anyhow::Result<Vec<FileInfo>>>()?;
let info = Arc::new(ManagedTorrentInfo {
span,
file_infos,
info: self.info,
info_hash: self.info_hash,
out_dir: self.output_folder,

4
crates/librqbit/src/torrent_state/paused.rs
View file

@ -2,14 +2,14 @@ use std::{collections::HashSet, sync::Arc};
use crate::{
chunk_tracker::{ChunkTracker, HaveNeededSelected},
type_aliases::OpenedFiles,
type_aliases::FileStorage,
};
use super::{streaming::TorrentStreams, ManagedTorrentInfo};
pub struct TorrentStatePaused {
pub(crate) info: Arc<ManagedTorrentInfo>,
pub(crate) files: OpenedFiles,
pub(crate) files: FileStorage,
pub(crate) chunk_tracker: ChunkTracker,
pub(crate) streams: Arc<TorrentStreams>,
}

51
crates/librqbit/src/torrent_state/streaming.rs
View file

@ -1,6 +1,6 @@
use std::{
collections::VecDeque,
io::{Read, Seek, SeekFrom},
io::SeekFrom,
sync::{
atomic::{AtomicUsize, Ordering},
Arc,
@ -10,11 +10,12 @@ use std::{
use anyhow::Context;
use dashmap::DashMap;
use librqbit_core::lengths::{Lengths, ValidPieceIndex};
use tokio::io::{AsyncRead, AsyncSeek};
use tracing::{debug, trace};
use crate::{opened_file::OpenedFile, type_aliases::OpenedFiles, ManagedTorrent};
use crate::{file_info::FileInfo, storage::TorrentStorage, ManagedTorrent};
use super::ManagedTorrentHandle;
@ -236,18 +237,22 @@ impl AsyncRead for FileStream {
"will write bytes"
);
poll_try_io!(poll_try_io!(self.torrent.with_opened_file(
poll_try_io!(poll_try_io!(self.torrent.with_storage_and_file(
self.file_id,
|fd| {
let mut g = fd.file.lock();
g.seek(SeekFrom::Start(self.position))?;
g.read_exact(buf)?;
|files, _fi| {
files.pread_exact(self.file_id, self.position, buf)?;
Ok::<_, anyhow::Error>(())
}
)));
self.as_mut().advance(bytes_to_read as u64);
tbuf.advance(bytes_to_read);
self.streams
.streams
.get_mut(&self.stream_id)
.unwrap()
.value_mut()
.position = self.position;
Poll::Ready(Ok(()))
}
@ -292,30 +297,25 @@ impl Drop for FileStream {
}
impl ManagedTorrent {
fn with_opened_files<F, R>(&self, f: F) -> anyhow::Result<R>
fn with_storage_and_file<F, R>(&self, file_id: usize, f: F) -> anyhow::Result<R>
where
F: FnOnce(&OpenedFiles) -> R,
F: FnOnce(&dyn TorrentStorage, &FileInfo) -> R,
{
self.with_state(|s| {
let files = match s {
crate::ManagedTorrentState::Paused(p) => &p.files,
crate::ManagedTorrentState::Live(l) => &l.files,
s => anyhow::bail!("with_opened_file: invalid state {}", s.name()),
_ => anyhow::bail!("invalid state"),
};
Ok(f(files))
let fi = self
.info()
.file_infos
.get(file_id)
.context("invalid file")?;
Ok(f(files, fi))
})
}
fn with_opened_file<F, R>(&self, file_id: usize, f: F) -> anyhow::Result<R>
where
F: FnOnce(&OpenedFile) -> R,
{
self.with_opened_files(|opened_files| {
let fd = opened_files.get(file_id).context("invalid file id")?;
Ok(f(fd))
})?
}
fn streams(&self) -> anyhow::Result<Arc<TorrentStreams>> {
self.with_state(|s| match s {
crate::ManagedTorrentState::Paused(p) => Ok(p.streams.clone()),
@ -326,7 +326,7 @@ impl ManagedTorrent {
fn maybe_reconnect_needed_peers_for_file(&self, file_id: usize) -> bool {
// If we have the full file, don't bother.
if let Ok(true) = self.with_opened_file(file_id, |f| f.approx_is_finished()) {
if self.is_file_finished(file_id) {
return false;
}
self.with_state(|state| {
@ -337,9 +337,14 @@ impl ManagedTorrent {
true
}
fn is_file_finished(&self, file_id: usize) -> bool {
self.with_chunk_tracker(|ct| ct.is_file_finished(&self.info.file_infos[file_id]))
.unwrap_or(false)
}
pub fn stream(self: Arc<Self>, file_id: usize) -> anyhow::Result<FileStream> {
let (fd_len, fd_offset) =
self.with_opened_file(file_id, |fd| (fd.len, fd.offset_in_torrent))?;
self.with_storage_and_file(file_id, |_fd, fi| (fi.len, fi.offset_in_torrent))?;
let streams = self.streams()?;
let s = FileStream {
stream_id: streams.next_id(),