diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs index 3837642..7038889 100644 --- a/src/peer/default_peer.rs +++ b/src/peer/default_peer.rs @@ -2,10 +2,10 @@ use crate::dispatcher::MessageDispatcher; use crate::message::peer::{FileSearch, FileSearchResponse}; use crate::message::server::MessageFactory; use crate::message::{Handlers, Message, MessageReader}; + use std::sync::mpsc; use std::sync::mpsc::{Receiver, Sender}; -use std::sync::{Arc, Barrier}; -use std::thread::{self}; +use std::thread::{self, JoinHandle}; use crate::peer::Peer; use crate::client::ClientOperation; @@ -17,8 +17,10 @@ use std::time::Duration; #[allow(dead_code)] pub struct DefaultPeer { peer: Peer, - // client_channel: Sender, peer_channel: Option>, + // Store thread handles for proper lifecycle management + read_thread: Option>, + write_thread: Option>, } #[allow(dead_code)] @@ -30,8 +32,10 @@ impl DefaultPeer { pub fn new(peer: Peer, _client_channel: Sender) -> Self { Self { peer, - // client_channel, peer_channel: None, + // Initialize handles as None + read_thread: None, + write_thread: None, } } pub fn connect(mut self) -> Result { @@ -49,7 +53,7 @@ impl DefaultPeer { .write_all(&MessageFactory::build_watch_user(token.as_str()).get_data()) .unwrap(); } - self.start_read_write_loops(stream).unwrap(); + self.start_read_write_loops(stream)?; Ok(self) } @@ -58,30 +62,23 @@ impl DefaultPeer { let (peer_sender, peer_reader): (Sender, Receiver) = mpsc::channel(); - let sender = peer_sender.clone(); - let barrier = Arc::new(Barrier::new(3)); - let read_barrier = barrier.clone(); - let write_barrier = barrier.clone(); - let done_barrier = barrier.clone(); - let mut read_stream = stream.try_clone()?; - let mut write_stream = stream.try_clone()?; + let mut write_stream = stream; // Use the original stream for writing let peer = self.peer.clone(); - thread::spawn(move || { - read_barrier.wait(); + // Spawn the reader thread + let read_handle = thread::spawn(move || { let mut handlers = Handlers::new(); handlers.register_handler(FileSearchResponse); - let dispatcher = MessageDispatcher::new(sender, handlers); + let dispatcher = MessageDispatcher::new(peer_sender, handlers); let mut buffered_reader = MessageReader::new(); loop { match buffered_reader.read_from_socket(&mut read_stream) { Ok(_) => {} Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, - Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { println!( "Read operation timed out in default peer {:}:{:}", @@ -90,7 +87,7 @@ impl DefaultPeer { continue; } Err(e) => { - eprintln!("Error reading from peer: {}", e); + eprintln!("Error reading from peer: {}. Terminating read loop.", e); break; } } @@ -98,35 +95,38 @@ impl DefaultPeer { match buffered_reader.extract_message() { Ok(Some(mut message)) => dispatcher.dispatch(&mut message), Err(e) => { - println!("Error extracting message in default peer: {}", e) + println!("Error extracting message in default peer: {}. Terminating read loop.", e); + break; } Ok(None) => continue, } } }); - thread::spawn(move || { - write_barrier.wait(); + // Spawn the writer thread + let write_handle = thread::spawn(move || { loop { - if let Ok(operation) = peer_reader.recv() { - match operation { + match peer_reader.recv() { + Ok(operation) => match operation { PeerOperation::SendMessage(message) => { - match write_stream.write_all(&message.get_buffer()) { - Ok(_) => {} - Err(e) => { - eprintln!("Error writing message to stream : {}", e); - break; - } + if let Err(e) = write_stream.write_all(&message.get_buffer()) { + eprintln!("Error writing message to stream: {}. Terminating write loop.", e); + break; } } PeerOperation::FileSearchResult(file_search) => { println!("{:?}", file_search) } + }, + Err(_) => { + // The sender has been dropped, the peer is shutting down. + println!("Peer channel closed. Terminating write loop."); + break; } } } }); - done_barrier.wait(); + Ok(()) } }