chore: improve read/write handling for default peer

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<ClientOperation>,
     peer_channel: Option<Sender<PeerOperation>>,
+    // Store thread handles for proper lifecycle management
+    read_thread: Option<JoinHandle<()>>,
+    write_thread: Option<JoinHandle<()>>,
 }
 
 #[allow(dead_code)]
@@ -30,8 +32,10 @@ impl DefaultPeer {
     pub fn new(peer: Peer, _client_channel: Sender<ClientOperation>) -> Self {
         Self {
             peer,
-            // client_channel,
             peer_channel: None,
+            // Initialize handles as None
+            read_thread: None,
+            write_thread: None,
         }
     }
     pub fn connect(mut self) -> Result<Self, io::Error> {
@@ -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<PeerOperation>, Receiver<PeerOperation>) =
             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(())
     }
 }