diff --git a/src/client.rs b/src/client.rs index ad6088a..fcce2e3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,7 +1,7 @@ use crate::{ peer::{DefaultPeer, Peer}, server::{PeerAddress, Server}, - utils::md5, + utils::{md5, thread_pool::ThreadPool}, }; use std::{ collections::HashMap, @@ -15,18 +15,22 @@ use std::{ sync::{mpsc, Arc}, time::{Duration, Instant}, }; + +const MAX_THREADS: usize = 100; pub enum ClientOperation { ConnectToPeer(Peer), } struct ClientContext { peers: HashMap, sender: Option>, + thread_pool: ThreadPool, } impl ClientContext { pub fn new() -> Self { Self { peers: HashMap::new(), sender: None, + thread_pool: ThreadPool::new(MAX_THREADS), } } } @@ -139,7 +143,7 @@ impl Client { if !unlocked_context.peers.contains_key(&peer.username) { let peer_clone = peer.clone(); let sender_clone = sender.clone(); - thread::spawn(move || { + unlocked_context.thread_pool.execute(move || { let default_peer = DefaultPeer::new(peer_clone, sender_clone); match default_peer.connect() { Ok(p) => { diff --git a/src/peer/mod.rs b/src/peer/mod.rs index 05c7ca4..8b43715 100644 --- a/src/peer/mod.rs +++ b/src/peer/mod.rs @@ -53,7 +53,7 @@ pub struct Peer { pub obfuscated_port: i8, } impl Peer { - #[allow(clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments, dead_code)] pub fn new( username: String, connection_type: ConnectionType, @@ -154,4 +154,5 @@ fn test_new_from_message2() { assert_eq!(peer.privileged, 0); assert_eq!(peer.unknown, 1); assert_eq!(peer.obfuscated_port, 0); -} \ No newline at end of file +} + diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 4b47c1d..a9f60ee 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,4 +1,5 @@ pub mod md5; +pub mod thread_pool; pub mod zlib; pub use md5::*; diff --git a/src/utils/thread_pool.rs b/src/utils/thread_pool.rs new file mode 100644 index 0000000..d453994 --- /dev/null +++ b/src/utils/thread_pool.rs @@ -0,0 +1,69 @@ +use std::sync::{mpsc, Arc, Mutex}; +use std::thread; + +type Job = Box; + +pub struct ThreadPool { + workers: Vec, + sender: mpsc::Sender, +} + +impl ThreadPool { + pub fn new(size: usize) -> ThreadPool { + assert!(size > 0); + + let (sender, receiver) = mpsc::channel(); + let receiver = Arc::new(Mutex::new(receiver)); + let mut workers = Vec::with_capacity(size); + + for id in 0..size { + workers.push(Worker::new(id, Arc::clone(&receiver))); + } + + ThreadPool { workers, sender } + } + + pub fn execute(&self, f: F) + where + F: FnOnce() + Send + 'static, + { + let job = Box::new(f); + self.sender.send(job).unwrap(); + } +} + +impl Drop for ThreadPool { + fn drop(&mut self) { + for worker in &mut self.workers { + if let Some(thread) = worker.thread.take() { + thread.join().unwrap(); + } + } + } +} + +struct Worker { + _id: usize, + thread: Option>, +} + +impl Worker { + fn new(id: usize, receiver: Arc>>) -> Worker { + let thread = thread::spawn(move || loop { + let job = receiver.lock().unwrap().recv(); + match job { + Ok(job) => { + job(); + } + Err(_) => { + break; + } + } + }); + + Worker { + _id: id, + thread: Some(thread), + } + } +} \ No newline at end of file