From 5027d8ccd153b4368dac0daeee678c76f1680793 Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Wed, 1 May 2024 22:14:34 +0100 Subject: [PATCH] All storage factories are generic now --- crates/librqbit/examples/custom_storage.rs | 7 ++-- crates/librqbit/src/session.rs | 6 ++-- crates/librqbit/src/storage/example.rs | 12 +++---- crates/librqbit/src/storage/filesystem/mod.rs | 8 +++-- crates/librqbit/src/storage/mmap.rs | 14 ++++---- crates/librqbit/src/storage/mod.rs | 33 +++++++++++++++++-- crates/librqbit/src/storage/slow.rs | 13 ++++---- crates/librqbit/src/storage/timing.rs | 13 ++++---- crates/librqbit/src/tests/e2e_stream.rs | 7 ++-- .../src/torrent_state/initializing.rs | 8 +++-- crates/librqbit/src/torrent_state/mod.rs | 10 +++--- crates/rqbit/src/main.rs | 9 +++-- 12 files changed, 88 insertions(+), 52 deletions(-) diff --git a/crates/librqbit/examples/custom_storage.rs b/crates/librqbit/examples/custom_storage.rs index f0fa5a4..3f568dc 100644 --- a/crates/librqbit/examples/custom_storage.rs +++ b/crates/librqbit/examples/custom_storage.rs @@ -1,6 +1,9 @@ use std::time::Duration; -use librqbit::{storage::mmap::MmapStorageFactory, SessionOptions}; +use librqbit::{ + storage::{mmap::MmapStorageFactory, StorageFactoryExt}, + SessionOptions, +}; use tracing::info; #[tokio::main] @@ -28,7 +31,7 @@ async fn main() -> anyhow::Result<()> { include_bytes!("../resources/ubuntu-21.04-live-server-amd64.iso.torrent").into(), ), Some(librqbit::AddTorrentOptions { - storage_factory: Some(Box::new(MmapStorageFactory {})), + storage_factory: Some(MmapStorageFactory::default().boxed()), paused: false, ..Default::default() }), diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs index 5486eaf..a7e660b 100644 --- a/crates/librqbit/src/session.rs +++ b/crates/librqbit/src/session.rs @@ -16,7 +16,7 @@ use crate::{ peer_connection::PeerConnectionOptions, read_buf::ReadBuf, spawn_utils::BlockingSpawner, - storage::{filesystem::FilesystemStorageFactory, StorageFactory}, + storage::{filesystem::FilesystemStorageFactory, BoxStorageFactory, StorageFactoryExt}, torrent_state::{ ManagedTorrentBuilder, ManagedTorrentHandle, ManagedTorrentState, TorrentStateLive, }, @@ -305,7 +305,7 @@ pub struct AddTorrentOptions { /// This is used to restore the session from serialized state. pub preferred_id: Option, - pub storage_factory: Option>, + pub storage_factory: Option, } pub struct ListOnlyResponse { @@ -1004,7 +1004,7 @@ impl Session { let storage_factory = opts .storage_factory .take() - .unwrap_or_else(|| Box::::default()); + .unwrap_or_else(|| FilesystemStorageFactory::default().boxed()); if opts.list_only { return Ok(AddTorrentResponse::ListOnly(ListOnlyResponse { diff --git a/crates/librqbit/src/storage/example.rs b/crates/librqbit/src/storage/example.rs index 8dc32a0..7fde403 100644 --- a/crates/librqbit/src/storage/example.rs +++ b/crates/librqbit/src/storage/example.rs @@ -19,21 +19,21 @@ impl InMemoryPiece { } } +#[derive(Default)] pub struct InMemoryExampleStorageFactory {} impl StorageFactory for InMemoryExampleStorageFactory { + type Storage = InMemoryExampleStorage; + fn init_storage( &self, info: &crate::torrent_state::ManagedTorrentInfo, - ) -> anyhow::Result> { - Ok(Box::new(InMemoryExampleStorage::new( - info.lengths, - info.file_infos.clone(), - )?)) + ) -> anyhow::Result { + InMemoryExampleStorage::new(info.lengths, info.file_infos.clone()) } } -struct InMemoryExampleStorage { +pub struct InMemoryExampleStorage { lengths: Lengths, file_infos: FileInfos, map: RwLock>, diff --git a/crates/librqbit/src/storage/filesystem/mod.rs b/crates/librqbit/src/storage/filesystem/mod.rs index bcb8bb5..3ce2315 100644 --- a/crates/librqbit/src/storage/filesystem/mod.rs +++ b/crates/librqbit/src/storage/filesystem/mod.rs @@ -18,7 +18,9 @@ use super::{StorageFactory, TorrentStorage}; pub struct FilesystemStorageFactory {} impl StorageFactory for FilesystemStorageFactory { - fn init_storage(&self, meta: &ManagedTorrentInfo) -> anyhow::Result> { + type Storage = FilesystemStorage; + + fn init_storage(&self, meta: &ManagedTorrentInfo) -> anyhow::Result { let mut files = Vec::::new(); let output_folder = &meta.options.output_folder; for file_details in meta.info.iter_file_details(&meta.lengths)? { @@ -54,10 +56,10 @@ impl StorageFactory for FilesystemStorageFactory { }; files.push(OpenedFile::new(file)); } - Ok(Box::new(FilesystemStorage { + Ok(FilesystemStorage { output_folder: output_folder.clone(), opened_files: files, - })) + }) } } diff --git a/crates/librqbit/src/storage/mmap.rs b/crates/librqbit/src/storage/mmap.rs index dbd1ee2..ffeb533 100644 --- a/crates/librqbit/src/storage/mmap.rs +++ b/crates/librqbit/src/storage/mmap.rs @@ -6,26 +6,26 @@ use crate::{FileInfos, ManagedTorrentInfo}; use super::{StorageFactory, TorrentStorage}; +#[derive(Default)] pub struct MmapStorageFactory {} -struct MmapStorage { +pub struct MmapStorage { mmap: RwLock, file_infos: FileInfos, } impl StorageFactory for MmapStorageFactory { - fn init_storage( - &self, - info: &ManagedTorrentInfo, - ) -> anyhow::Result> { - Ok(Box::new(MmapStorage { + type Storage = MmapStorage; + + fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result { + Ok(MmapStorage { mmap: RwLock::new( MmapOptions::new() .len(info.lengths.total_length().try_into()?) .map_anon()?, ), file_infos: info.file_infos.clone(), - })) + }) } } diff --git a/crates/librqbit/src/storage/mod.rs b/crates/librqbit/src/storage/mod.rs index a557617..6e0094d 100644 --- a/crates/librqbit/src/storage/mod.rs +++ b/crates/librqbit/src/storage/mod.rs @@ -9,11 +9,40 @@ use std::{any::Any, path::Path}; use crate::torrent_state::ManagedTorrentInfo; pub trait StorageFactory: Send + Sync + Any { - fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result>; + type Storage: TorrentStorage; + + fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result; +} + +pub type BoxStorageFactory = Box>>; + +pub trait StorageFactoryExt { + fn boxed(self) -> BoxStorageFactory; +} + +impl StorageFactoryExt for SF { + fn boxed(self) -> BoxStorageFactory { + struct BoxFactory { + sf: SF, + } + + impl StorageFactory for BoxFactory { + type Storage = Box; + + fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result { + let s = self.sf.init_storage(info)?; + Ok(Box::new(s)) + } + } + + Box::new(BoxFactory { sf: self }) + } } impl StorageFactory for Box { - fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result> { + type Storage = U::Storage; + + fn init_storage(&self, info: &ManagedTorrentInfo) -> anyhow::Result { (**self).init_storage(info) } } diff --git a/crates/librqbit/src/storage/slow.rs b/crates/librqbit/src/storage/slow.rs index e822d7b..9443a1d 100644 --- a/crates/librqbit/src/storage/slow.rs +++ b/crates/librqbit/src/storage/slow.rs @@ -17,17 +17,16 @@ impl SlowStorageFactory { } impl StorageFactory for SlowStorageFactory { - fn init_storage( - &self, - info: &crate::ManagedTorrentInfo, - ) -> anyhow::Result> { - Ok(Box::new(SlowStorage { + type Storage = SlowStorage; + + fn init_storage(&self, info: &crate::ManagedTorrentInfo) -> anyhow::Result { + Ok(SlowStorage { underlying: self.underlying_factory.init_storage(info)?, - })) + }) } } -struct SlowStorage { +pub struct SlowStorage { underlying: U, } diff --git a/crates/librqbit/src/storage/timing.rs b/crates/librqbit/src/storage/timing.rs index b3cf7cd..7893765 100644 --- a/crates/librqbit/src/storage/timing.rs +++ b/crates/librqbit/src/storage/timing.rs @@ -15,18 +15,17 @@ impl TimingStorageFactory { } impl StorageFactory for TimingStorageFactory { - fn init_storage( - &self, - info: &crate::ManagedTorrentInfo, - ) -> anyhow::Result> { - Ok(Box::new(TimingStorage { + type Storage = TimingStorage; + + fn init_storage(&self, info: &crate::ManagedTorrentInfo) -> anyhow::Result { + Ok(TimingStorage { name: self.name.clone(), underlying: self.underlying_factory.init_storage(info)?, - })) + }) } } -struct TimingStorage { +pub struct TimingStorage { name: String, underlying: U, } diff --git a/crates/librqbit/src/tests/e2e_stream.rs b/crates/librqbit/src/tests/e2e_stream.rs index 1125e6d..c06857f 100644 --- a/crates/librqbit/src/tests/e2e_stream.rs +++ b/crates/librqbit/src/tests/e2e_stream.rs @@ -5,8 +5,9 @@ use tokio::{io::AsyncReadExt, time::timeout}; use tracing::info; use crate::{ - create_torrent, storage::example::InMemoryExampleStorageFactory, AddTorrent, - CreateTorrentOptions, Session, + create_torrent, + storage::{example::InMemoryExampleStorageFactory, StorageFactoryExt}, + AddTorrent, CreateTorrentOptions, Session, }; use super::test_util::create_default_random_dir_with_torrents; @@ -85,7 +86,7 @@ async fn e2e_stream() -> anyhow::Result<()> { AddTorrent::from_bytes(torrent.as_bytes()?), Some(crate::AddTorrentOptions { paused: false, - storage_factory: Some(Box::new(InMemoryExampleStorageFactory {})), + storage_factory: Some(InMemoryExampleStorageFactory::default().boxed()), initial_peers: Some(vec![peer]), ..Default::default() }), diff --git a/crates/librqbit/src/torrent_state/initializing.rs b/crates/librqbit/src/torrent_state/initializing.rs index 7944b6b..684d130 100644 --- a/crates/librqbit/src/torrent_state/initializing.rs +++ b/crates/librqbit/src/torrent_state/initializing.rs @@ -8,7 +8,11 @@ use anyhow::Context; use size_format::SizeFormatterBinary as SF; use tracing::{debug, info, warn}; -use crate::{chunk_tracker::ChunkTracker, file_ops::FileOps, storage::StorageFactory}; +use crate::{ + chunk_tracker::ChunkTracker, + file_ops::FileOps, + storage::{BoxStorageFactory, StorageFactory}, +}; use super::{paused::TorrentStatePaused, ManagedTorrentInfo}; @@ -34,7 +38,7 @@ impl TorrentStateInitializing { pub async fn check( &self, - storage_factory: &dyn StorageFactory, + storage_factory: &BoxStorageFactory, ) -> anyhow::Result { let files = storage_factory.init_storage(&self.meta)?; info!("Doing initial checksum validation, this might take a while..."); diff --git a/crates/librqbit/src/torrent_state/mod.rs b/crates/librqbit/src/torrent_state/mod.rs index a0f0e8d..65b806d 100644 --- a/crates/librqbit/src/torrent_state/mod.rs +++ b/crates/librqbit/src/torrent_state/mod.rs @@ -36,7 +36,7 @@ use tracing::warn; use crate::chunk_tracker::ChunkTracker; use crate::file_info::FileInfo; use crate::spawn_utils::BlockingSpawner; -use crate::storage::StorageFactory; +use crate::storage::BoxStorageFactory; use crate::torrent_state::stats::LiveStats; use crate::type_aliases::FileInfos; use crate::type_aliases::PeerStream; @@ -108,7 +108,7 @@ pub struct ManagedTorrentInfo { pub struct ManagedTorrent { pub info: Arc, - pub(crate) storage_factory: Box, + pub(crate) storage_factory: BoxStorageFactory, state_change_notify: Notify, locked: RwLock, @@ -273,7 +273,7 @@ impl ManagedTorrent { error_span!(parent: span.clone(), "initialize_and_start"), token.clone(), async move { - match init.check(&*t.storage_factory).await { + match init.check(&t.storage_factory).await { Ok(paused) => { let mut g = t.locked.write(); if let ManagedTorrentState::Initializing(_) = &g.state { @@ -504,7 +504,7 @@ pub(crate) struct ManagedTorrentBuilder { peer_id: Option, spawner: Option, allow_overwrite: bool, - storage_factory: Box, + storage_factory: BoxStorageFactory, } impl ManagedTorrentBuilder { @@ -512,7 +512,7 @@ impl ManagedTorrentBuilder { info: TorrentMetaV1Info, info_hash: Id20, output_folder: PathBuf, - storage_factory: Box, + storage_factory: BoxStorageFactory, ) -> Self { Self { info, diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs index f288ad1..793bae5 100644 --- a/crates/rqbit/src/main.rs +++ b/crates/rqbit/src/main.rs @@ -8,8 +8,7 @@ use librqbit::{ http_api::{HttpApi, HttpApiOptions}, http_api_client, librqbit_spawn, storage::{ - filesystem::FilesystemStorageFactory, slow::SlowStorageFactory, - timing::TimingStorageFactory, + filesystem::FilesystemStorageFactory, timing::TimingStorageFactory, StorageFactoryExt, }, tracing_subscriber_config_utils::{init_logging, InitLoggingOptions}, AddTorrent, AddTorrentOptions, AddTorrentResponse, Api, ListOnlyResponse, @@ -382,9 +381,9 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> { initial_peers: download_opts.initial_peers.clone().map(|p| p.0), disable_trackers: download_opts.disable_trackers, storage_factory: Some({ - let sf = Box::::default(); - // let sf = Box::new(SlowStorageFactory::new(sf)); - Box::new(TimingStorageFactory::new("fs".to_owned(), sf)) + let sf = FilesystemStorageFactory::default(); + // let sf = SlowStorageFactory::new(sf); + TimingStorageFactory::new("fs".to_owned(), sf).boxed() }), ..Default::default() };