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

Merge pull request #195 from ikatson/root-span

Browse source 
Fix e2e test to break less
This commit is contained in:
Igor Katson 2024-08-18 16:43:14 +01:00 • committed by GitHub
commit 2ee8366db0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 27 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -54,7 +54,7 @@ use serde::{Deserialize, Serialize};
use tokio::net::{TcpListener, TcpStream};
use tokio_stream::StreamExt;
use tokio_util::sync::{CancellationToken, DropGuard};
use tracing::{debug, error, error_span, info, trace, warn, Instrument};
use tracing::{debug, error, error_span, info, trace, warn, Instrument, Span};
use tracker_comms::TrackerComms;
pub const SUPPORTED_SCHEMES: [&str; 3] = ["http:", "https:", "magnet:"];
@ -114,6 +114,8 @@ pub struct Session {
concurrent_initialize_semaphore: Arc<tokio::sync::Semaphore>,
root_span: Option<Span>,
// This is stored for all tasks to stop when session is dropped.
_cancellation_token_drop_guard: DropGuard,
}
@ -380,6 +382,9 @@ pub struct SessionOptions {
// how many concurrent torrent initializations can happen
pub concurrent_init_limit: Option<usize>,
// the root span to use. If not set will be None.
pub root_span: Option<Span>,
}
async fn create_tcp_listener(
@ -567,11 +572,12 @@ impl Session {
default_storage_factory: opts.default_storage_factory,
reqwest_client,
connector: stream_connector,
root_span: opts.root_span,
concurrent_initialize_semaphore: Arc::new(tokio::sync::Semaphore::new(opts.concurrent_init_limit.unwrap_or(3)))
});
if let Some(mut disk_write_rx) = disk_write_rx {
session.spawn(error_span!("disk_writer"), async move {
session.spawn(error_span!(parent: session.rs(), "disk_writer"), async move {
while let Some(work) = disk_write_rx.recv().await {
trace!(disk_write_rx_queue_len = disk_write_rx.len());
spawner.spawn_block_in_place(work);
@ -582,7 +588,7 @@ impl Session {
if let Some(tcp_listener) = tcp_listener {
session.spawn(
error_span!("tcp_listen", port = tcp_listen_port),
error_span!(parent: session.rs(), "tcp_listen", port = tcp_listen_port),
session.clone().task_tcp_listener(tcp_listener),
);
}
@ -590,7 +596,7 @@ impl Session {
if let Some(listen_port) = tcp_listen_port {
if opts.enable_upnp_port_forwarding {
session.spawn(
error_span!("upnp_forward", port = listen_port),
error_span!(parent: session.rs(), "upnp_forward", port = listen_port),
session.clone().task_upnp_port_forwarder(listen_port),
);
}
@ -613,7 +619,7 @@ impl Session {
st = ps.next(), if !added_all => {
if let Some(st) = st {
let (id, st) = st?;
let span = error_span!("add_torrent", info_hash=?st.info_hash());
let span = error_span!(parent: session.rs(), "add_torrent", info_hash=?st.info_hash());
let (add_torrent, mut opts) = st.into_add_torrent()?;
opts.preferred_id = Some(id);
let fut = session.add_torrent(add_torrent, Some(opts)).instrument(span);
@ -698,7 +704,7 @@ impl Session {
debug!("error checking incoming connection: {e:#}");
e
})
.instrument(error_span!("incoming", addr=%addr))
.instrument(error_span!(parent: self.rs(), "incoming", addr=%addr))
);
}
Err(e) => {
@ -750,6 +756,10 @@ impl Session {
spawn_with_cancel(span, self.cancellation_token.clone(), fut);
}
fn rs(&self) -> Option<tracing::Id> {
self.root_span.as_ref().and_then(|s| s.id())
}
/// Stop the session and all managed tasks.
pub async fn stop(&self) {
let torrents = self
@ -926,7 +936,7 @@ impl Session {
self.main_torrent_info(add_res, opts).await
}
.instrument(error_span!("add_torrent"))
.instrument(error_span!(parent: self.rs(), "add_torrent"))
.boxed()
}
@ -1067,7 +1077,7 @@ impl Session {
}) {
return Ok(AddTorrentResponse::AlreadyManaged(id, handle));
}
let managed_torrent = builder.build(error_span!(parent: None, "torrent", id))?;
let managed_torrent = builder.build(error_span!(parent: self.rs(), "torrent", id))?;
g.add_torrent(managed_torrent.clone(), id);
managed_torrent
};

9
crates/librqbit/src/tests/e2e.rs
View file

@ -60,6 +60,9 @@ async fn test_e2e_download() {
max_random_sleep_ms: rand::thread_rng().gen_range(0u8..16),
}
.as_peer_id();
let listen_range_start = 15100u16 + i as u16;
let listen_range_end = listen_range_start + 1;
let listen_range = listen_range_start..listen_range_end;
let session = crate::Session::new_with_opts(
std::env::temp_dir().join("does_not_exist"),
SessionOptions {
@ -69,10 +72,11 @@ async fn test_e2e_download() {
persistence: None,
peer_id: Some(peer_id),
peer_opts: None,
listen_port_range: Some(15100..17000),
listen_port_range: Some(listen_range),
enable_upnp_port_forwarding: false,
default_storage_factory: None,
defer_writes_up_to: None,
root_span: Some(error_span!(parent: None, "server", id = i)),
..Default::default()
},
)
@ -121,7 +125,7 @@ async fn test_e2e_download() {
session.tcp_listen_port().unwrap(),
))
}
.instrument(error_span!("server", server = i)),
.instrument(error_span!("server", id = i)),
);
futs.push(timeout(Duration::from_secs(30), rx));
}
@ -152,6 +156,7 @@ async fn test_e2e_download() {
persistence: None,
listen_port_range: None,
enable_upnp_port_forwarding: false,
root_span: Some(error_span!("client")),
..Default::default()
},
)

2
crates/peer_binary_protocol/src/extended/handshake.rs
View file

@ -52,7 +52,7 @@ where
ByteBuf: Eq + std::hash::Hash + std::borrow::Borrow<[u8]>,
{
fn get_msgid(&self, msg_type: &'a [u8]) -> Option<u8> {
self.m.get(msg_type).map(|v| *v)
self.m.get(msg_type).copied()
}
pub fn ut_metadata(&self) -> Option<u8> {

1
crates/rqbit/src/main.rs
View file

@ -340,6 +340,7 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
}),
socks_proxy_url: socks_url,
concurrent_init_limit: Some(opts.concurrent_init_limit),
root_span: None,
};
let stats_printer = |session: Arc<Session>| async move {