Add timing storage

crates/librqbit/src/storage/mod.rs

@@ -2,6 +2,7 @@ pub mod example;
 pub mod filesystem;
 pub mod mmap;
 pub mod slow;
+pub mod timing;
 
 use std::{any::Any, path::Path};
 

crates/librqbit/src/storage/slow.rs

@@ -34,7 +34,7 @@ struct SlowStorage {
 fn random_sleep() {
     let sl = rand::thread_rng().gen_range(0f64..0.1f64);
     let sl = Duration::from_secs_f64(sl);
-    tracing::debug!(duration = ?sl, "sleeping");
+    tracing::trace!(duration = ?sl, "sleeping");
     std::thread::sleep(sl)
 }
 

crates/librqbit/src/storage/timing.rs

@@ -0,0 +1,83 @@
+use super::{StorageFactory, TorrentStorage};
+
+pub struct TimingStorageFactory {
+    name: String,
+    underlying_factory: Box<dyn StorageFactory>,
+}
+
+impl TimingStorageFactory {
+    pub fn new(name: String, underlying: Box<dyn StorageFactory>) -> Self {
+        Self {
+            name,
+            underlying_factory: underlying,
+        }
+    }
+}
+
+impl StorageFactory for TimingStorageFactory {
+    fn init_storage(
+        &self,
+        info: &crate::ManagedTorrentInfo,
+    ) -> anyhow::Result<Box<dyn TorrentStorage>> {
+        Ok(Box::new(TimingStorage {
+            name: self.name.clone(),
+            underlying: self.underlying_factory.init_storage(info)?,
+        }))
+    }
+}
+
+struct TimingStorage {
+    name: String,
+    underlying: Box<dyn TorrentStorage>,
+}
+
+macro_rules! timeit {
+    ($name:expr, $op:expr, $($rest:tt),*) => {
+        {
+            let start = std::time::Instant::now();
+            let r = $op;
+            let elapsed = start.elapsed();
+            tracing::debug!(name = $name, $($rest),*, ?elapsed, "timeit");
+            r
+        }
+    };
+}
+
+impl TorrentStorage for TimingStorage {
+    fn pread_exact(&self, file_id: usize, offset: u64, buf: &mut [u8]) -> anyhow::Result<()> {
+        let storage = &self.name;
+        timeit!(
+            "pread_exact",
+            self.underlying.pread_exact(file_id, offset, buf),
+            file_id,
+            offset,
+            storage
+        )
+    }
+
+    fn pwrite_all(&self, file_id: usize, offset: u64, buf: &[u8]) -> anyhow::Result<()> {
+        let storage = &self.name;
+        timeit!(
+            "pwrite_all",
+            self.underlying.pwrite_all(file_id, offset, buf),
+            file_id,
+            offset,
+            storage
+        )
+    }
+
+    fn remove_file(&self, file_id: usize, filename: &std::path::Path) -> anyhow::Result<()> {
+        self.underlying.remove_file(file_id, filename)
+    }
+
+    fn ensure_file_length(&self, file_id: usize, length: u64) -> anyhow::Result<()> {
+        self.underlying.ensure_file_length(file_id, length)
+    }
+
+    fn take(&self) -> anyhow::Result<Box<dyn TorrentStorage>> {
+        Ok(Box::new(Self {
+            underlying: self.underlying.take()?,
+            name: self.name.clone(),
+        }))
+    }
+}

crates/rqbit/src/main.rs

@@ -7,7 +7,10 @@ use librqbit::{
     api::ApiAddTorrentResponse,
     http_api::{HttpApi, HttpApiOptions},
     http_api_client, librqbit_spawn,
-    storage::{filesystem::FilesystemStorageFactory, slow::SlowStorageFactory},
+    storage::{
+        filesystem::FilesystemStorageFactory, slow::SlowStorageFactory,
+        timing::TimingStorageFactory,
+    },
     tracing_subscriber_config_utils::{init_logging, InitLoggingOptions},
     AddTorrent, AddTorrentOptions, AddTorrentResponse, Api, ListOnlyResponse,
     PeerConnectionOptions, Session, SessionOptions, TorrentStatsState,
@@ -368,6 +371,7 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
             }
             let http_api_url = format!("http://{}", opts.http_api_listen_addr);
             let client = http_api_client::HttpApiClient::new(&http_api_url)?;
+
             let torrent_opts = || AddTorrentOptions {
                 only_files_regex: download_opts.only_files_matching_regex.clone(),
                 overwrite: download_opts.overwrite,
@@ -377,10 +381,11 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
                 sub_folder: download_opts.sub_folder.clone(),
                 initial_peers: download_opts.initial_peers.clone().map(|p| p.0),
                 disable_trackers: download_opts.disable_trackers,
-                storage_factory: Some(Box::new(SlowStorageFactory::new(Box::<
-                    FilesystemStorageFactory,
-                >::default(
-                )))),
+                storage_factory: Some({
+                    let sf = Box::<FilesystemStorageFactory>::default();
+                    let sf = Box::new(SlowStorageFactory::new(sf));
+                    Box::new(TimingStorageFactory::new("hdd".to_owned(), sf))
+                }),
                 ..Default::default()
             };
             let connect_to_existing = match client.validate_rqbit_server().await {