From 30c5457ec2cbe72637ebccf407cc1627c8d8cc9b Mon Sep 17 00:00:00 2001 From: Royyan Zahir Date: Tue, 16 Jun 2026 07:06:29 +0400 Subject: [PATCH 1/4] feat(stream): add Control::open_stream_on_connection `open_stream` picks one of a peer's connections at random. This adds a variant that opens on a caller-specified connection, so a flow can be routed over a chosen connection/path. `Shared` already tracked per-connection senders; this exposes one via `Shared::sender_on`. Signed-off-by: Royyan Zahir --- protocols/stream/CHANGELOG.md | 3 ++ protocols/stream/src/control.rs | 33 ++++++++++++++++- protocols/stream/src/shared.rs | 5 +++ protocols/stream/tests/lib.rs | 65 ++++++++++++++++++++++++++++++++- 4 files changed, 104 insertions(+), 2 deletions(-) diff --git a/protocols/stream/CHANGELOG.md b/protocols/stream/CHANGELOG.md index 6b9bbc51325..298ec14c77f 100644 --- a/protocols/stream/CHANGELOG.md +++ b/protocols/stream/CHANGELOG.md @@ -1,5 +1,8 @@ ## 0.5.0-alpha +- Add `Control::open_stream_on_connection` to open an outbound stream on a + specific connection instead of a randomly chosen one. + See [PR 6487](https://github.com/libp2p/rust-libp2p/pull/6487). - Raise MSRV to 1.88.0. See [PR 6273](https://github.com/libp2p/rust-libp2p/pull/6273). diff --git a/protocols/stream/src/control.rs b/protocols/stream/src/control.rs index bea15778a18..fa94c6b956f 100644 --- a/protocols/stream/src/control.rs +++ b/protocols/stream/src/control.rs @@ -11,7 +11,7 @@ use futures::{ channel::{mpsc, oneshot}, }; use libp2p_identity::PeerId; -use libp2p_swarm::{Stream, StreamProtocol}; +use libp2p_swarm::{ConnectionId, Stream, StreamProtocol}; use crate::{AlreadyRegistered, handler::NewStream, shared::Shared}; @@ -64,6 +64,37 @@ impl Control { Ok(stream) } + /// Like [`Control::open_stream`] but opens on a specific `connection` + /// instead of picking one of the peer's connections. Errors if that + /// connection is not, or no longer, established. + pub async fn open_stream_on_connection( + &mut self, + peer: PeerId, + connection: ConnectionId, + protocol: StreamProtocol, + ) -> Result { + tracing::debug!(%peer, %connection, "Requesting new stream on connection"); + + let mut new_stream_sender = Shared::lock(&self.shared) + .sender_on(connection) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::NotConnected, "connection not established") + })?; + + let (sender, receiver) = oneshot::channel(); + + new_stream_sender + .send(NewStream { protocol, sender }) + .await + .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?; + + let stream = receiver + .await + .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))??; + + Ok(stream) + } + /// Accept inbound streams for the provided protocol. /// /// To stop accepting streams, simply drop the returned [`IncomingStreams`] handle. diff --git a/protocols/stream/src/shared.rs b/protocols/stream/src/shared.rs index d75f0504689..92ab2a36b3b 100644 --- a/protocols/stream/src/shared.rs +++ b/protocols/stream/src/shared.rs @@ -122,6 +122,11 @@ impl Shared { } } + /// Sender for a specific connection, if it is established. + pub(crate) fn sender_on(&self, connection: ConnectionId) -> Option> { + self.senders.get(&connection).cloned() + } + pub(crate) fn sender(&mut self, peer: PeerId) -> mpsc::Sender { let maybe_sender = self .connections diff --git a/protocols/stream/tests/lib.rs b/protocols/stream/tests/lib.rs index a00a00a8f62..76517b01481 100644 --- a/protocols/stream/tests/lib.rs +++ b/protocols/stream/tests/lib.rs @@ -3,7 +3,7 @@ use std::io; use futures::{AsyncReadExt as _, AsyncWriteExt as _, StreamExt as _}; use libp2p_identity::PeerId; use libp2p_stream as stream; -use libp2p_swarm::{StreamProtocol, Swarm}; +use libp2p_swarm::{ConnectionId, StreamProtocol, Swarm, SwarmEvent}; use libp2p_swarm_test::SwarmExt as _; use stream::OpenStreamError; use tracing::level_filters::LevelFilter; @@ -78,3 +78,66 @@ async fn dial_errors_are_propagated() { assert_eq!(e.kind(), io::ErrorKind::NotConnected); assert_eq!("Dial error: no addresses for peer.", e.to_string()); } + +#[tokio::test] +async fn open_stream_on_connection_targets_an_existing_connection() { + let mut swarm1 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new()); + let mut swarm2 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new()); + + let mut control = swarm1.behaviour().new_control(); + let mut incoming = swarm2.behaviour().new_control().accept(PROTOCOL).unwrap(); + + let (addr, _) = swarm2.listen().with_memory_addr_external().await; + let swarm2_peer_id = *swarm2.local_peer_id(); + + tokio::spawn(async move { + while let Some((_, mut stream)) = incoming.next().await { + stream.write_all(&[42]).await.unwrap(); + stream.close().await.unwrap(); + } + }); + tokio::spawn(swarm2.loop_on_next()); + + // Dial and capture the established connection id. + swarm1.dial(addr).unwrap(); + let conn_id = loop { + if let SwarmEvent::ConnectionEstablished { connection_id, .. } = + swarm1.next_swarm_event().await + { + break connection_id; + } + }; + tokio::spawn(swarm1.loop_on_next()); + + let mut stream = control + .open_stream_on_connection(swarm2_peer_id, conn_id, PROTOCOL) + .await + .unwrap(); + let mut buf = [0u8; 1]; + stream.read_exact(&mut buf).await.unwrap(); + assert_eq!([42], buf); +} + +#[tokio::test] +async fn open_stream_on_connection_errors_on_unknown_connection() { + let mut swarm1 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new()); + let mut swarm2 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new()); + + let mut control = swarm1.behaviour().new_control(); + swarm2.listen().with_memory_addr_external().await; + swarm1.connect(&mut swarm2).await; + let swarm2_peer_id = *swarm2.local_peer_id(); + + tokio::spawn(swarm1.loop_on_next()); + tokio::spawn(swarm2.loop_on_next()); + + let bogus = ConnectionId::new_unchecked(99999); + let error = control + .open_stream_on_connection(swarm2_peer_id, bogus, PROTOCOL) + .await + .unwrap_err(); + let OpenStreamError::Io(e) = error else { + panic!("unexpected error: {error}") + }; + assert_eq!(e.kind(), io::ErrorKind::NotConnected); +} From 612d0ba048837dc9b630bbe2e3283ecfadb93a5b Mon Sep 17 00:00:00 2001 From: Royyan Zahir Date: Fri, 26 Jun 2026 05:38:45 +0400 Subject: [PATCH 2/4] refactor(stream): dedup open_stream channel dance into request_stream No API change. Both open_stream and open_stream_on_connection shared the oneshot/send/await sequence verbatim; factor it out. Also note the ConnectionId source in the doc. --- protocols/stream/src/control.rs | 55 +++++++++++++++------------------ 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/protocols/stream/src/control.rs b/protocols/stream/src/control.rs index fa94c6b956f..5c42746af7f 100644 --- a/protocols/stream/src/control.rs +++ b/protocols/stream/src/control.rs @@ -48,25 +48,15 @@ impl Control { ) -> Result { tracing::debug!(%peer, "Requesting new stream"); - let mut new_stream_sender = Shared::lock(&self.shared).sender(peer); - - let (sender, receiver) = oneshot::channel(); - - new_stream_sender - .send(NewStream { protocol, sender }) - .await - .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?; - - let stream = receiver - .await - .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))??; - - Ok(stream) + let new_stream_sender = Shared::lock(&self.shared).sender(peer); + request_stream(new_stream_sender, protocol).await } /// Like [`Control::open_stream`] but opens on a specific `connection` - /// instead of picking one of the peer's connections. Errors if that - /// connection is not, or no longer, established. + /// instead of picking one of the peer's connections. The `connection` is + /// the [`ConnectionId`] from a [`libp2p_swarm::SwarmEvent::ConnectionEstablished`]. + /// Errors with [`io::ErrorKind::NotConnected`] if that connection is not, or + /// no longer, established. pub async fn open_stream_on_connection( &mut self, peer: PeerId, @@ -75,24 +65,12 @@ impl Control { ) -> Result { tracing::debug!(%peer, %connection, "Requesting new stream on connection"); - let mut new_stream_sender = Shared::lock(&self.shared) + let new_stream_sender = Shared::lock(&self.shared) .sender_on(connection) .ok_or_else(|| { io::Error::new(io::ErrorKind::NotConnected, "connection not established") })?; - - let (sender, receiver) = oneshot::channel(); - - new_stream_sender - .send(NewStream { protocol, sender }) - .await - .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?; - - let stream = receiver - .await - .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))??; - - Ok(stream) + request_stream(new_stream_sender, protocol).await } /// Accept inbound streams for the provided protocol. @@ -106,6 +84,23 @@ impl Control { } } +/// Send a [`NewStream`] request on `new_stream_sender` and await the opened stream. +async fn request_stream( + mut new_stream_sender: mpsc::Sender, + protocol: StreamProtocol, +) -> Result { + let (sender, receiver) = oneshot::channel(); + + new_stream_sender + .send(NewStream { protocol, sender }) + .await + .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?; + + receiver + .await + .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))? +} + /// Errors while opening a new stream. #[derive(Debug)] #[non_exhaustive] From f401fd76fb4878446232939e3b65b79bcdfd33ce Mon Sep 17 00:00:00 2001 From: Royyan Zahir Date: Mon, 29 Jun 2026 17:58:59 +0400 Subject: [PATCH 3/4] refactor(stream): inline sender lookup --- protocols/stream/src/control.rs | 4 +++- protocols/stream/src/shared.rs | 7 +------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/protocols/stream/src/control.rs b/protocols/stream/src/control.rs index 5c42746af7f..671f76aff9b 100644 --- a/protocols/stream/src/control.rs +++ b/protocols/stream/src/control.rs @@ -66,7 +66,9 @@ impl Control { tracing::debug!(%peer, %connection, "Requesting new stream on connection"); let new_stream_sender = Shared::lock(&self.shared) - .sender_on(connection) + .senders + .get(&connection) + .cloned() .ok_or_else(|| { io::Error::new(io::ErrorKind::NotConnected, "connection not established") })?; diff --git a/protocols/stream/src/shared.rs b/protocols/stream/src/shared.rs index 92ab2a36b3b..5dee2e273d7 100644 --- a/protocols/stream/src/shared.rs +++ b/protocols/stream/src/shared.rs @@ -20,7 +20,7 @@ pub(crate) struct Shared { supported_inbound_protocols: HashMap>, connections: HashMap, - senders: HashMap>, + pub(crate) senders: HashMap>, /// Tracks channel pairs for a peer whilst we are dialing them. pending_channels: HashMap, mpsc::Receiver)>, @@ -122,11 +122,6 @@ impl Shared { } } - /// Sender for a specific connection, if it is established. - pub(crate) fn sender_on(&self, connection: ConnectionId) -> Option> { - self.senders.get(&connection).cloned() - } - pub(crate) fn sender(&mut self, peer: PeerId) -> mpsc::Sender { let maybe_sender = self .connections From 9595d98fb5561b05870cd8942f08a474e8f02940 Mon Sep 17 00:00:00 2001 From: Royyan Zahir Date: Mon, 29 Jun 2026 18:00:17 +0400 Subject: [PATCH 4/4] docs(stream): cross-reference open_stream variants --- protocols/stream/src/control.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/protocols/stream/src/control.rs b/protocols/stream/src/control.rs index 671f76aff9b..36fb3713e56 100644 --- a/protocols/stream/src/control.rs +++ b/protocols/stream/src/control.rs @@ -41,6 +41,9 @@ impl Control { /// time which is enforced by requiring `&mut self`. /// /// This backpressure mechanism breaks if you clone [`Control`]s excessively. + /// + /// > Alternatively, [`Control::open_stream_on_connection`] opens on a + /// > specific connection instead of a randomly chosen one. pub async fn open_stream( &mut self, peer: PeerId,