Initial impl of up/down ratelimits

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

@@ -83,6 +83,7 @@ use tracing::{debug, error, error_span, info, trace, warn, Instrument};
 use crate::{
     chunk_tracker::{ChunkMarkingResult, ChunkTracker, HaveNeededSelected},
     file_ops::FileOps,
+    limits::Limits,
     peer_connection::{
         PeerConnection, PeerConnectionHandler, PeerConnectionOptions, WriterRequest,
     },
@@ -201,6 +202,12 @@ pub struct TorrentStateLive {
 
     pub(crate) streams: Arc<TorrentStreams>,
     have_broadcast_tx: tokio::sync::broadcast::Sender<ValidPieceIndex>,
+
+    ratelimit_upload_tx: tokio::sync::mpsc::UnboundedSender<(
+        tokio::sync::mpsc::UnboundedSender<WriterRequest>,
+        ChunkInfo,
+    )>,
+    ratelimits: Limits,
 }
 
 impl TorrentStateLive {
@@ -238,6 +245,12 @@ impl TorrentStateLive {
 
         let (have_broadcast_tx, _) = tokio::sync::broadcast::channel(128);
 
+        let (ratelimit_upload_tx, ratelimit_upload_rx) = tokio::sync::mpsc::unbounded_channel::<(
+            tokio::sync::mpsc::UnboundedSender<WriterRequest>,
+            ChunkInfo,
+        )>();
+        let ratelimits = Limits::new(paused.shared.options.ratelimits);
+
         let state = Arc::new(TorrentStateLive {
             torrent: paused.shared.clone(),
             peers: PeerStates {
@@ -272,6 +285,8 @@ impl TorrentStateLive {
             per_piece_locks: (0..lengths.total_pieces())
                 .map(|_| RwLock::new(()))
                 .collect(),
+            ratelimit_upload_tx,
+            ratelimits,
         });
 
         state.spawn(
@@ -304,6 +319,11 @@ impl TorrentStateLive {
             error_span!(parent: state.torrent.span.clone(), "peer_adder"),
             state.clone().task_peer_adder(peer_queue_rx),
         );
+
+        state.spawn(
+            error_span!(parent: state.torrent.span.clone(), "upload_scheduler"),
+            state.clone().task_upload_scheduler(ratelimit_upload_rx),
+        );
         Ok(state)
     }
 
@@ -388,6 +408,26 @@ impl TorrentStateLive {
         Ok(())
     }
 
+    async fn task_upload_scheduler(
+        self: Arc<Self>,
+        mut rx: tokio::sync::mpsc::UnboundedReceiver<(
+            tokio::sync::mpsc::UnboundedSender<WriterRequest>,
+            ChunkInfo,
+        )>,
+    ) -> anyhow::Result<()> {
+        while let Some((tx, ci)) = rx.recv().await {
+            self.ratelimits.prepare_for_upload(ci.size as usize).await;
+            if let Some(session) = self.torrent.session.upgrade() {
+                session
+                    .ratelimits
+                    .prepare_for_upload(ci.size as usize)
+                    .await;
+            }
+            let _ = tx.send(WriterRequest::ReadChunkRequest(ci));
+        }
+        Ok(())
+    }
+
     async fn task_manage_incoming_peer(
         self: Arc<Self>,
         checked_peer: CheckedIncomingConnection,
@@ -1227,12 +1267,10 @@ impl PeerHandler {
             );
         }
 
-        // TODO: this is not super efficient as it does copying multiple times.
-        // Theoretically, this could be done in the sending code, so that it reads straight into
-        // the send buffer.
-        let request = WriterRequest::ReadChunkRequest(chunk_info);
-        trace!("sending {:?}", &request);
-        Ok::<_, anyhow::Error>(self.tx.send(request)?)
+        self.state
+            .ratelimit_upload_tx
+            .send((self.tx.clone(), chunk_info))?;
+        Ok(())
     }
 
     fn on_have(&self, have: u32) {
@@ -1406,6 +1444,18 @@ impl PeerHandler {
                     None => return Ok(()),
                 };
 
+                self.state
+                    .ratelimits
+                    .prepare_for_download(request.length as usize)
+                    .await;
+
+                if let Some(session) = self.state.torrent().session.upgrade() {
+                    session
+                        .ratelimits
+                        .prepare_for_download(request.length as usize)
+                        .await;
+                }
+
                 loop {
                     match aframe!(tokio::time::timeout(
                         Duration::from_secs(5),

crates/librqbit/src/torrent_state/mod.rs

@@ -36,6 +36,7 @@ use tracing::trace;
 use tracing::warn;
 
 use crate::chunk_tracker::ChunkTracker;
+use crate::limits::LimitsConfig;
 use crate::session::TorrentId;
 use crate::spawn_utils::BlockingSpawner;
 use crate::storage::BoxStorageFactory;
@@ -103,6 +104,7 @@ pub(crate) struct ManagedTorrentOptions {
     pub allow_overwrite: bool,
     pub output_folder: PathBuf,
     pub disk_write_queue: Option<DiskWorkQueueSender>,
+    pub ratelimits: LimitsConfig,
     #[cfg(feature = "disable-upload")]
     pub _disable_upload: bool,
 }