Mmap custom storage example

Cargo.lock

@@ -1287,6 +1287,7 @@ dependencies = [
  "librqbit-sha1-wrapper",
  "librqbit-tracker-comms",
  "librqbit-upnp",
+ "memmap2",
  "openssl",
  "parking_lot",
  "rand",
@@ -1484,6 +1485,15 @@ version = "2.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d"
 
+[[package]]
+name = "memmap2"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"

crates/librqbit/Cargo.toml

@@ -78,3 +78,4 @@ tracing-subscriber = "0.3"
 tokio-test = "0.4"
 tempfile = "3"
 rand = { version = "0.8", features = ["small_rng"] }
+memmap2 = "0.9.4"

crates/librqbit/examples/custom_storage.rs

@@ -0,0 +1,109 @@
+use std::time::Duration;
+
+use anyhow::Context;
+use librqbit::{
+    storage::{StorageFactory, TorrentStorage},
+    FileInfos, ManagedTorrentInfo, SessionOptions,
+};
+use memmap2::{MmapMut, MmapOptions};
+use parking_lot::RwLock;
+use tracing::info;
+
+struct MmapStorageFactory {}
+
+struct MmapStorage {
+    mmap: RwLock<MmapMut>,
+    file_infos: FileInfos,
+}
+
+impl StorageFactory for MmapStorageFactory {
+    fn init_storage(
+        &self,
+        info: &ManagedTorrentInfo,
+    ) -> anyhow::Result<Box<dyn librqbit::storage::TorrentStorage>> {
+        Ok(Box::new(MmapStorage {
+            mmap: RwLock::new(
+                MmapOptions::new()
+                    .len(info.lengths.total_length().try_into()?)
+                    .map_anon()?,
+            ),
+            file_infos: info.file_infos.clone(),
+        }))
+    }
+}
+
+impl TorrentStorage for MmapStorage {
+    fn pread_exact(&self, file_id: usize, offset: u64, buf: &mut [u8]) -> anyhow::Result<()> {
+        let start: usize = (self.file_infos[file_id].offset_in_torrent + offset).try_into()?;
+        let end = start + buf.len();
+        buf.copy_from_slice(self.mmap.read().get(start..end).context("bad range")?);
+        Ok(())
+    }
+
+    fn pwrite_all(&self, file_id: usize, offset: u64, buf: &[u8]) -> anyhow::Result<()> {
+        let start: usize = (self.file_infos[file_id].offset_in_torrent + offset).try_into()?;
+        let end = start + buf.len();
+        let mut g = self.mmap.write();
+        let target = g.get_mut(start..end).context("bad range")?;
+        target.copy_from_slice(buf);
+        Ok(())
+    }
+
+    fn remove_file(&self, _file_id: usize, _filename: &std::path::Path) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    fn ensure_file_length(&self, _file_id: usize, _length: u64) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    fn take(&self) -> anyhow::Result<Box<dyn TorrentStorage>> {
+        anyhow::bail!("not implemented")
+    }
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    // Output logs to console.
+    match std::env::var("RUST_LOG") {
+        Ok(_) => {}
+        Err(_) => std::env::set_var("RUST_LOG", "info"),
+    }
+    tracing_subscriber::fmt::init();
+    let s = librqbit::Session::new_with_opts(
+        Default::default(),
+        SessionOptions {
+            disable_dht_persistence: true,
+            persistence: false,
+            listen_port_range: None,
+            enable_upnp_port_forwarding: false,
+            ..Default::default()
+        },
+    )
+    .await?;
+    let handle = s
+        .add_torrent(
+            librqbit::AddTorrent::TorrentFileBytes(
+                include_bytes!("../resources/ubuntu-21.04-live-server-amd64.iso.torrent").into(),
+            ),
+            Some(librqbit::AddTorrentOptions {
+                storage_factory: Some(Box::new(MmapStorageFactory {})),
+                paused: false,
+                ..Default::default()
+            }),
+        )
+        .await?
+        .into_handle()
+        .unwrap();
+    tokio::spawn({
+        let h = handle.clone();
+        async move {
+            loop {
+                info!("{}", h.stats());
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+        }
+    });
+    handle.wait_until_completed().await?;
+    Ok(())
+}

crates/librqbit/examples/storage.rs

@@ -1,74 +0,0 @@
-use librqbit::{
-    storage::{StorageFactory, TorrentStorage},
-    ManagedTorrentInfo, SessionOptions,
-};
-
-struct DummyStorage {}
-
-impl StorageFactory for DummyStorage {
-    fn init_storage(
-        &self,
-        _info: &ManagedTorrentInfo,
-    ) -> anyhow::Result<Box<dyn librqbit::storage::TorrentStorage>> {
-        Ok(Box::new(DummyStorage {}))
-    }
-}
-
-impl TorrentStorage for DummyStorage {
-    fn pread_exact(&self, _file_id: usize, _offset: u64, _buf: &mut [u8]) -> anyhow::Result<()> {
-        anyhow::bail!("pread_exact")
-    }
-
-    fn pwrite_all(&self, _file_id: usize, _offset: u64, _buf: &[u8]) -> anyhow::Result<()> {
-        anyhow::bail!("pwrite_all")
-    }
-
-    fn remove_file(&self, _file_id: usize, _filename: &std::path::Path) -> anyhow::Result<()> {
-        anyhow::bail!("remove_file")
-    }
-
-    fn ensure_file_length(&self, _file_id: usize, _length: u64) -> anyhow::Result<()> {
-        anyhow::bail!("ensure_file_length")
-    }
-
-    fn take(&self) -> anyhow::Result<Box<dyn TorrentStorage>> {
-        Ok(Box::new(Self {}))
-    }
-}
-
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    // Output logs to console.
-    match std::env::var("RUST_LOG") {
-        Ok(_) => {}
-        Err(_) => std::env::set_var("RUST_LOG", "info"),
-    }
-    tracing_subscriber::fmt::init();
-    let s = librqbit::Session::new_with_opts(
-        "/does-not-matter".into(),
-        SessionOptions {
-            disable_dht: true,
-            persistence: false,
-            listen_port_range: None,
-            enable_upnp_port_forwarding: false,
-            ..Default::default()
-        },
-    )
-    .await?;
-    let handle = s
-        .add_torrent(
-            librqbit::AddTorrent::TorrentFileBytes(
-                include_bytes!("../resources/ubuntu-21.04-live-server-amd64.iso.torrent").into(),
-            ),
-            Some(librqbit::AddTorrentOptions {
-                storage_factory: Some(Box::new(DummyStorage {})),
-                paused: true,
-                ..Default::default()
-            }),
-        )
-        .await?
-        .into_handle()
-        .unwrap();
-    handle.wait_until_initialized().await?;
-    Ok(())
-}

crates/librqbit/src/lib.rs

@@ -58,6 +58,7 @@ pub use spawn_utils::spawn as librqbit_spawn;
 pub use torrent_state::{
     ManagedTorrent, ManagedTorrentInfo, ManagedTorrentState, TorrentStats, TorrentStatsState,
 };
+pub use type_aliases::FileInfos;
 
 pub use buffers::*;
 pub use clone_to_owned::CloneToOwned;

crates/librqbit/src/type_aliases.rs

@@ -8,6 +8,6 @@ pub type BF = bitvec::boxed::BitBox<u8, bitvec::order::Msb0>;
 
 pub type PeerHandle = SocketAddr;
 pub type PeerStream = BoxStream<'static, SocketAddr>;
-pub(crate) type FileInfos = Vec<FileInfo>;
+pub type FileInfos = Vec<FileInfo>;
 pub(crate) type FileStorage = Box<dyn TorrentStorage>;
 pub(crate) type FilePriorities = Vec<usize>;