more work

Author: da2ce7

examples/get_metadata/Cargo.toml

@@ -24,3 +24,5 @@ hex = "0"
 futures = "0"
 tokio = { version = "1", features = ["full"] }
 tokio-util = {version = "0", features = ["codec"]}
+tracing = "0"
+tracing-subscriber = "0"

examples/get_metadata/src/main.rs

@@ -118,22 +118,22 @@ async fn main() {
     extensions.add(Extension::ExtensionProtocol);
 
     // Create a handshaker that can initiate connections with peers
-    let (handshaker_send, mut handshaker_recv) = HandshakerBuilder::new()
+    let (handshaker, mut tasks) = HandshakerBuilder::new()
         .with_extensions(extensions)
         .with_config(
             // Set a low handshake timeout so we don't wait on peers that aren't listening on tcp
             HandshakerConfig::default().with_connect_timeout(Duration::from_millis(500)),
         )
         .build(TcpTransport)
         .await
-        .expect("it should build a handshaker pair")
-        .into_parts();
+        .expect("it should build a handshaker pair");
+    let (handshaker_send, mut handshaker_recv) = handshaker.into_parts();
 
     // Create a peer manager that will hold our peers and heartbeat/send messages to them
     let (mut peer_manager_send, peer_manager_recv) = PeerManagerBuilder::new().build().into_parts();
 
     // Hook up a future that feeds incoming (handshaken) peers over to the peer manager
-    tokio::spawn(async move {
+    tasks.spawn(async move {
         while let Some(complete_msg) = handshaker_recv.next().await {
             let (_, extensions, hash, pid, addr, sock) = complete_msg.unwrap().into_parts();
 
@@ -185,7 +185,7 @@ async fn main() {
     let mut merged_recv = futures::stream::select(peer_manager_recv.map(Either::Left), timer.map(Either::Right)).boxed();
 
     // Hook up a future that receives messages from the peer manager
-    tokio::spawn(async move {
+    tasks.spawn(async move {
         let mut uber_send = uber_send.clone();
 
         while let Some(item) = merged_recv.next().await {
@@ -268,4 +268,6 @@ async fn main() {
         .expect("Failed to create output file")
         .write_all(&metainfo.to_bytes())
         .expect("Failed to write metainfo to file");
+
+    tasks.shutdown().await;
 }

examples/simple_torrent/Cargo.toml

@@ -24,3 +24,5 @@ clap = "4"
 futures = "0"
 tokio = { version = "1", features = ["full"] }
 tokio-util = {version = "0", features = ["codec"]}
+tracing = "0"
+tracing-subscriber = "0"

examples/simple_torrent/src/main.rs

@@ -2,7 +2,7 @@ use std::cmp;
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::Read;
-use std::sync::Arc;
+use std::sync::{Arc, Once};
 
 use disk::fs::NativeFileSystem;
 use disk::fs_cache::FileHandleCache;
@@ -16,8 +16,7 @@ use futures::lock::Mutex;
 use futures::{stream, SinkExt as _, StreamExt as _};
 use handshake::transports::TcpTransport;
 use handshake::{
-    Extensions, Handshaker, HandshakerBuilder, HandshakerConfig, HandshakerSink, HandshakerStream, InitiateMessage, PeerId,
-    Protocol,
+    Extensions, Handshaker, HandshakerBuilder, HandshakerConfig, HandshakerStream, InitiateMessage, PeerId, Protocol,
 };
 use metainfo::{Info, Metainfo};
 use peer::messages::{BitFieldMessage, HaveMessage, PeerWireProtocolMessage, PieceMessage, RequestMessage};
@@ -31,10 +30,13 @@ use tokio::signal;
 use tokio::task::JoinSet;
 use tokio_util::bytes::BytesMut;
 use tokio_util::codec::{Decoder, Framed};
+use tracing::level_filters::LevelFilter;
 
 // Maximum number of requests that can be in flight at once.
 const MAX_PENDING_BLOCKS: usize = 50;
 
+pub static INIT: Once = Once::new();
+
 // Enum to store our selection state updates
 #[allow(dead_code)]
 #[derive(Debug)]
@@ -59,8 +61,30 @@ enum Downloader {
     Interrupted,
 }
 
+enum Setup {
+    Finished((NativeDiskManager, PeerManager, TcpHandshaker), JoinSet<()>),
+    Interrupted,
+}
+
+pub fn tracing_stdout_init(filter: LevelFilter) {
+    let builder = tracing_subscriber::fmt().with_max_level(filter).with_ansi(true);
+
+    builder.pretty().with_file(true).init();
+
+    tracing::info!("Logging initialized");
+}
+
+async fn ctrl_c() {
+    signal::ctrl_c().await.expect("failed to listen for event");
+    println!("Ctrl-C received, shutting down...");
+}
+
 #[tokio::main]
 async fn main() {
+    INIT.call_once(|| {
+        tracing_stdout_init(LevelFilter::TRACE);
+    });
+
     // Parse command-line arguments
     let matched_arguments = parse_arguments();
     let (torrent_file_path, download_directory, peer_address) = extract_arguments(&matched_arguments);
@@ -71,28 +95,51 @@ async fn main() {
     // Create a JoinSet to manage background tasks
     let tasks = Arc::new(Mutex::new(JoinSet::new()));
 
-    let ctrl_c = async {
-        signal::ctrl_c().await.expect("failed to listen for event");
-        println!("Ctrl-C received, shutting down...");
+    // Setup the managers.
+    let setup = setup(download_directory);
+
+    // Await either the completion of the setup or the Ctrl-C signal
+    let setup = tokio::select! {
+        setup = setup => Setup::Finished(setup.0, setup.1),
+        () = ctrl_c() => Setup::Interrupted,
+    };
+
+    let (managers, mut handshaker_tasks) = match setup {
+        Setup::Finished(managers, handshaker_tasks) => (managers, handshaker_tasks),
+        Setup::Interrupted => {
+            tracing::warn!("setup was canceled...");
+            return;
+        }
     };
 
-    let downloader = downloader(tasks.clone(), download_directory, peer_address, metainfo, info_hash);
+    let downloader = downloader(tasks.clone(), managers, peer_address, metainfo, info_hash);
 
-    // Await either the completion of all tasks or the Ctrl-C signal
+    // Await either the completion of the downloader or the Ctrl-C signal
     let status = tokio::select! {
         () = downloader => Downloader::Finished,
-        () = ctrl_c => Downloader::Interrupted,
+        () = ctrl_c() => Downloader::Interrupted,
     };
 
     match status {
         Downloader::Finished => {
-            while let Some(result) = tasks.lock().await.join_next().await {
+            while let Some(result) = handshaker_tasks.try_join_next() {
                 if let Err(e) = result {
                     eprintln!("Task failed: {e:?}");
                 }
             }
+            handshaker_tasks.shutdown().await;
+
+            while let Some(result) = tasks.lock().await.try_join_next() {
+                if let Err(e) = result {
+                    eprintln!("Task failed: {e:?}");
+                }
+            }
+            tasks.lock().await.shutdown().await;
+        }
+        Downloader::Interrupted => {
+            handshaker_tasks.shutdown().await;
+            tasks.lock().await.shutdown().await;
         }
-        Downloader::Interrupted => tasks.lock().await.shutdown().await,
     }
 }
 
@@ -142,21 +189,36 @@ fn load_and_parse_torrent_file(torrent_file_path: &str) -> (Metainfo, InfoHash)
     (metainfo, info_hash)
 }
 
+async fn setup(download_directory: String) -> ((NativeDiskManager, PeerManager, TcpHandshaker), JoinSet<()>) {
+    // Setup disk manager for handling file operations
+    let disk_manager = setup_disk_manager(&download_directory);
+
+    // Setup peer manager for managing peer communication
+    let peer_manager = setup_peer_manager();
+
+    // Setup handshaker for managing peer connections
+    let (handshaker, handshaker_tasks) = setup_handshaker().await;
+
+    ((disk_manager, peer_manager, handshaker), handshaker_tasks)
+}
+
 async fn downloader(
     tasks: Arc<Mutex<JoinSet<()>>>,
-    download_directory: String,
+    managers: (NativeDiskManager, PeerManager, TcpHandshaker),
     peer_address: String,
     metainfo: Metainfo,
     info_hash: InfoHash,
 ) {
+    let (disk_manager, peer_manager, handshaker) = managers;
+
     // Setup disk manager for handling file operations
-    let (mut disk_manager_sender, disk_manager_receiver) = setup_disk_manager(&download_directory);
+    let (mut disk_manager_sender, disk_manager_receiver) = disk_manager.into_parts();
 
     // Setup peer manager for managing peer communication
-    let (peer_manager_sender, peer_manager_receiver) = setup_peer_manager();
+    let (peer_manager_sender, peer_manager_receiver) = peer_manager.into_parts();
 
     // Setup handshaker for managing peer connections
-    let (mut handshaker_sender, handshaker_receiver) = setup_handshaker().await;
+    let (mut handshaker_sender, handshaker_receiver) = handshaker.into_parts();
 
     // Handle new incoming connections
     tasks
@@ -231,44 +293,39 @@ async fn downloader(
     .await;
 }
 
-fn setup_disk_manager(download_directory: &str) -> (DiskManagerSink<FileHandleCache<NativeFileSystem>>, DiskManagerStream) {
+type NativeDiskManager = DiskManager<FileHandleCache<NativeFileSystem>>;
+
+fn setup_disk_manager(download_directory: &str) -> DiskManager<FileHandleCache<NativeFileSystem>> {
     let filesystem = FileHandleCache::new(NativeFileSystem::with_directory(download_directory), 100);
-    let disk_manager: DiskManager<FileHandleCache<NativeFileSystem>> = DiskManagerBuilder::new()
+
+    DiskManagerBuilder::new()
         .with_sink_buffer_capacity(1)
         .with_stream_buffer_capacity(0)
-        .build(Arc::new(filesystem));
-
-    disk_manager.into_parts()
+        .build(Arc::new(filesystem))
 }
 
-async fn setup_handshaker() -> (HandshakerSink, HandshakerStream<TcpStream>) {
-    let handshaker: Handshaker<TcpStream> = HandshakerBuilder::new()
+type TcpHandshaker = Handshaker<TcpStream>;
+
+async fn setup_handshaker() -> (TcpHandshaker, JoinSet<()>) {
+    HandshakerBuilder::new()
         .with_peer_id(PeerId::from_hash("-BI0000-000000000000".as_bytes()).unwrap())
         .with_config(HandshakerConfig::default().with_wait_buffer_size(0).with_done_buffer_size(0))
         .build(TcpTransport)
         .await
-        .unwrap();
-
-    handshaker.into_parts()
+        .unwrap()
 }
 
+type PeerManager = peer::PeerManager<
+    Framed<TcpStream, PeerProtocolCodec<PeerWireProtocol<NullProtocol>>>,
+    PeerWireProtocolMessage<NullProtocol>,
+>;
+
 #[allow(clippy::type_complexity)]
-fn setup_peer_manager() -> (
-    PeerManagerSink<Framed<TcpStream, PeerProtocolCodec<PeerWireProtocol<NullProtocol>>>, PeerWireProtocolMessage<NullProtocol>>,
-    PeerManagerStream<
-        Framed<TcpStream, PeerProtocolCodec<PeerWireProtocol<NullProtocol>>>,
-        PeerWireProtocolMessage<NullProtocol>,
-    >,
-) {
-    let peer_manager: peer::PeerManager<
-        Framed<TcpStream, PeerProtocolCodec<PeerWireProtocol<NullProtocol>>>,
-        PeerWireProtocolMessage<NullProtocol>,
-    > = PeerManagerBuilder::new()
+fn setup_peer_manager() -> PeerManager {
+    PeerManagerBuilder::new()
         .with_sink_buffer_capacity(0)
         .with_stream_buffer_capacity(0)
-        .build();
-
-    peer_manager.into_parts()
+        .build()
 }
 
 async fn handle_new_connections(

packages/handshake/examples/handshake_torrent.rs

@@ -25,7 +25,7 @@ async fn main() -> std::io::Result<()> {
 
     // Show up as a uTorrent client...
     let peer_id = (*b"-UT2060-000000000000").into();
-    let mut handshaker = HandshakerBuilder::new()
+    let (mut handshaker, mut tasks) = HandshakerBuilder::new()
         .with_peer_id(peer_id)
         .build(TcpTransport)
         .await
@@ -40,6 +40,8 @@ async fn main() -> std::io::Result<()> {
 
     sleep(Duration::from_secs(10)).await;
 
+    tasks.shutdown().await;
+
     Ok(())
 }
 

packages/handshake/src/handshake/handshaker.rs

@@ -114,7 +114,7 @@ impl HandshakerBuilder {
     /// # Errors
     ///
     /// Returns a IO error if unable to build.
-    pub async fn build<T>(&self, transport: T) -> std::io::Result<Handshaker<T::Socket>>
+    pub async fn build<T>(&self, transport: T) -> std::io::Result<(Handshaker<T::Socket>, JoinSet<()>)>
     where
         T: Transport + Send + Sync + 'static,
         <T as Transport>::Socket: AsyncWrite + AsyncRead + std::fmt::Debug + Send + Sync,
@@ -157,7 +157,7 @@ impl<S> Handshaker<S>
 where
     S: AsyncRead + AsyncWrite + std::fmt::Debug + Send + Sync + Unpin + 'static,
 {
-    async fn with_builder<T>(builder: &HandshakerBuilder, transport: T) -> std::io::Result<Handshaker<T::Socket>>
+    async fn with_builder<T>(builder: &HandshakerBuilder, transport: T) -> std::io::Result<(Handshaker<T::Socket>, JoinSet<()>)>
     where
         T: Transport<Socket = S> + Send + Sync + 'static,
         <T as Transport>::Listener: Send,
@@ -208,7 +208,7 @@ where
         let sink = HandshakerSink::new(addr_send, open_port, builder.pid, filters);
         let stream = HandshakerStream::new(sock_recv);
 
-        Ok(Handshaker { sink, stream })
+        Ok((Handshaker { sink, stream }, tasks))
     }
 }
 

packages/handshake/src/message/complete.rs

@@ -7,6 +7,7 @@ use crate::message::protocol::Protocol;
 
 /// Message containing completed handshaking information.
 #[allow(clippy::module_name_repetitions)]
+#[derive(Debug)]
 pub struct CompleteMessage<S> {
     prot: Protocol,
     ext: Extensions,

packages/handshake/tests/test_byte_after_handshake.rs

@@ -1,6 +1,5 @@
 use std::io::{Read, Write};
 use std::net::TcpStream;
-use std::thread;
 
 use common::{tracing_stdout_init, INIT};
 use futures::stream::StreamExt;
@@ -21,7 +20,7 @@ async fn positive_recover_bytes() {
     let mut handshaker_one_addr = "127.0.0.1:0".parse().unwrap();
     let handshaker_one_pid = [4u8; bt::PEER_ID_LEN].into();
 
-    let mut handshaker_one = HandshakerBuilder::new()
+    let (mut handshaker_one, mut tasks_one) = HandshakerBuilder::new()
         .with_bind_addr(handshaker_one_addr)
         .with_peer_id(handshaker_one_pid)
         .build(TcpTransport)
@@ -30,7 +29,7 @@ async fn positive_recover_bytes() {
 
     handshaker_one_addr.set_port(handshaker_one.port());
 
-    thread::spawn(move || {
+    tasks_one.spawn_blocking(move || {
         let mut stream = TcpStream::connect(handshaker_one_addr).unwrap();
         let mut write_buffer = Vec::new();
 
@@ -46,15 +45,21 @@ async fn positive_recover_bytes() {
         stream.read_exact(&mut vec![0u8; expect_read_length][..]).unwrap();
     });
 
-    if let Some(message) = handshaker_one.next().await {
-        let (_, _, _, _, _, mut sock) = message.unwrap().into_parts();
+    let test = tokio::spawn(async move {
+        if let Some(message) = handshaker_one.next().await {
+            let (_, _, _, _, _, mut sock) = message.unwrap().into_parts();
 
-        let mut recv_buffer = vec![0u8; 1];
-        sock.read_exact(&mut recv_buffer).await.unwrap();
+            let mut recv_buffer = vec![0u8; 1];
+            sock.read_exact(&mut recv_buffer).await.unwrap();
 
-        // Assert that our buffer contains the bytes after the handshake
-        assert_eq!(vec![55], recv_buffer);
-    } else {
-        panic!("Failed to receive handshake message");
-    }
+            // Assert that our buffer contains the bytes after the handshake
+            assert_eq!(vec![55], recv_buffer);
+        } else {
+            panic!("Failed to receive handshake message");
+        }
+    });
+
+    let res = test.await;
+    tasks_one.shutdown().await;
+    res.unwrap();
 }