From d3413370d214a472bfa63b20f17e4524d8664991 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Wed, 1 May 2024 20:47:33 +0100
Subject: [PATCH] Slow storage impl

---
 crates/librqbit/src/storage/mod.rs  |  1 +
 crates/librqbit/src/storage/slow.rs | 65 +++++++++++++++++++++++++++++
 crates/rqbit/src/main.rs            |  5 +++
 3 files changed, 71 insertions(+)
 create mode 100644 crates/librqbit/src/storage/slow.rs

diff --git a/crates/librqbit/src/storage/mod.rs b/crates/librqbit/src/storage/mod.rs
index 9401385..82354cf 100644
--- a/crates/librqbit/src/storage/mod.rs
+++ b/crates/librqbit/src/storage/mod.rs
@@ -1,6 +1,7 @@
 pub mod example;
 pub mod filesystem;
 pub mod mmap;
+pub mod slow;
 
 use std::{any::Any, path::Path};
 
diff --git a/crates/librqbit/src/storage/slow.rs b/crates/librqbit/src/storage/slow.rs
new file mode 100644
index 0000000..6a5d63e
--- /dev/null
+++ b/crates/librqbit/src/storage/slow.rs
@@ -0,0 +1,65 @@
+use std::time::Duration;
+
+use rand::Rng;
+
+use super::{StorageFactory, TorrentStorage};
+
+pub struct SlowStorageFactory {
+    underlying_factory: Box<dyn StorageFactory>,
+}
+
+impl SlowStorageFactory {
+    pub fn new(underlying: Box<dyn StorageFactory>) -> Self {
+        Self {
+            underlying_factory: underlying,
+        }
+    }
+}
+
+impl StorageFactory for SlowStorageFactory {
+    fn init_storage(
+        &self,
+        info: &crate::ManagedTorrentInfo,
+    ) -> anyhow::Result<Box<dyn TorrentStorage>> {
+        Ok(Box::new(SlowStorage {
+            underlying: self.underlying_factory.init_storage(info)?,
+        }))
+    }
+}
+
+struct SlowStorage {
+    underlying: Box<dyn TorrentStorage>,
+}
+
+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");
+    std::thread::sleep(sl)
+}
+
+impl TorrentStorage for SlowStorage {
+    fn pread_exact(&self, file_id: usize, offset: u64, buf: &mut [u8]) -> anyhow::Result<()> {
+        random_sleep();
+        self.underlying.pread_exact(file_id, offset, buf)
+    }
+
+    fn pwrite_all(&self, file_id: usize, offset: u64, buf: &[u8]) -> anyhow::Result<()> {
+        random_sleep();
+        self.underlying.pwrite_all(file_id, offset, buf)
+    }
+
+    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()?,
+        }))
+    }
+}
diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs
index 3ca442b..c2f7103 100644
--- a/crates/rqbit/src/main.rs
+++ b/crates/rqbit/src/main.rs
@@ -7,6 +7,7 @@ use librqbit::{
     api::ApiAddTorrentResponse,
     http_api::{HttpApi, HttpApiOptions},
     http_api_client, librqbit_spawn,
+    storage::{filesystem::FilesystemStorageFactory, slow::SlowStorageFactory},
     tracing_subscriber_config_utils::{init_logging, InitLoggingOptions},
     AddTorrent, AddTorrentOptions, AddTorrentResponse, Api, ListOnlyResponse,
     PeerConnectionOptions, Session, SessionOptions, TorrentStatsState,
@@ -376,6 +377,10 @@ 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(
+                )))),
                 ..Default::default()
             };
             let connect_to_existing = match client.validate_rqbit_server().await {