leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

perf optimization: avoid tokio::io::split

Browse source
This commit is contained in:
Igor Katson 2025-01-30 11:19:19 +00:00
parent b26c3daac4
commit 62d5288398
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
2 changed files with 44 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

35
crates/librqbit/src/peer_connection.rs
View file

@ -151,11 +151,14 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
self.handler.on_handshake(handshake)?; self.handler.on_handshake(handshake)?;
let (read, write) = conn.into_split();
self.manage_peer( self.manage_peer(
h_supports_extended, h_supports_extended,
read_buf, read_buf,
write_buf, write_buf,
conn, read,
write,
outgoing_chan, outgoing_chan,
have_broadcast, have_broadcast,
) )
@ -179,23 +182,23 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
.unwrap_or_else(|| Duration::from_secs(10)); .unwrap_or_else(|| Duration::from_secs(10));
let now = Instant::now(); let now = Instant::now();
let conn = self.connector.connect(self.addr); let (mut read, mut write) =
let mut conn = with_timeout(connect_timeout, conn) with_timeout(connect_timeout, self.connector.connect(self.addr))
.await .await
.context("error connecting")?; .context("error connecting")?;
self.handler.on_connected(now.elapsed()); self.handler.on_connected(now.elapsed());
let mut write_buf = Vec::<u8>::with_capacity(PIECE_MESSAGE_DEFAULT_LEN); let mut write_buf = Vec::<u8>::with_capacity(PIECE_MESSAGE_DEFAULT_LEN);
let handshake = Handshake::new(self.info_hash, self.peer_id); let handshake = Handshake::new(self.info_hash, self.peer_id);
handshake.serialize(&mut write_buf); handshake.serialize(&mut write_buf);
with_timeout(rwtimeout, conn.write_all(&write_buf)) with_timeout(rwtimeout, write.write_all(&write_buf))
.await .await
.context("error writing handshake")?; .context("error writing handshake")?;
write_buf.clear(); write_buf.clear();
let mut read_buf = ReadBuf::new(); let mut read_buf = ReadBuf::new();
let h = read_buf let h = read_buf
.read_handshake(&mut conn, rwtimeout) .read_handshake(&mut read, rwtimeout)
.await .await
.context("error reading handshake")?; .context("error reading handshake")?;
let h_supports_extended = h.supports_extended(); let h_supports_extended = h.supports_extended();
@ -218,7 +221,8 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
h_supports_extended, h_supports_extended,
read_buf, read_buf,
write_buf, write_buf,
conn, read,
write,
outgoing_chan, outgoing_chan,
have_broadcast, have_broadcast,
) )
@ -230,7 +234,8 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
handshake_supports_extended: bool, handshake_supports_extended: bool,
mut read_buf: ReadBuf, mut read_buf: ReadBuf,
mut write_buf: Vec<u8>, mut write_buf: Vec<u8>,
mut conn: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, mut read: impl tokio::io::AsyncRead + Unpin + Send,
mut write: impl tokio::io::AsyncWrite + Unpin + Send,
mut outgoing_chan: tokio::sync::mpsc::UnboundedReceiver<WriterRequest>, mut outgoing_chan: tokio::sync::mpsc::UnboundedReceiver<WriterRequest>,
mut have_broadcast: tokio::sync::broadcast::Receiver<ValidPieceIndex>, mut have_broadcast: tokio::sync::broadcast::Receiver<ValidPieceIndex>,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
@ -256,14 +261,12 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
my_extended my_extended
.serialize(&mut write_buf, &Default::default) .serialize(&mut write_buf, &Default::default)
.unwrap(); .unwrap();
with_timeout(rwtimeout, conn.write_all(&write_buf)) with_timeout(rwtimeout, write.write_all(&write_buf))
.await .await
.context("error writing extended handshake")?; .context("error writing extended handshake")?;
write_buf.clear(); write_buf.clear();
} }
let (mut read_half, mut write_half) = tokio::io::split(conn);
let writer = async move { let writer = async move {
let keep_alive_interval = self let keep_alive_interval = self
.options .options
@ -274,14 +277,14 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
let len = self let len = self
.handler .handler
.serialize_bitfield_message_to_buf(&mut write_buf)?; .serialize_bitfield_message_to_buf(&mut write_buf)?;
with_timeout(rwtimeout, write_half.write_all(&write_buf[..len])) with_timeout(rwtimeout, write.write_all(&write_buf[..len]))
.await .await
.context("error writing bitfield to peer")?; .context("error writing bitfield to peer")?;
trace!("sent bitfield"); trace!("sent bitfield");
} }
let len = MessageOwned::Unchoke.serialize(&mut write_buf, &Default::default)?; let len = MessageOwned::Unchoke.serialize(&mut write_buf, &Default::default)?;
with_timeout(rwtimeout, write_half.write_all(&write_buf[..len])) with_timeout(rwtimeout, write.write_all(&write_buf[..len]))
.await .await
.context("error writing unchoke")?; .context("error writing unchoke")?;
trace!("sent unchoke"); trace!("sent unchoke");
@ -378,7 +381,7 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
} }
}; };
with_timeout(rwtimeout, write_half.write_all(&write_buf[..len])) with_timeout(rwtimeout, write.write_all(&write_buf[..len]))
.await .await
.context("error writing the message to peer")?; .context("error writing the message to peer")?;
@ -395,7 +398,7 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
let reader = async move { let reader = async move {
loop { loop {
let message = read_buf let message = read_buf
.read_message(&mut read_half, rwtimeout) .read_message(&mut read, rwtimeout)
.await .await
.context("error reading message")?; .context("error reading message")?;
trace!("received: {:?}", &message); trace!("received: {:?}", &message);

44
crates/librqbit/src/stream_connect.rs
View file

@ -30,10 +30,13 @@ impl SocksProxyConfig {
async fn connect( async fn connect(
&self, &self,
addr: SocketAddr, addr: SocketAddr,
) -> anyhow::Result<impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin> { ) -> anyhow::Result<(
impl tokio::io::AsyncRead + Unpin,
impl tokio::io::AsyncWrite + Unpin,
)> {
let proxy_addr = (self.host.as_str(), self.port); let proxy_addr = (self.host.as_str(), self.port);
if let Some((username, password)) = self.username_password.as_ref() { let stream = if let Some((username, password)) = self.username_password.as_ref() {
tokio_socks::tcp::Socks5Stream::connect_with_password( tokio_socks::tcp::Socks5Stream::connect_with_password(
proxy_addr, proxy_addr,
addr, addr,
@ -41,12 +44,14 @@ impl SocksProxyConfig {
password.as_str(), password.as_str(),
) )
.await .await
.context("error connecting to proxy") .context("error connecting to proxy")?
} else { } else {
tokio_socks::tcp::Socks5Stream::connect(proxy_addr, addr) tokio_socks::tcp::Socks5Stream::connect(proxy_addr, addr)
.await .await
.context("error connecting to proxy") .context("error connecting to proxy")?
} };
Ok(tokio::io::split(stream))
} }
} }
@ -61,22 +66,23 @@ impl From<Option<SocksProxyConfig>> for StreamConnector {
} }
} }
pub(crate) trait AsyncReadWrite:
tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin
{
}
impl<T> AsyncReadWrite for T where T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin {}
impl StreamConnector { impl StreamConnector {
pub async fn connect(&self, addr: SocketAddr) -> anyhow::Result<Box<dyn AsyncReadWrite>> { pub async fn connect(
&self,
addr: SocketAddr,
) -> anyhow::Result<(
Box<dyn tokio::io::AsyncRead + Send + Unpin>,
Box<dyn tokio::io::AsyncWrite + Send + Unpin>,
)> {
if let Some(proxy) = self.proxy_config.as_ref() { if let Some(proxy) = self.proxy_config.as_ref() {
return Ok(Box::new(proxy.connect(addr).await?)); let (r, w) = proxy.connect(addr).await?;
return Ok((Box::new(r), Box::new(w)));
} }
Ok(Box::new(
tokio::net::TcpStream::connect(addr) let (r, w) = tokio::net::TcpStream::connect(addr)
.await .await
.context("error connecting")?, .context("error connecting")?
)) .into_split();
Ok((Box::new(r), Box::new(w)))
} }
} }