From f58df6c244826a43897efe9471b165ddf8d7480a Mon Sep 17 00:00:00 2001 From: Lee Smet Date: Mon, 10 Feb 2025 15:22:33 +0100 Subject: [PATCH] Close #529: Make peer connection flush cancellable Investigation has shown that connection file descriptors could leak in situations where both peers disconnected with data left on the socket. While it's not entirely clear how this happened, it seems that the select was blocked while flushing. To alleviate this, make the flush itself a branch on the select, and therefore cancellable. Signed-off-by: Lee Smet --- mycelium/src/peer.rs | 96 ++++++++++++++++++++++++++------------------ 1 file changed, 58 insertions(+), 38 deletions(-) diff --git a/mycelium/src/peer.rs b/mycelium/src/peer.rs index 7eb1c99..5b769b4 100644 --- a/mycelium/src/peer.rs +++ b/mycelium/src/peer.rs @@ -92,16 +92,18 @@ impl Peer { // Framed for peer // Used to send and receive packets from a TCP stream - let mut framed = Framed::new(connection, packet::Codec::new()); + let framed = Framed::with_capacity(connection, packet::Codec::new(), 128 << 10); + let (mut sink, mut stream) = framed.split(); { let peer = peer.clone(); tokio::spawn(async move { + let mut needs_flush = false; loop { select! { // Received over the TCP stream - frame = framed.next() => { + frame = stream.next() => { match frame { Some(Ok(packet)) => { match packet { @@ -135,60 +137,78 @@ impl Peer { } } - Some(packet) = from_routing_data.recv() => { - if let Err(e) = framed.feed(Packet::DataPacket(packet)).await { - error!("Failed to feed data packet to connection: {e}"); - break - } + rv = from_routing_data.recv(), if !needs_flush => { + match rv { + None => break, + Some(packet) => { - for _ in 1..PACKET_COALESCE_WINDOW { - // There can be 2 cases of errors here, empty channel and no more - // senders. In both cases we don't really care at this point. - if let Ok(packet) = from_routing_data.try_recv() { - if let Err(e) = framed.feed(Packet::DataPacket(packet)).await { + needs_flush = true; + + if let Err(e) = sink.feed(Packet::DataPacket(packet)).await { error!("Failed to feed data packet to connection: {e}"); break } - trace!("Instantly queued ready packet to transfer to peer"); - } else { - // No packets ready, flush currently buffered ones - break - } - } - if let Err(e) = framed.flush().await { - error!("Failed to flush buffered peer connection data packets: {e}"); - break + + for _ in 1..PACKET_COALESCE_WINDOW { + // There can be 2 cases of errors here, empty channel and no more + // senders. In both cases we don't really care at this point. + if let Ok(packet) = from_routing_data.try_recv() { + if let Err(e) = sink.feed(Packet::DataPacket(packet)).await { + error!("Failed to feed data packet to connection: {e}"); + break + } + trace!("Instantly queued ready packet to transfer to peer"); + } else { + // No packets ready, flush currently buffered ones + break + } + } + } } } - Some(packet) = from_routing_control.recv() => { - if let Err(e) = framed.feed(Packet::ControlPacket(packet)).await { - error!("Failed to feed control packet to connection: {e}"); - break - } + rv = from_routing_control.recv(), if !needs_flush => { + match rv { + None => break, + Some(packet) => { - for _ in 1..PACKET_COALESCE_WINDOW { - // There can be 2 cases of errors here, empty channel and no more - // senders. In both cases we don't really care at this point. - if let Ok(packet) = from_routing_control.try_recv() { - if let Err(e) = framed.feed(Packet::ControlPacket(packet)).await { - error!("Failed to feed data packet to connection: {e}"); + needs_flush = true; + + if let Err(e) = sink.feed(Packet::ControlPacket(packet)).await { + error!("Failed to feed control packet to connection: {e}"); break } - } else { - // No packets ready, flush currently buffered ones - break + + for _ in 1..PACKET_COALESCE_WINDOW { + // There can be 2 cases of errors here, empty channel and no more + // senders. In both cases we don't really care at this point. + if let Ok(packet) = from_routing_control.try_recv() { + if let Err(e) = sink.feed(Packet::ControlPacket(packet)).await { + error!("Failed to feed data packet to connection: {e}"); + break + } + } else { + // No packets ready, flush currently buffered ones + break + } + } } } + } - if let Err(e) = framed.flush().await { - error!("Failed to flush buffered peer connection control packets: {e}"); + r = sink.flush(), if needs_flush => { + if let Err(err) = r { + error!("Failed to flush peer connection: {err}"); break } + needs_flush = false; } _ = death_watcher.notified() => { + // Attempt gracefull shutdown + let mut framed = sink.reunite(stream).expect("SplitSink and SplitStream here can only be part of the same original Framned; Qed"); + let _ = framed.close().await; break; } } @@ -203,7 +223,7 @@ impl Peer { error!("Peer {remote_id} could not notify router of termination: {e}"); } }); - } + }; Ok(peer) }