Add storage example

crates/librqbit/src/chunk_tracker.rs

@@ -171,23 +171,6 @@ impl ChunkTracker {
         }
     }
 
-    pub fn new_empty(lengths: Lengths, file_infos: &FileInfos) -> anyhow::Result<Self> {
-        let have = BF::from_boxed_slice(vec![0; lengths.piece_bitfield_bytes()].into_boxed_slice());
-        let selected = have.clone();
-        let chunk_status =
-            BF::from_boxed_slice(vec![0; lengths.chunk_bitfield_bytes()].into_boxed_slice());
-        let queued = have.clone();
-        Ok(Self {
-            queue_pieces: queued,
-            chunk_status,
-            have,
-            selected,
-            lengths,
-            per_file_bytes: vec![0; file_infos.len()],
-            hns: Default::default(),
-        })
-    }
-
     pub fn get_lengths(&self) -> &Lengths {
         &self.lengths
     }

crates/librqbit/src/lib.rs

@@ -41,7 +41,7 @@ mod peer_info_reader;
 mod read_buf;
 mod session;
 mod spawn_utils;
-mod storage;
+pub mod storage;
 mod torrent_state;
 pub mod tracing_subscriber_config_utils;
 mod type_aliases;
@@ -56,7 +56,9 @@ pub use session::{
     SUPPORTED_SCHEMES,
 };
 pub use spawn_utils::spawn as librqbit_spawn;
-pub use torrent_state::{ManagedTorrent, ManagedTorrentState, TorrentStats, TorrentStatsState};
+pub use torrent_state::{
+    ManagedTorrent, ManagedTorrentInfo, ManagedTorrentState, TorrentStats, TorrentStatsState,
+};
 
 pub use buffers::*;
 pub use clone_to_owned::CloneToOwned;

crates/librqbit/src/session.rs

@@ -15,6 +15,7 @@ use crate::{
     peer_connection::PeerConnectionOptions,
     read_buf::ReadBuf,
     spawn_utils::BlockingSpawner,
+    storage::StorageFactory,
     torrent_state::{
         ManagedTorrentBuilder, ManagedTorrentHandle, ManagedTorrentState, TorrentStateLive,
     },
@@ -43,7 +44,6 @@ use librqbit_core::{
 use parking_lot::RwLock;
 use peer_binary_protocol::Handshake;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
-use serde_with::serde_as;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_stream::StreamExt;
 use tokio_util::sync::{CancellationToken, DropGuard};
@@ -270,9 +270,7 @@ fn merge_two_optional_streams<T>(
 }
 
 /// Options for adding new torrents to the session.
-#[serde_as]
-#[derive(Default, Clone, Serialize, Deserialize)]
-#[serde(default)]
+#[derive(Default)]
 pub struct AddTorrentOptions {
     /// Start in paused state.
     pub paused: bool,
@@ -295,7 +293,6 @@ pub struct AddTorrentOptions {
     pub peer_opts: Option<PeerConnectionOptions>,
 
     /// Force a refresh interval for polling trackers.
-    #[serde_as(as = "Option<serde_with::DurationSeconds>")]
     pub force_tracker_interval: Option<Duration>,
 
     pub disable_trackers: bool,
@@ -304,8 +301,9 @@ pub struct AddTorrentOptions {
     pub initial_peers: Option<Vec<SocketAddr>>,
 
     /// This is used to restore the session from serialized state.
-    #[serde(skip)]
     pub preferred_id: Option<usize>,
+
+    pub storage_factory: Option<Box<dyn StorageFactory>>,
 }
 
 pub struct ListOnlyResponse {
@@ -976,7 +974,7 @@ impl Session {
         trackers: Vec<String>,
         peer_rx: Option<PeerStream>,
         initial_peers: Vec<SocketAddr>,
-        opts: AddTorrentOptions,
+        mut opts: AddTorrentOptions,
     ) -> anyhow::Result<AddTorrentResponse> {
         debug!("Torrent info: {:#?}", &info);
 
@@ -1031,6 +1029,10 @@ impl Session {
             builder.peer_read_write_timeout(t);
         }
 
+        if let Some(storage_factory) = opts.storage_factory.take() {
+            builder.storage_factory(storage_factory);
+        }
+
         let (managed_torrent, id) = {
             let mut g = self.db.write();
             if let Some((id, handle)) = g.torrents.iter().find(|(_, t)| t.info_hash() == info_hash)

crates/librqbit/src/storage/example.rs

@@ -6,7 +6,7 @@ use parking_lot::RwLock;
 
 use crate::type_aliases::FileInfos;
 
-use super::TorrentStorage;
+use super::{StorageFactory, TorrentStorage};
 
 struct InMemoryPiece {
     bytes: Box<[u8]>,
@@ -19,7 +19,21 @@ impl InMemoryPiece {
     }
 }
 
-pub struct InMemoryExampleStorage {
+pub struct InMemoryExampleStorageFactory {}
+
+impl StorageFactory for InMemoryExampleStorageFactory {
+    fn init_storage(
+        &self,
+        info: &crate::torrent_state::ManagedTorrentInfo,
+    ) -> anyhow::Result<Box<dyn TorrentStorage>> {
+        Ok(Box::new(InMemoryExampleStorage::new(
+            info.lengths,
+            info.file_infos.clone(),
+        )?))
+    }
+}
+
+struct InMemoryExampleStorage {
     lengths: Lengths,
     file_infos: FileInfos,
     map: RwLock<HashMap<ValidPieceIndex, InMemoryPiece>>,
@@ -28,7 +42,7 @@ pub struct InMemoryExampleStorage {
 }
 
 impl InMemoryExampleStorage {
-    pub fn new(lengths: Lengths, file_infos: FileInfos) -> anyhow::Result<Self> {
+    fn new(lengths: Lengths, file_infos: FileInfos) -> anyhow::Result<Self> {
         // Max memory 128MiB. Make it tunable
         let max_pieces = 128 * 1024 * 1024 / lengths.default_piece_length();
         if max_pieces == 0 {

crates/librqbit/src/torrent_state/live/mod.rs

@@ -86,7 +86,6 @@ use crate::{
         PeerConnection, PeerConnectionHandler, PeerConnectionOptions, WriterRequest,
     },
     session::CheckedIncomingConnection,
-    storage::TorrentStorage,
     torrent_state::{peer::Peer, utils::atomic_inc},
     type_aliases::{FilePriorities, FileStorage, PeerHandle, BF},
 };