diff --git a/soulseek-rs-lib/src/client/context.rs b/soulseek-rs-lib/src/client/context.rs index 8606abf..ff98041 100644 --- a/soulseek-rs-lib/src/client/context.rs +++ b/soulseek-rs-lib/src/client/context.rs @@ -42,6 +42,7 @@ impl ClientContext { private_messages: Vec::new(), pending_connect_tokens: HashMap::new(), max_peers: Arc::new(AtomicUsize::new(DEFAULT_MAX_PEERS)), + download_speed_limit: Arc::default(), shares: Arc::new(Shares::empty()), shared_directories: Vec::new(), peer_addresses: HashMap::new(), diff --git a/soulseek-rs-lib/src/client/downloads.rs b/soulseek-rs-lib/src/client/downloads.rs index 28c8177..e47db5d 100644 --- a/soulseek-rs-lib/src/client/downloads.rs +++ b/soulseek-rs-lib/src/client/downloads.rs @@ -14,6 +14,12 @@ impl Client { .unwrap_or_default() } + pub fn set_download_speed_limit(&self, bytes_per_second: u64) { + if let Ok(ctx) = self.context.read_safe() { + ctx.download_speed_limit.set(bytes_per_second); + } + } + #[must_use] pub fn pause_download(&self, username: &str, filename: &str) -> bool { match self.context.write_safe() { diff --git a/soulseek-rs-lib/src/client/mod.rs b/soulseek-rs-lib/src/client/mod.rs index 21d7dfd..7a146b1 100644 --- a/soulseek-rs-lib/src/client/mod.rs +++ b/soulseek-rs-lib/src/client/mod.rs @@ -9,6 +9,7 @@ use crate::types::{ SimilarUser, UserInfo, UserInterests, UserPresence, UserStats, UserStatus, }; use crate::utils::logger; +use crate::utils::rate_limit::RateLimit; use crate::{ Transfer, actor::{ActorSystem, peer_registry::PeerRegistry}, @@ -486,6 +487,7 @@ pub struct ClientContext { /// a token we sent in a ConnectToPeer to the peer we expect back. pending_connect_tokens: HashMap, max_peers: Arc, + pub(crate) download_speed_limit: Arc, /// Files we share with peers (read-only after connect). pub shares: Arc, /// The directories the current share index was built from. diff --git a/soulseek-rs-lib/src/peer/download_peer.rs b/soulseek-rs-lib/src/peer/download_peer.rs index 034e54f..fd2e231 100644 --- a/soulseek-rs-lib/src/peer/download_peer.rs +++ b/soulseek-rs-lib/src/peer/download_peer.rs @@ -363,6 +363,11 @@ impl DownloadPeer { stream .set_read_timeout(Some(CANCEL_POLL)) .map_err(DownloadError::ConnectionFailed)?; + let speed_limit = client_context + .read() + .map_err(|_| DownloadError::LockPoisoned)? + .download_speed_limit + .clone(); let mut last_data = Instant::now(); loop { @@ -384,7 +389,10 @@ impl DownloadPeer { break; } - match stream.read(&mut read_buffer) { + let allowance = speed_limit.take(READ_BUFFER_SIZE); + let read = stream.read(&mut read_buffer[..allowance]); + speed_limit.give_back(allowance - read.as_ref().map_or(0, |&n| n)); + match read { Ok(0) => { trace!( "[download_peer:{}] connection closed by peer", diff --git a/soulseek-rs-lib/src/utils/mod.rs b/soulseek-rs-lib/src/utils/mod.rs index e32ce7c..b7ea3bd 100644 --- a/soulseek-rs-lib/src/utils/mod.rs +++ b/soulseek-rs-lib/src/utils/mod.rs @@ -4,6 +4,7 @@ pub(crate) mod keepalive; pub mod lock; pub mod md5; pub mod path; +pub(crate) mod rate_limit; pub(crate) mod semaphore; pub mod zlib; diff --git a/soulseek-rs-lib/src/utils/rate_limit.rs b/soulseek-rs-lib/src/utils/rate_limit.rs new file mode 100644 index 0000000..a5a261a --- /dev/null +++ b/soulseek-rs-lib/src/utils/rate_limit.rs @@ -0,0 +1,136 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Mutex, PoisonError}; +use std::time::{Duration, Instant}; + +const TICK: Duration = Duration::from_millis(100); +const TICKS_PER_SECOND: u64 = 1000 / TICK.as_millis() as u64; + +#[derive(Default)] +pub struct RateLimit { + bytes_per_second: AtomicU64, + bucket: Mutex, +} + +#[derive(Default)] +struct Bucket { + tokens: u64, + refilled: Option, +} + +impl RateLimit { + pub fn set(&self, bytes_per_second: u64) { + self.bytes_per_second + .store(bytes_per_second, Ordering::Relaxed); + } + + pub fn take(&self, want: usize) -> usize { + loop { + let rate = self.bytes_per_second.load(Ordering::Relaxed); + if rate == 0 { + return want; + } + let mut bucket = + self.bucket.lock().unwrap_or_else(PoisonError::into_inner); + let now = Instant::now(); + let since = bucket.refilled.map_or(TICK, |at| now - at); + if since >= TICK { + bucket.tokens = (rate / TICKS_PER_SECOND).max(1); + bucket.refilled = Some(now); + } + if bucket.tokens > 0 { + let grant = bucket.tokens.min(want as u64); + bucket.tokens -= grant; + return grant as usize; + } + drop(bucket); + std::thread::sleep(TICK.saturating_sub(since)); + } + } + + pub fn give_back(&self, unused: usize) { + if unused == 0 || self.bytes_per_second.load(Ordering::Relaxed) == 0 { + return; + } + self.bucket + .lock() + .unwrap_or_else(PoisonError::into_inner) + .tokens += unused as u64; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn no_limit_grants_everything_at_once() { + let limit = RateLimit::default(); + let started = Instant::now(); + assert_eq!(limit.take(1 << 20), 1 << 20); + assert!(started.elapsed() < TICK); + } + + #[test] + fn a_limit_paces_what_is_taken() { + let limit = RateLimit::default(); + limit.set(20_000); + let started = Instant::now(); + let mut taken = 0; + while taken < 6_000 { + let grant = limit.take(4096); + assert!((1..=4096).contains(&grant), "granted {grant}"); + taken += grant; + } + let elapsed = started.elapsed(); + assert!( + (2 * TICK..10 * TICK).contains(&elapsed), + "{taken} bytes at 20 kB/s took {elapsed:?}" + ); + } + + #[test] + fn concurrent_takers_share_one_budget() { + let limit = RateLimit::default(); + limit.set(20_000); + let started = Instant::now(); + std::thread::scope(|s| { + for _ in 0..3 { + s.spawn(|| { + let mut taken = 0; + while taken < 2_000 { + taken += limit.take(4096); + } + }); + } + }); + assert!( + started.elapsed() >= 2 * TICK, + "6 kB at 20 kB/s across three takers took {:?}", + started.elapsed() + ); + } + + #[test] + fn unused_allowance_goes_back_to_the_budget() { + let limit = RateLimit::default(); + limit.set(20_000); + let started = Instant::now(); + let grant = limit.take(4096); + limit.give_back(grant); + assert_eq!(limit.take(4096), grant); + assert!(started.elapsed() < TICK); + } + + #[test] + fn lifting_a_limit_frees_a_take_that_is_waiting() { + let limit = RateLimit::default(); + limit.set(100); + assert_eq!(limit.take(1000), 10); + std::thread::scope(|s| { + let waiting = s.spawn(|| limit.take(1000)); + std::thread::sleep(Duration::from_millis(10)); + limit.set(0); + assert_eq!(waiting.join().unwrap(), 1000); + }); + } +} diff --git a/soulseek-rs-lib/tests/e2e.rs b/soulseek-rs-lib/tests/e2e.rs index d859fa5..8e6cf70 100644 --- a/soulseek-rs-lib/tests/e2e.rs +++ b/soulseek-rs-lib/tests/e2e.rs @@ -936,6 +936,82 @@ fn a_download_completes_when_the_f_connection_beats_token_registration() { ); } +#[test] +fn a_download_speed_limit_paces_a_transfer_until_it_is_lifted() { + let server = server_or_skip!(); + let listen_port = free_port().expect("free listen port"); + let mut client = Client::with_settings(server.listening_settings( + "e2e_cappedpeer_dl", + "pw", + listen_port, + )); + client.connect().expect("connect"); + assert!(client.login().expect("login")); + client.set_download_speed_limit(64 * 1024); + + let filename = "capped.mp3"; + let content: Vec = + (0..1024 * 1024u32).map(|i| (i % 251) as u8).collect(); + let download_dir = unique_download_dir(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + let cfg = MockUpload { + listen_addr: format!("127.0.0.1:{listen_port}"), + peer_username: "e2e_cappedpeer".to_string(), + filename: filename.to_string(), + content: content.clone(), + token: 434_343, + ready: ready_tx, + token_delay: Duration::ZERO, + }; + let uploader = std::thread::spawn(move || { + if let Err(e) = run_mock_uploader(&cfg) { + eprintln!("[mock uploader] {e}"); + } + }); + ready_rx + .recv_timeout(Duration::from_secs(5)) + .expect("mock uploader P connection"); + + let started = Instant::now(); + let (_download, status_rx) = client + .download( + filename.to_string(), + "e2e_cappedpeer".to_string(), + content.len() as u64, + download_dir.display().to_string(), + ) + .expect("start download"); + let in_progress = || { + client + .get_all_downloads() + .iter() + .any(|d| matches!(d.status, DownloadStatus::InProgress { .. })) + }; + assert!(wait_for(in_progress), "the download should start"); + std::thread::sleep(Duration::from_secs(1)); + assert!( + in_progress(), + "1 MiB at 64 KiB/s must still be arriving a second in" + ); + + client.set_download_speed_limit(0); + assert!( + wait_for_completion(&client, &status_rx, Duration::from_secs(20)), + "the download should reach Completed" + ); + let elapsed = started.elapsed(); + let _ = uploader.join(); + assert!( + elapsed < Duration::from_secs(8), + "capped for the whole file this takes 16s; lifted it took {elapsed:?}" + ); + assert_eq!( + std::fs::read(download_dir.join(filename)).expect("downloaded file"), + content + ); + let _ = std::fs::remove_dir_all(&download_dir); +} + fn cancel_mid_transfer(peer_username: &str, stall_after: Option) { let server = server_or_skip!();