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

Merge pull request #320 from ikatson/stream-traits

Browse source 
[perf] don't use tokio::io::split
This commit is contained in:
Igor Katson 2025-01-30 14:11:36 +00:00 • committed by GitHub
commit ed32b899ed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 78 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/librqbit/src/http_api/handlers/mod.rs
View file

@ -26,7 +26,7 @@ async fn h_api_root(parts: Parts) -> impl IntoResponse {
.headers
.get("Accept")
.and_then(|h| h.to_str().ok())
.map_or(false, |h| h.contains("text/html"))
.is_some_and(|h| h.contains("text/html"))
{
return Redirect::temporary("./web/").into_response();
}

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

@ -88,6 +88,16 @@ where
}
}
struct ManagePeerArgs<R, W> {
handshake_supports_extended: bool,
read_buf: ReadBuf,
write_buf: Vec<u8>,
read: R,
write: W,
outgoing_chan: tokio::sync::mpsc::UnboundedReceiver<WriterRequest>,
have_broadcast: tokio::sync::broadcast::Receiver<ValidPieceIndex>,
}
impl<H: PeerConnectionHandler> PeerConnection<H> {
pub fn new(
addr: SocketAddr,
@ -147,18 +157,21 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
.context("error writing handshake")?;
write_buf.clear();
let h_supports_extended = handshake.supports_extended();
let handshake_supports_extended = handshake.supports_extended();
self.handler.on_handshake(handshake)?;
self.manage_peer(
h_supports_extended,
let (read, write) = conn.into_split();
self.manage_peer(ManagePeerArgs {
handshake_supports_extended,
read_buf,
write_buf,
conn,
read,
write,
outgoing_chan,
have_broadcast,
)
})
.await
}
@ -179,26 +192,26 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
.unwrap_or_else(|| Duration::from_secs(10));
let now = Instant::now();
let conn = self.connector.connect(self.addr);
let mut conn = with_timeout(connect_timeout, conn)
.await
.context("error connecting")?;
let (mut read, mut write) =
with_timeout(connect_timeout, self.connector.connect(self.addr))
.await
.context("error connecting")?;
self.handler.on_connected(now.elapsed());
let mut write_buf = Vec::<u8>::with_capacity(PIECE_MESSAGE_DEFAULT_LEN);
let handshake = Handshake::new(self.info_hash, self.peer_id);
handshake.serialize(&mut write_buf);
with_timeout(rwtimeout, conn.write_all(&write_buf))
with_timeout(rwtimeout, write.write_all(&write_buf))
.await
.context("error writing handshake")?;
write_buf.clear();
let mut read_buf = ReadBuf::new();
let h = read_buf
.read_handshake(&mut conn, rwtimeout)
.read_handshake(&mut read, rwtimeout)
.await
.context("error reading handshake")?;
let h_supports_extended = h.supports_extended();
let handshake_supports_extended = h.supports_extended();
trace!(
peer_id=?Id20::new(h.peer_id),
decoded_id=?try_decode_peer_id(Id20::new(h.peer_id)),
@ -214,26 +227,35 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
self.handler.on_handshake(h)?;
self.manage_peer(
h_supports_extended,
self.manage_peer(ManagePeerArgs {
handshake_supports_extended,
read_buf,
write_buf,
conn,
read,
write,
outgoing_chan,
have_broadcast,
)
})
.await
}
async fn manage_peer(
&self,
handshake_supports_extended: bool,
mut read_buf: ReadBuf,
mut write_buf: Vec<u8>,
mut conn: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
mut outgoing_chan: tokio::sync::mpsc::UnboundedReceiver<WriterRequest>,
mut have_broadcast: tokio::sync::broadcast::Receiver<ValidPieceIndex>,
args: ManagePeerArgs<
impl tokio::io::AsyncRead + Send + Unpin,
impl tokio::io::AsyncWrite + Send + Unpin,
>,
) -> anyhow::Result<()> {
let ManagePeerArgs {
handshake_supports_extended,
mut read_buf,
mut write_buf,
mut read,
mut write,
mut outgoing_chan,
mut have_broadcast,
} = args;
use tokio::io::AsyncWriteExt;
let rwtimeout = self
@ -256,14 +278,12 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
my_extended
.serialize(&mut write_buf, &Default::default)
.unwrap();
with_timeout(rwtimeout, conn.write_all(&write_buf))
with_timeout(rwtimeout, write.write_all(&write_buf))
.await
.context("error writing extended handshake")?;
write_buf.clear();
}
let (mut read_half, mut write_half) = tokio::io::split(conn);
let writer = async move {
let keep_alive_interval = self
.options
@ -274,14 +294,14 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
let len = self
.handler
.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
.context("error writing bitfield to peer")?;
trace!("sent bitfield");
}
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
.context("error writing unchoke")?;
trace!("sent unchoke");
@ -378,7 +398,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
.context("error writing the message to peer")?;
@ -395,7 +415,7 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
let reader = async move {
loop {
let message = read_buf
.read_message(&mut read_half, rwtimeout)
.read_message(&mut read, rwtimeout)
.await
.context("error reading message")?;
trace!("received: {:?}", &message);

2
crates/librqbit/src/session.rs
View file

@ -996,7 +996,7 @@ impl Session {
name,
} = add_res;
let private = metadata.as_ref().map_or(false, |m| m.info.private);
let private = metadata.as_ref().is_some_and(|m| m.info.private);
let make_peer_rx = || {
self.make_peer_rx(

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

@ -30,10 +30,13 @@ impl SocksProxyConfig {
async fn connect(
&self,
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);
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(
proxy_addr,
addr,
@ -41,12 +44,14 @@ impl SocksProxyConfig {
password.as_str(),
)
.await
.context("error connecting to proxy")
.context("error connecting to proxy")?
} else {
tokio_socks::tcp::Socks5Stream::connect(proxy_addr, addr)
.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 {
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() {
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)
.await
.context("error connecting")?,
))
let (r, w) = tokio::net::TcpStream::connect(addr)
.await
.context("error connecting")?
.into_split();
Ok((Box::new(r), Box::new(w)))
}
}

5
crates/librqbit/src/torrent_state/live/mod.rs
View file

@ -1253,10 +1253,7 @@ impl PeerHandler {
///
/// If this returns, an existing in-flight piece was marked to be ours.
fn try_steal_old_slow_piece(&self, threshold: f64) -> Option<ValidPieceIndex> {
let my_avg_time = match self.counters.average_piece_download_time() {
Some(t) => t,
None => return None,
};
let my_avg_time = self.counters.average_piece_download_time()?;
let (stolen_idx, from_peer) = {
let mut g = self.state.lock_write("try_steal_old_slow_piece");

4
crates/librqbit/src/torrent_state/live/peer/mod.rs
View file

@ -263,8 +263,6 @@ impl LivePeerState {
}
pub fn has_full_torrent(&self, total_pieces: usize) -> bool {
self.bitfield
.get(0..total_pieces)
.map_or(false, |s| s.all())
self.bitfield.get(0..total_pieces).is_some_and(|s| s.all())
}
}