diff --git a/soulseek-rs-lib/src/actor/peer_registry.rs b/soulseek-rs-lib/src/actor/peer_registry.rs index 7eec579..3819830 100644 --- a/soulseek-rs-lib/src/actor/peer_registry.rs +++ b/soulseek-rs-lib/src/actor/peer_registry.rs @@ -6,11 +6,12 @@ use crate::peer::Peer; use crate::utils::lock::MutexExt; use crate::{debug, error}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::net::TcpStream; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc::Sender; use std::sync::{Arc, Mutex}; +use std::time::Instant; /// Source of unique per-actor ids so terminal-outcome eviction can be made /// identity-aware (a replaced actor must not evict its replacement). @@ -18,13 +19,14 @@ static NEXT_PEER_ID: AtomicU64 = AtomicU64::new(1); /// Registered peers keyed by username, each stored with the unique id of the /// actor currently occupying the slot. -type PeerMap = HashMap)>; +type PeerMap = HashMap, Instant)>; pub struct PeerRegistry { peers: Arc>, actor_system: Arc, client_channel: Sender, own_username: String, + capacity: Arc, } impl PeerRegistry { @@ -33,12 +35,28 @@ impl PeerRegistry { actor_system: Arc, client_channel: Sender, own_username: String, + ) -> Self { + Self::with_capacity( + actor_system, + client_channel, + own_username, + Arc::new(AtomicUsize::new(usize::MAX)), + ) + } + + #[must_use] + pub(crate) fn with_capacity( + actor_system: Arc, + client_channel: Sender, + own_username: String, + capacity: Arc, ) -> Self { Self { peers: Arc::new(Mutex::new(HashMap::new())), actor_system, client_channel, own_username, + capacity, } } @@ -47,6 +65,16 @@ impl PeerRegistry { peer: Peer, stream: Option, reader: Option, + ) -> Result, String> { + self.register_peer_protected(peer, stream, reader, &HashSet::new()) + } + + pub(crate) fn register_peer_protected( + &self, + peer: Peer, + stream: Option, + reader: Option, + protected: &HashSet, ) -> Result, String> { let username = peer.username.clone(); let id = NEXT_PEER_ID.fetch_add(1, Ordering::Relaxed); @@ -71,6 +99,28 @@ impl PeerRegistry { .lock_safe() .map_err(|e| format!("peer registry lock poisoned: {e}"))?; + let capacity = self.capacity.load(Ordering::Relaxed).max(1); + if peers.len() >= capacity && !peers.contains_key(&username) { + let victim = peers + .iter() + .filter(|(name, _)| !protected.contains(name.as_str())) + .min_by_key(|(_, (_, _, registered))| *registered) + .map(|(name, _)| name.clone()); + let Some(victim) = victim else { + return Err(format!( + "peer registry at capacity ({capacity}) and every peer \ + is busy; refusing {username}" + )); + }; + if let Some((_, handle, _)) = peers.remove(&victim) { + let _ = handle.stop(); + debug!( + "[peer_registry] evicted idle peer {} to admit {}", + victim, username + ); + } + } + let handle = self .actor_system .try_spawn_with_handle(actor, |actor, handle| { @@ -81,8 +131,8 @@ impl PeerRegistry { // become an orphan pinning a pool worker forever. Eviction on the // replaced actor's later shutdown is identity-aware (keyed on its id), // so stopping it here cannot evict this new connection. - if let Some((_, old_handle)) = - peers.insert(username.clone(), (id, handle.clone())) + if let Some((_, old_handle, _)) = + peers.insert(username.clone(), (id, handle.clone(), Instant::now())) { let _ = old_handle.stop(); debug!( @@ -97,7 +147,9 @@ impl PeerRegistry { #[must_use] pub fn get_peer(&self, username: &str) -> Option> { match self.peers.lock_safe() { - Ok(peers) => peers.get(username).map(|(_, handle)| handle.clone()), + Ok(peers) => { + peers.get(username).map(|(_, handle, _)| handle.clone()) + } Err(e) => { error!("[peer_registry] get_peer: {}", e); None @@ -123,7 +175,7 @@ impl PeerRegistry { debug!("[peer_registry] Removed peer actor for {}", username); } - removed.map(|(_, handle)| handle) + removed.map(|(_, handle, _)| handle) } /// Remove and return the actor for `username` only if it is still the actor @@ -142,8 +194,11 @@ impl PeerRegistry { return None; } }; - if peers.get(username).is_some_and(|(stored, _)| *stored == id) { - let removed = peers.remove(username).map(|(_, handle)| handle); + if peers + .get(username) + .is_some_and(|(stored, _, _)| *stored == id) + { + let removed = peers.remove(username).map(|(_, handle, _)| handle); debug!( "[peer_registry] Removed peer actor {} for {}", id, username @@ -192,6 +247,7 @@ impl Clone for PeerRegistry { actor_system: self.actor_system.clone(), client_channel: self.client_channel.clone(), own_username: self.own_username.clone(), + capacity: self.capacity.clone(), } } } @@ -201,8 +257,112 @@ mod tests { use super::PeerRegistry; use crate::actor::ActorSystem; use crate::peer::{ConnectionType, Peer}; + use std::collections::HashSet; use std::net::{TcpListener, TcpStream}; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; + + fn loopback_peer(name: &str) -> (Peer, TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let stream = TcpStream::connect(addr).unwrap(); + stream.set_nonblocking(true).unwrap(); + let far_end = listener.accept().unwrap().0; + let peer = Peer::new( + name.to_string(), + ConnectionType::P, + "127.0.0.1".to_string(), + u32::from(addr.port()), + None, + 0, + 0, + 0, + ); + (peer, stream, far_end) + } + + fn capped_registry( + capacity: usize, + ) -> ( + PeerRegistry, + std::sync::mpsc::Receiver, + ) { + let system = Arc::new(ActorSystem::new()); + let (tx, rx) = std::sync::mpsc::channel(); + let registry = PeerRegistry::with_capacity( + system, + tx, + "me".to_string(), + Arc::new(AtomicUsize::new(capacity)), + ); + (registry, rx) + } + + #[test] + fn capacity_evicts_the_oldest_idle_peer() { + let (registry, _rx) = capped_registry(2); + + for name in ["a", "b", "c"] { + let (peer, stream, _far_end) = loopback_peer(name); + registry.register_peer(peer, Some(stream), None).unwrap(); + std::thread::sleep(Duration::from_millis(5)); + } + + assert!(!registry.contains("a"), "the oldest peer must be evicted"); + assert!(registry.contains("b")); + assert!(registry.contains("c")); + } + + #[test] + fn capacity_refuses_when_every_peer_is_protected() { + let (registry, _rx) = capped_registry(1); + + let (busy, busy_stream, _busy_far) = loopback_peer("busy"); + registry + .register_peer(busy, Some(busy_stream), None) + .unwrap(); + + let protected: HashSet = + std::iter::once("busy".to_string()).collect(); + let (newcomer, stream, _far) = loopback_peer("newcomer"); + let result = registry.register_peer_protected( + newcomer, + Some(stream), + None, + &protected, + ); + + assert!(result.is_err(), "a full registry of busy peers must refuse"); + assert!(registry.contains("busy")); + assert!(!registry.contains("newcomer")); + } + + #[test] + fn a_returning_username_is_replaced_not_refused_at_capacity() { + let (registry, _rx) = capped_registry(1); + + let (first, first_stream, _first_far) = loopback_peer("bob"); + registry + .register_peer(first, Some(first_stream), None) + .unwrap(); + + let protected: HashSet = + std::iter::once("bob".to_string()).collect(); + let (again, stream, _far) = loopback_peer("bob"); + let result = registry.register_peer_protected( + again, + Some(stream), + None, + &protected, + ); + + assert!( + result.is_ok(), + "a reconnecting username must replace itself" + ); + assert!(registry.contains("bob")); + } #[test] fn remove_peer_if_respects_actor_identity() { diff --git a/soulseek-rs-lib/src/client/connection.rs b/soulseek-rs-lib/src/client/connection.rs index da2bd69..bc4a6a4 100644 --- a/soulseek-rs-lib/src/client/connection.rs +++ b/soulseek-rs-lib/src/client/connection.rs @@ -50,10 +50,11 @@ impl Client { ) = mpsc::channel(); let mut ctx = self.context.write_safe()?; - let peer_registry = PeerRegistry::new( + let peer_registry = PeerRegistry::with_capacity( ctx.actor_system.clone(), sender.clone(), self.username.clone(), + ctx.max_peers.clone(), ); ctx.peer_registry = Some(peer_registry); @@ -175,6 +176,13 @@ impl Client { Ok(()) } + pub fn set_max_peers(&self, max_peers: usize) { + if let Ok(ctx) = self.context.read_safe() { + ctx.max_peers + .store(max_peers.max(1), std::sync::atomic::Ordering::Relaxed); + } + } + #[allow(dead_code)] pub fn remove_peer(&self, username: &str) { let context = match self.context.read_safe() { @@ -214,7 +222,10 @@ impl Client { } }; if let Some(ref registry) = context.peer_registry { - match registry.register_peer(peer_clone, stream, None) { + let protected = context.protected_peers(); + match registry.register_peer_protected( + peer_clone, stream, None, &protected, + ) { Ok(_) => (), Err(e) => { trace!( diff --git a/soulseek-rs-lib/src/client/mod.rs b/soulseek-rs-lib/src/client/mod.rs index f473408..b935ee9 100644 --- a/soulseek-rs-lib/src/client/mod.rs +++ b/soulseek-rs-lib/src/client/mod.rs @@ -23,7 +23,7 @@ use std::{ net::TcpStream, sync::{ Arc, RwLock, - atomic::{AtomicBool, AtomicU32, Ordering}, + atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}, mpsc::{self, Receiver, Sender}, }, thread::{self, sleep}, @@ -58,6 +58,10 @@ pub const DEFAULT_UPLOAD_SLOTS: usize = 10; /// before giving up and failing the download. Matches the direct-dial timeout. const BROKER_CONNECT_TIMEOUT: Duration = Duration::from_secs(20); +const DEFAULT_MAX_PEERS: usize = 512; + +const BROWSE_PROTECT_WINDOW: Duration = Duration::from_mins(1); + /// Source of non-zero correlation tokens for server-brokered connections. static NEXT_CONNECT_TOKEN: AtomicU32 = AtomicU32::new(1); @@ -295,7 +299,8 @@ pub struct ClientContext { private_messages: Vec, /// Correlation tokens for server-brokered (firewalled) connections, mapping /// a token we sent in a ConnectToPeer to the peer we expect back. - pending_connect_tokens: HashMap, + pending_connect_tokens: HashMap, + max_peers: Arc, /// Files we share with peers (read-only after connect). pub shares: Arc, /// The directories the current share index was built from. @@ -311,6 +316,7 @@ pub struct ClientContext { pending_serves: HashMap>, /// Shared-file listings received from peers we browsed. browse_results: HashMap>, + pending_browses: HashMap, /// Latest snapshot of the public chat-room list (from `RoomList`, code 64). room_list: Vec, /// Chat-room events awaiting consumption by the client/UI. @@ -395,6 +401,7 @@ impl ClientContext { searches: HashMap::new(), private_messages: Vec::new(), pending_connect_tokens: HashMap::new(), + max_peers: Arc::new(AtomicUsize::new(DEFAULT_MAX_PEERS)), shares: Arc::new(Shares::empty()), shared_directories: Vec::new(), peer_addresses: HashMap::new(), @@ -403,6 +410,7 @@ impl ClientContext { active_uploads: HashMap::new(), pending_serves: HashMap::new(), browse_results: HashMap::new(), + pending_browses: HashMap::new(), room_list: Vec::new(), room_events: Vec::new(), room_members: HashMap::new(), @@ -563,9 +571,17 @@ impl ClientContext { username: String, directories: Vec, ) { + self.pending_browses.remove(&username); self.browse_results.insert(username, directories); } + pub(crate) fn mark_browse_pending(&mut self, username: &str) { + let now = Instant::now(); + self.pending_browses.retain(|_, deadline| *deadline > now); + self.pending_browses + .insert(username.to_string(), now + BROWSE_PROTECT_WINDOW); + } + /// Remove and return the shared-file listing browsed from `username`. pub fn take_browse_result( &mut self, @@ -577,12 +593,57 @@ impl ClientContext { /// Remember that a server-brokered connection to `username` is pending under /// `token`; the peer will quote it back in a PierceFirewall. pub fn add_pending_connect(&mut self, token: u32, username: String) { - self.pending_connect_tokens.insert(token, username); + self.pending_connect_tokens + .insert(token, (username, Instant::now() + BROKER_CONNECT_TIMEOUT)); } /// Resolve and consume the peer expected for a brokered connection `token`. pub fn take_pending_connect(&mut self, token: u32) -> Option { - self.pending_connect_tokens.remove(&token) + self.pending_connect_tokens + .remove(&token) + .map(|(username, _)| username) + } + + pub fn take_expired_connects(&mut self, now: Instant) -> Vec { + let mut expired = Vec::new(); + self.pending_connect_tokens + .retain(|_, (username, deadline)| { + if now >= *deadline { + expired.push(std::mem::take(username)); + return false; + } + true + }); + expired + } + + pub(crate) fn protected_peers(&self) -> HashSet { + let mut protected: HashSet = self + .downloads + .list() + .iter() + .filter(|d| { + matches!( + d.status, + DownloadStatus::Queued + | DownloadStatus::InProgress { .. } + | DownloadStatus::Paused { .. } + ) + }) + .map(|d| d.username.clone()) + .collect(); + protected + .extend(self.active_uploads.values().map(|u| u.username.clone())); + protected.extend(self.pending_serves.keys().cloned()); + protected.extend(self.pending_peer_messages.keys().cloned()); + let now = Instant::now(); + protected.extend( + self.pending_browses + .iter() + .filter(|(_, deadline)| **deadline > now) + .map(|(username, _)| username.clone()), + ); + protected } /// Record a private message received from another user. diff --git a/soulseek-rs-lib/src/client/operations.rs b/soulseek-rs-lib/src/client/operations.rs index 7bdc715..02b67cd 100644 --- a/soulseek-rs-lib/src/client/operations.rs +++ b/soulseek-rs-lib/src/client/operations.rs @@ -1,12 +1,14 @@ use super::{ - Arc, BROKER_CONNECT_TIMEOUT, Client, ClientContext, ClientOperation, - ConnectionType, Download, DownloadPeer, DownloadStatus, Peer, PeerMessage, + Arc, Client, ClientContext, ClientOperation, ConnectionType, Download, + DownloadPeer, DownloadStatus, Duration, Instant, Peer, PeerMessage, PeerRegistry, Receiver, RwLock, RwLockExt, ServerMessage, - build_search_response, debug, error, info, next_connect_token, sleep, + build_search_response, debug, error, info, mpsc, next_connect_token, thread, trace, warn, }; use crate::message::server::MessageFactory; +const CONNECT_SWEEP_INTERVAL: Duration = Duration::from_secs(1); + impl Client { pub(crate) fn listen_to_client_operations( reader: Receiver, @@ -14,20 +16,40 @@ impl Client { own_username: String, ) { thread::spawn(move || { - for operation in reader { + let mut last_sweep = Instant::now(); + loop { + let next = reader.recv_timeout(CONNECT_SWEEP_INTERVAL); + if last_sweep.elapsed() >= CONNECT_SWEEP_INTERVAL { + last_sweep = Instant::now(); + Self::sweep_expired_connects(&client_context); + } + let operation = match next { + Ok(operation) => operation, + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => break, + }; match operation { ClientOperation::ConnectToPeer(peer) => { - let client_context_clone = client_context.clone(); - let own_username_clone = own_username.clone(); + if matches!(peer.connection_type, ConnectionType::F) { + let client_context_clone = client_context.clone(); + let own_username_clone = own_username.clone(); - thread::spawn(move || { + thread::spawn(move || { + Self::connect_to_peer( + peer, + client_context_clone, + own_username_clone, + None, + ); + }); + } else { Self::connect_to_peer( peer, - client_context_clone, - own_username_clone, + client_context.clone(), + own_username.clone(), None, ); - }); + } } ClientOperation::SearchResult(search_result) => { trace!("[client] SearchResult {:?}", search_result); @@ -122,6 +144,28 @@ impl Client { ); continue; }; + if matches!( + download.status, + DownloadStatus::InProgress { .. } + | DownloadStatus::Completed + ) { + debug!( + "[client] transfer token {} already claimed; \ + ignoring replayed TransferResponse", + token + ); + continue; + } + if let Ok(mut ctx) = client_context.write_safe() { + ctx.update_download_with_status( + token, + DownloadStatus::InProgress { + bytes_downloaded: 0, + total_bytes: download.size, + speed_bytes_per_sec: 0.0, + }, + ); + } thread::spawn(move || { let download_peer = DownloadPeer::new( @@ -285,17 +329,12 @@ impl Client { u8::try_from(obfuscation_type).unwrap_or(0), obfuscated_port, ); - let client_context_clone = client_context.clone(); - let own_username_clone = own_username.clone(); - - thread::spawn(move || { - Self::connect_to_peer( - peer, - client_context_clone, - own_username_clone, - None, - ); - }); + Self::connect_to_peer( + peer, + client_context.clone(), + own_username.clone(), + None, + ); } } ClientOperation::UpdateDownloadTokens( @@ -751,30 +790,19 @@ impl Client { ConnectionType::P, ); let _ = sender.send(ServerMessage::SendMessage(msg)); - - // Bound the brokered attempt: if no PierceFirewall - // consumes the token, fail the peer's queued - // downloads (so the caller's Receiver unblocks) - // and reclaim the token. A successful pierce - // takes the token first, making this a no-op. - let timeout_ctx = client_context.clone(); - let timeout_user = username.clone(); - thread::spawn(move || { - sleep(BROKER_CONNECT_TIMEOUT); - let still_pending = - timeout_ctx.write_safe().is_ok_and(|mut c| { - c.take_pending_connect(token).is_some() - }); - if still_pending { - Self::fail_queued_downloads( - &timeout_ctx, - &timeout_user, - ); - } - }); } } } }); } + + fn sweep_expired_connects(client_context: &Arc>) { + let expired = match client_context.write_safe() { + Ok(mut ctx) => ctx.take_expired_connects(Instant::now()), + Err(_) => return, + }; + for username in expired { + Self::fail_queued_downloads(client_context, &username); + } + } } diff --git a/soulseek-rs-lib/src/client/rooms.rs b/soulseek-rs-lib/src/client/rooms.rs index 41d4d5b..8afc871 100644 --- a/soulseek-rs-lib/src/client/rooms.rs +++ b/soulseek-rs-lib/src/client/rooms.rs @@ -177,7 +177,8 @@ impl Client { let request = crate::message::server::MessageFactory::build_get_share_file_list(); let (connected, registry) = { - let ctx = self.context.read_safe()?; + let mut ctx = self.context.write_safe()?; + ctx.mark_browse_pending(username); ( ctx.peer_registry .as_ref() diff --git a/soulseek-rs-lib/src/client/tests.rs b/soulseek-rs-lib/src/client/tests.rs index 9179497..9b5efcb 100644 --- a/soulseek-rs-lib/src/client/tests.rs +++ b/soulseek-rs-lib/src/client/tests.rs @@ -395,3 +395,51 @@ fn a_clean_disconnect_keeps_queued_uploads_an_error_drops_them() { thread::sleep(Duration::from_millis(10)); } } + +#[test] +fn a_pending_connect_expires_after_the_broker_timeout() { + let mut context = ClientContext::new(); + context.add_pending_connect(7, "ghost".to_string()); + + assert!(context.take_expired_connects(Instant::now()).is_empty()); + + let after_deadline = + Instant::now() + BROKER_CONNECT_TIMEOUT + Duration::from_secs(1); + assert_eq!( + context.take_expired_connects(after_deadline), + vec!["ghost".to_string()] + ); + assert!(context.take_pending_connect(7).is_none()); +} + +#[test] +fn protected_peers_covers_downloads_uploads_and_pending_serves() { + let mut context = ClientContext::new(); + context.add_download(download( + "downloader", + "song.mp3", + 1, + DownloadStatus::Queued, + mpsc::channel().0, + )); + context.add_download(download( + "done", + "old.mp3", + 2, + DownloadStatus::Completed, + mpsc::channel().0, + )); + context + .pending_serves + .insert("waiting".to_string(), vec![9]); + context.mark_browse_pending("browsed"); + + let protected = context.protected_peers(); + assert!(protected.contains("downloader")); + assert!(protected.contains("waiting")); + assert!(protected.contains("browsed")); + assert!(!protected.contains("done")); + + context.store_browse_result("browsed".to_string(), Vec::new()); + assert!(!context.protected_peers().contains("browsed")); +} diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index ed77a24..739e3a0 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/soulseek-rs-lib/src/peer/listen.rs @@ -11,6 +11,7 @@ use crate::message::{Message, MessageReader}; use crate::peer::{ConnectionType, DownloadPeer, Peer}; use crate::types::Download; use crate::utils::lock::RwLockExt; +use crate::utils::semaphore::{Permit, Semaphore}; use crate::{DownloadStatus, debug, error, info, trace}; /// How long to wait before accepting again after a failure. @@ -25,6 +26,8 @@ const PEER_INIT_MESSAGE_CODE: u8 = 1; /// How long an accepted peer gets to send its peer-init handshake. const PEER_INIT_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_HANDSHAKES: usize = 128; + #[derive(Clone)] struct ConnectionContext { client_sender: Sender, @@ -140,7 +143,13 @@ fn handle_peer_connection( } }; if let Some(ref registry) = client_context.peer_registry { - match registry.register_peer(peer.clone(), Some(stream), Some(reader)) { + let protected = client_context.protected_peers(); + match registry.register_peer_protected( + peer.clone(), + Some(stream), + Some(reader), + &protected, + ) { Ok(_) => (), Err(e) => { error!( @@ -283,6 +292,7 @@ fn handle_pierce_firewall( fn handle_incoming_connection( mut stream: TcpStream, context: ConnectionContext, + permit: Permit, ) { let Ok(peer_addr) = stream.peer_addr() else { error!("[listener] failed to get peer address"); @@ -303,6 +313,7 @@ fn handle_incoming_connection( return; } }; + drop(permit); // A firewalled peer brokered through the server connects back with a // PierceFirewall (code 0) instead of a PeerInit (code 1). @@ -406,6 +417,7 @@ impl Listen { client_context, own_username, }; + let handshakes = Arc::new(Semaphore::new(MAX_HANDSHAKES)); for stream in listener.incoming() { let Ok(stream) = stream else { @@ -421,11 +433,14 @@ impl Listen { continue; }; + let permit = handshakes.acquire(); let context = context.clone(); // One thread per connection: the peer-init handshake blocks, and a // peer that is slow to send one must not stop us accepting anybody // else — a wedged accept loop makes us unreachable to every peer. - thread::spawn(move || handle_incoming_connection(stream, context)); + thread::spawn(move || { + handle_incoming_connection(stream, context, permit); + }); } } } diff --git a/soulseek-rs-lib/src/utils/mod.rs b/soulseek-rs-lib/src/utils/mod.rs index 6b42e19..5bb6848 100644 --- a/soulseek-rs-lib/src/utils/mod.rs +++ b/soulseek-rs-lib/src/utils/mod.rs @@ -3,6 +3,7 @@ pub mod logger; pub mod lock; pub mod md5; pub mod path; +pub(crate) mod semaphore; pub mod zlib; // Re-export commonly used items diff --git a/soulseek-rs-lib/src/utils/semaphore.rs b/soulseek-rs-lib/src/utils/semaphore.rs new file mode 100644 index 0000000..0508551 --- /dev/null +++ b/soulseek-rs-lib/src/utils/semaphore.rs @@ -0,0 +1,73 @@ +use std::sync::{Arc, Condvar, Mutex, PoisonError}; + +pub struct Semaphore { + permits: Mutex, + available: Condvar, +} + +pub struct Permit(Arc); + +impl Semaphore { + #[must_use] + pub const fn new(permits: usize) -> Self { + Self { + permits: Mutex::new(permits), + available: Condvar::new(), + } + } + + #[must_use] + pub fn acquire(self: &Arc) -> Permit { + let mut permits = + self.permits.lock().unwrap_or_else(PoisonError::into_inner); + while *permits == 0 { + permits = self + .available + .wait(permits) + .unwrap_or_else(PoisonError::into_inner); + } + *permits -= 1; + Permit(Arc::clone(self)) + } +} + +impl Drop for Permit { + fn drop(&mut self) { + let mut permits = self + .0 + .permits + .lock() + .unwrap_or_else(PoisonError::into_inner); + *permits += 1; + self.0.available.notify_one(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn permits_bound_concurrent_holders() { + let semaphore = Arc::new(Semaphore::new(2)); + let first = semaphore.acquire(); + let _second = semaphore.acquire(); + + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let _third = semaphore.acquire(); + tx.send(()).unwrap(); + }); + + assert!( + rx.recv_timeout(Duration::from_millis(100)).is_err(), + "a third acquire must block while both permits are held" + ); + drop(first); + assert!( + rx.recv_timeout(Duration::from_secs(2)).is_ok(), + "releasing a permit must wake the waiter" + ); + } +}